From d4d889febbf6b6de40194e6f85b364b908aec817 Mon Sep 17 00:00:00 2001 From: zhan3916 <92444497+zhan3916@users.noreply.github.com> Date: Mon, 6 Jul 2026 14:38:08 +0800 Subject: [PATCH] Update flashMLA implementation on MACA (#1) --- LICENSE | 7 +- benchmark/bench_flash_mla.py | 520 - benchmark/visualize.py | 29 - csrc/flash_api/flash_api.cpp | 400 +- ...576_hdimv512_m16n16_bf16_True_True_sm80.cu | 18 - ...576_hdimv512_m32n16_bf16_True_True_sm80.cu | 18 - ...576_hdimv512_m32n16_fp16_True_True_sm80.cu | 18 - csrc/flash_api/flash_mla.h | 120 +- .../flash_fwd_dispatch_template.h | 112 +- .../flash_fwd_launch_template.h | 156 +- csrc/flash_dispatch/flash_fwd_run_template.h | 69 + .../flash_run_fwd_template_impl.h | 28 - csrc/flash_kernel/feature/alibi.h | 76 - csrc/flash_kernel/feature/attn_mask.h | 166 - csrc/flash_kernel/feature/dropout.h | 206 - csrc/flash_kernel/feature/mask.h | 159 +- csrc/flash_kernel/feature/rotary.h | 529 - csrc/flash_kernel/feature/softmax.h | 245 +- .../flash_dense_mla_decode_kernel.h | 87 + .../flash_fwd_split_kernel_k64_V1x8.h | 686 - ...=> flash_fwd_splitkv_mla_combine_kernel.h} | 90 +- .../flash_sparse_mla_decode_kernel.h | 63 + csrc/flash_kernel/kernel_traits.h | 166 +- ...kernel_k64_16x16_4waves_stage1_xcore1000.h | 470 + ...d_mla_kernel_k64_16x16_4waves_xcore1000.h} | 598 +- ...d_mla_kernel_k64_32x16_4waves_xcore1000.h} | 626 +- ...d_mla_kernel_k64_64x16_8waves_xcore1000.h} | 527 +- ...se_mla_kernel_k64_64x16_8waves_xcore1000.h | 275 + ...se_prefill_kernel_64x16_8waves_xcore1000.h | 347 + ...wd_mla_kernel_k64_64x16_8waves_xcore1500.h | 341 + ...wd_mla_kernel_k64_64x32_8waves_xcore1500.h | 336 + csrc/flash_run/run_mha_fwd.cpp | 12 - csrc/flash_run/{run_mha.h => run_mla.h} | 3 +- csrc/flash_run/run_mla_fwd.cpp | 24 + csrc/instantiations/flash_fwd_mla_metadata.cu | 97 + ..._m16n16_bf16_True_True_1stage_xcore1000.cu | 18 + ..._m16n16_bf16_True_True_split_xcore1000.cu} | 6 +- ...m16n16_fp16_True_True_1stage_xcore1000.cu} | 10 +- ..._m16n16_fp16_True_True_split_xcore1000.cu} | 6 +- ..._m32n16_bf16_True_True_split_xcore1000.cu} | 6 +- ..._m32n16_fp16_True_True_split_xcore1000.cu} | 6 +- ..._m64n16_bf16_True_True_split_xcore1000.cu} | 6 +- ..._m64n16_fp16_True_True_split_xcore1000.cu} | 6 +- ...2_m64n16_bf16_True_True_split_xcore1000.cu | 18 + ..._m64n16_bf16_True_True_split_xcore1000.cu} | 10 +- ..._m64n16_fp16_True_True_split_xcore1000.cu} | 10 +- ...2_m64n16_bf16_True_True_split_xcore1500.cu | 19 + ...2_m64n16_fp16_True_True_split_xcore1500.cu | 19 + ...2_m64n32_bf16_True_True_split_xcore1500.cu | 19 + ...2_m64n32_fp16_True_True_split_xcore1500.cu | 19 + ...2_m64n16_bf16_True_True_split_xcore1500.cu | 19 + ...2_m64n16_bf16_True_True_split_xcore1500.cu | 19 + ...2_m64n16_fp16_True_True_split_xcore1500.cu | 19 + csrc/mctlass/.gitignore | 42 + csrc/mctlass/include/cute/algorithm/axpby.hpp | 79 + csrc/mctlass/include/cute/algorithm/clear.hpp | 66 + csrc/mctlass/include/cute/algorithm/copy.hpp | 523 + csrc/mctlass/include/cute/algorithm/fill.hpp | 87 + .../include/cute/algorithm/functional.hpp | 198 + csrc/mctlass/include/cute/algorithm/gemm.hpp | 744 + .../mctlass/include/cute/algorithm/prefer.hpp | 46 + .../cute/algorithm/tensor_algorithms.hpp | 123 + .../cute/algorithm/tuple_algorithms.hpp | 875 + .../include/cute/arch/cluster_sm90.hpp | 243 + csrc/mctlass/include/cute/arch/copy.hpp | 71 + csrc/mctlass/include/cute/arch/copy_sm75.hpp | 322 + csrc/mctlass/include/cute/arch/copy_sm80.hpp | 201 + csrc/mctlass/include/cute/arch/copy_sm90.hpp | 225 + .../include/cute/arch/copy_sm90_desc.hpp | 201 + .../include/cute/arch/copy_sm90_tma.hpp | 861 + csrc/mctlass/include/cute/arch/mma.hpp | 64 + csrc/mctlass/include/cute/arch/mma_sm75.hpp | 120 + csrc/mctlass/include/cute/arch/mma_sm80.hpp | 2714 +++ csrc/mctlass/include/cute/arch/mma_sm90.hpp | 961 ++ .../include/cute/arch/mma_sm90_desc.hpp | 135 + .../include/cute/arch/mma_sm90_gmma.hpp | 12780 ++++++++++++++ csrc/mctlass/include/cute/arch/util.hpp | 249 + csrc/mctlass/include/cute/atom/copy_atom.hpp | 707 + .../mctlass/include/cute/atom/copy_traits.hpp | 131 + .../include/cute/atom/copy_traits_sm75.hpp | 160 + .../include/cute/atom/copy_traits_sm80.hpp | 130 + .../include/cute/atom/copy_traits_sm90.hpp | 132 + .../cute/atom/copy_traits_sm90_tma.hpp | 973 ++ csrc/mctlass/include/cute/atom/mma_atom.hpp | 1063 ++ csrc/mctlass/include/cute/atom/mma_traits.hpp | 208 + .../include/cute/atom/mma_traits_sm75.hpp | 81 + .../include/cute/atom/mma_traits_sm80.hpp | 604 + .../include/cute/atom/mma_traits_sm90.hpp | 132 + .../cute/atom/mma_traits_sm90_gmma.hpp | 3185 ++++ csrc/mctlass/include/cute/config.hpp | 169 + .../include/cute/container/alignment.hpp | 70 + csrc/mctlass/include/cute/container/array.hpp | 334 + .../include/cute/container/array_aligned.hpp | 42 + .../include/cute/container/array_subbyte.hpp | 633 + .../include/cute/container/bit_field.hpp | 131 + .../include/cute/container/cuda_types.hpp | 186 + csrc/mctlass/include/cute/container/tuple.hpp | 702 + .../include/cute/container/type_list.hpp | 136 + csrc/mctlass/include/cute/int_tuple.hpp | 875 + csrc/mctlass/include/cute/layout.hpp | 1687 ++ .../include/cute/numeric/arithmetic_tuple.hpp | 436 + csrc/mctlass/include/cute/numeric/bfloat.hpp | 53 + csrc/mctlass/include/cute/numeric/complex.hpp | 193 + csrc/mctlass/include/cute/numeric/float8.hpp | 43 + csrc/mctlass/include/cute/numeric/half.hpp | 41 + csrc/mctlass/include/cute/numeric/int.hpp | 135 + .../include/cute/numeric/integer_sequence.hpp | 151 + .../include/cute/numeric/integer_subbyte.hpp | 233 + .../cute/numeric/integral_constant.hpp | 436 + csrc/mctlass/include/cute/numeric/math.hpp | 313 + csrc/mctlass/include/cute/numeric/real.hpp | 56 + csrc/mctlass/include/cute/numeric/tfloat.hpp | 53 + csrc/mctlass/include/cute/numeric/uint128.hpp | 259 + csrc/mctlass/include/cute/pointer.hpp | 325 + csrc/mctlass/include/cute/stride.hpp | 442 + csrc/mctlass/include/cute/swizzle.hpp | 527 + csrc/mctlass/include/cute/swizzle_layout.hpp | 1012 ++ csrc/mctlass/include/cute/swizzle_ptr.hpp | 284 + csrc/mctlass/include/cute/tensor.hpp | 956 ++ .../mctlass/include/cute/tensor_predicate.hpp | 63 + csrc/mctlass/include/cute/tile.hpp | 58 + csrc/mctlass/include/cute/underscore.hpp | 150 + csrc/mctlass/include/cute/util/debug.hpp | 163 + csrc/mctlass/include/cute/util/print.hpp | 141 + .../mctlass/include/cute/util/type_traits.hpp | 207 + csrc/mctlass/include/mctlass/aligned_buffer.h | 129 + csrc/mctlass/include/mctlass/arch/arch.h | 111 + csrc/mctlass/include/mctlass/arch/barrier.h | 444 + .../include/mctlass/arch/cache_operation.h | 66 + csrc/mctlass/include/mctlass/arch/memory.h | 684 + .../include/mctlass/arch/memory_sm75.h | 354 + .../include/mctlass/arch/memory_sm80.h | 459 + csrc/mctlass/include/mctlass/arch/mma.h | 233 + csrc/mctlass/include/mctlass/arch/mma_sm50.h | 432 + csrc/mctlass/include/mctlass/arch/mma_sm60.h | 252 + csrc/mctlass/include/mctlass/arch/mma_sm61.h | 144 + csrc/mctlass/include/mctlass/arch/mma_sm70.h | 659 + csrc/mctlass/include/mctlass/arch/mma_sm75.h | 1285 ++ csrc/mctlass/include/mctlass/arch/mma_sm80.h | 2339 +++ csrc/mctlass/include/mctlass/arch/mma_sm90.h | 266 + .../include/mctlass/arch/mma_sparse_sm80.h | 1646 ++ .../include/mctlass/arch/reg_reconfig.h | 68 + csrc/mctlass/include/mctlass/arch/simd.h | 125 + csrc/mctlass/include/mctlass/arch/simd_sm60.h | 110 + csrc/mctlass/include/mctlass/arch/simd_sm61.h | 147 + csrc/mctlass/include/mctlass/arch/wmma.h | 230 + csrc/mctlass/include/mctlass/arch/wmma_sm70.h | 137 + csrc/mctlass/include/mctlass/arch/wmma_sm72.h | 210 + csrc/mctlass/include/mctlass/arch/wmma_sm75.h | 209 + csrc/mctlass/include/mctlass/array.h | 2461 +++ .../include/mctlass/array_planar_complex.h | 103 + csrc/mctlass/include/mctlass/array_subbyte.h | 564 + csrc/mctlass/include/mctlass/barrier.h | 185 + csrc/mctlass/include/mctlass/bfloat16.h | 514 + csrc/mctlass/include/mctlass/blas3.h | 176 + csrc/mctlass/include/mctlass/block_striped.h | 267 + .../include/mctlass/cluster_launch.hpp | 235 + csrc/mctlass/include/mctlass/complex.h | 693 + csrc/mctlass/include/mctlass/constants.h | 1239 ++ .../mctlass/conv/conv2d_problem_size.h | 645 + .../mctlass/conv/conv3d_problem_size.h | 477 + .../include/mctlass/conv/convolution.h | 168 + .../mctlass/conv/device/direct_convolution.h | 269 + .../conv/device/implicit_gemm_convolution.h | 328 + .../device/implicit_gemm_convolution_fusion.h | 268 + .../mctlass/conv/kernel/default_conv2d.h | 272 + .../conv/kernel/default_conv2d_dgrad.h | 1927 +++ .../conv/kernel/default_conv2d_fprop.h | 1989 +++ .../conv/kernel/default_conv2d_fprop_fusion.h | 357 + .../default_conv2d_fprop_with_broadcast.h | 130 + .../default_conv2d_fprop_with_reduction.h | 130 + .../conv/kernel/default_conv2d_group_fprop.h | 622 + .../conv/kernel/default_conv2d_wgrad.h | 1011 ++ .../conv/kernel/default_conv2d_wgrad_fusion.h | 325 + .../conv/kernel/default_conv3d_dgrad.h | 303 + .../conv/kernel/default_conv3d_fprop.h | 515 + .../conv/kernel/default_conv3d_fprop_fusion.h | 360 + .../conv/kernel/default_conv3d_wgrad.h | 509 + .../conv/kernel/default_depthwise_fprop.h | 588 + .../mctlass/conv/kernel/direct_convolution.h | 505 + .../conv/kernel/implicit_gemm_convolution.h | 456 + .../kernel/implicit_gemm_convolution_fusion.h | 463 + .../implicit_gemm_convolution_strided_dgrad.h | 492 + ...cit_gemm_convolution_with_fused_epilogue.h | 499 + .../mctlass/conv/thread/depthwise_mma.h | 325 + ...rad_filter_tile_access_iterator_analytic.h | 485 + ...ad_filter_tile_access_iterator_optimized.h | 619 + ...t_gradient_tile_access_iterator_analytic.h | 606 + ..._gradient_tile_access_iterator_optimized.h | 821 + ...activation_tile_access_iterator_analytic.h | 332 + ...vation_tile_access_iterator_few_channels.h | 360 + ...tion_tile_access_iterator_fixed_channels.h | 353 + ...ctivation_tile_access_iterator_optimized.h | 422 + ...rop_filter_tile_access_iterator_analytic.h | 319 + ...filter_tile_access_iterator_few_channels.h | 289 + ...lter_tile_access_iterator_fixed_channels.h | 275 + ...op_filter_tile_access_iterator_optimized.h | 317 + .../mctlass/conv/threadblock/conv2d_params.h | 893 + .../conv/threadblock/conv2d_tile_iterator.h | 337 + ...activation_tile_access_iterator_analytic.h | 285 + ...ctivation_tile_access_iterator_optimized.h | 321 + ...t_gradient_tile_access_iterator_analytic.h | 260 + ..._gradient_tile_access_iterator_optimized.h | 310 + ...rad_filter_tile_access_iterator_analytic.h | 268 + ...ad_filter_tile_access_iterator_optimized.h | 289 + ...t_gradient_tile_access_iterator_analytic.h | 343 + ..._gradient_tile_access_iterator_optimized.h | 490 + ...activation_tile_access_iterator_analytic.h | 291 + ...ctivation_tile_access_iterator_optimized.h | 479 + ...rop_filter_tile_access_iterator_analytic.h | 253 + ...op_filter_tile_access_iterator_optimized.h | 277 + .../mctlass/conv/threadblock/conv3d_params.h | 508 + ...activation_tile_access_iterator_analytic.h | 289 + ...ctivation_tile_access_iterator_optimized.h | 319 + ...t_gradient_tile_access_iterator_analytic.h | 267 + ..._gradient_tile_access_iterator_optimized.h | 310 + .../depthwise_direct_conv_params.h | 230 + ...erator_direct_conv_fixed_stride_dilation.h | 314 + ...le_access_iterator_direct_conv_optimized.h | 291 + .../depthwise_fprop_direct_conv_multistage.h | 551 + ...le_access_iterator_direct_conv_optimized.h | 261 + .../threadblock/depthwise_fprop_pipelined.h | 336 + .../conv/threadblock/depthwise_mma_base.h | 229 + ...depthwise_mma_core_with_lane_access_size.h | 952 ++ .../implicit_gemm_fprop_fusion_multistage.h | 802 + .../threadblock/implicit_gemm_multistage.h | 539 + .../threadblock/implicit_gemm_pipelined.h | 320 + .../implicit_gemm_wgrad_fusion_multistage.h | 729 + ...icated_scale_bias_vector_access_iterator.h | 471 + .../predicated_scale_bias_vector_iterator.h | 371 + .../conv/threadblock/threadblock_swizzle.h | 193 + .../mctlass/conv/warp/mma_depthwise_simt.h | 380 + .../warp/mma_depthwise_simt_tile_iterator.h | 862 + .../conv/warp/scale_bias_relu_transform.h | 225 + csrc/mctlass/include/mctlass/coord.h | 480 + csrc/mctlass/include/mctlass/core_io.h | 287 + .../mctlass/detail/dependent_false.hpp | 86 + csrc/mctlass/include/mctlass/device_kernel.h | 120 + .../collective/builders/sm90_builder.inl | 545 + .../collective/collective_builder.hpp | 77 + .../collective/collective_epilogue.hpp | 55 + .../epilogue/collective/default_epilogue.hpp | 216 + .../mctlass/epilogue/collective/detail.hpp | 211 + .../collective/epilogue_tensor_broadcast.hpp | 245 + .../collective/sm70_epilogue_vectorized.hpp | 335 + .../sm90_epilogue_tma_warpspecialized.hpp | 606 + ...e_tma_warpspecialized_bias_elementwise.hpp | 679 + .../mctlass/epilogue/dispatch_policy.hpp | 136 + .../mctlass/epilogue/thread/activation.h | 705 + .../mctlass/epilogue/thread/conversion_op.h | 132 + .../mctlass/epilogue/thread/detail.hpp | 52 + .../epilogue/thread/linear_combination.h | 330 + .../linear_combination_bias_elementwise.h | 270 + .../thread/linear_combination_bias_relu.h | 471 + .../thread/linear_combination_clamp.h | 689 + .../thread/linear_combination_dgelu.h | 250 + .../thread/linear_combination_drelu.h | 452 + .../epilogue/thread/linear_combination_gelu.h | 70 + .../thread/linear_combination_generic.h | 208 + .../thread/linear_combination_hardswish.h | 69 + .../thread/linear_combination_leaky_relu.h | 231 + .../thread/linear_combination_params.h | 75 + .../linear_combination_planar_complex.h | 237 + .../epilogue/thread/linear_combination_relu.h | 572 + .../thread/linear_combination_relu0.h | 543 + .../linear_combination_residual_block.h | 301 + .../thread/linear_combination_sigmoid.h | 70 + .../epilogue/thread/linear_combination_silu.h | 69 + .../linear_combination_tensor_broadcast.hpp | 252 + .../linear_combination_with_elementwise.h | 234 + .../mctlass/epilogue/thread/reduction_op.h | 97 + .../mctlass/epilogue/thread/scale_type.h | 62 + .../default_epilogue_complex_tensor_op.h | 255 + ...default_epilogue_complex_tensor_op_blas3.h | 264 + .../default_epilogue_direct_store.h | 74 + .../default_epilogue_planar_complex.h | 241 + .../threadblock/default_epilogue_simt.h | 422 + .../threadblock/default_epilogue_tensor_op.h | 848 + .../default_epilogue_tensor_op_blas3.h | 175 + ...default_epilogue_tensor_op_row_broadcast.h | 183 + .../default_epilogue_volta_tensor_op.h | 337 + .../default_epilogue_with_broadcast.h | 245 + .../default_epilogue_with_reduction.h | 177 + .../default_epilogue_wmma_tensor_op.h | 165 + .../threadblock/default_thread_map_simt.h | 127 + .../default_thread_map_tensor_op.h | 208 + .../default_thread_map_volta_tensor_op.h | 228 + .../default_thread_map_wmma_tensor_op.h | 113 + .../direct_store_epilogue_iterator.h | 142 + .../mctlass/epilogue/threadblock/epilogue.h | 652 + .../epilogue/threadblock/epilogue_base.h | 240 + .../threadblock/epilogue_base_streamk.h | 197 + .../epilogue/threadblock/epilogue_depthwise.h | 335 + .../threadblock/epilogue_direct_store.h | 374 + .../threadblock/epilogue_gemm_k_reduction.h | 212 + .../threadblock/epilogue_planar_complex.h | 441 + .../threadblock/epilogue_smem_accumulator.h | 230 + .../epilogue_streamk_with_broadcast.h | 443 + .../epilogue_visitor_with_softmax.h | 513 + .../threadblock/epilogue_with_broadcast.h | 1778 ++ .../threadblock/epilogue_with_reduction.h | 827 + .../threadblock/epilogue_with_visitor.h | 411 + .../epilogue/threadblock/epilogue_workspace.h | 197 + .../threadblock/interleaved_epilogue.h | 432 + .../threadblock/output_iterator_parameter.h | 92 + .../threadblock/output_tile_thread_map.h | 730 + .../threadblock/predicated_tile_iterator.h | 1343 ++ .../predicated_tile_iterator_affine.h | 616 + ...cated_tile_iterator_affine_layout_params.h | 156 + .../predicated_tile_iterator_blas3.h | 635 + .../predicated_tile_iterator_direct_conv.h | 445 + .../predicated_tile_iterator_params.h | 475 + .../predicated_tile_iterator_predicates.h | 309 + .../predicated_tile_iterator_row_broadcast.h | 519 + .../predicated_tile_iterator_strided_dgrad.h | 480 + .../threadblock/shared_load_iterator.h | 223 + .../threadblock/shared_load_iterator_mixed.h | 594 + .../shared_load_iterator_pitch_liner.h | 194 + .../fragment_iterator_complex_tensor_op.h | 187 + ...ment_iterator_gaussian_complex_tensor_op.h | 194 + .../epilogue/warp/fragment_iterator_simt.h | 164 + .../warp/fragment_iterator_tensor_op.h | 277 + .../warp/fragment_iterator_volta_tensor_op.h | 269 + .../warp/fragment_iterator_wmma_tensor_op.h | 165 + .../mctlass/epilogue/warp/simt_policy.h | 107 + .../mctlass/epilogue/warp/tensor_op_policy.h | 148 + .../epilogue/warp/tile_iterator_simt.h | 790 + .../epilogue/warp/tile_iterator_tensor_op.h | 761 + .../warp/tile_iterator_tensor_op_mixed.h | 817 + .../warp/tile_iterator_volta_tensor_op.h | 440 + .../warp/tile_iterator_wmma_tensor_op.h | 237 + .../epilogue/warp/volta_tensor_op_policy.h | 195 + .../epilogue/warp/wmma_tensor_op_policy.h | 101 + csrc/mctlass/include/mctlass/fast_math.h | 971 ++ csrc/mctlass/include/mctlass/float8.h | 1207 ++ .../include/mctlass/floating_point_nvrtc.h | 65 + csrc/mctlass/include/mctlass/functional.h | 533 + .../collective/builders/sm90_gmma_builder.inl | 832 + .../gemm/collective/collective_builder.hpp | 82 + .../gemm/collective/collective_mma.hpp | 74 + .../gemm/collective/sm70_mma_twostage.hpp | 596 + .../gemm/collective/sm80_mma_multistage.hpp | 688 + .../sm90_mma_multistage_gmma_ss.hpp | 600 + .../sm90_mma_tma_gmma_rs_warpspecialized.hpp | 594 + .../gemm/collective/sm90_mma_tma_gmma_ss.hpp | 497 + .../sm90_mma_tma_gmma_ss_warpspecialized.hpp | 501 + .../mctlass/gemm/device/base_grouped.h | 478 + .../gemm/device/default_gemm_configuration.h | 815 + .../include/mctlass/gemm/device/ell_gemm.h | 848 + .../include/mctlass/gemm/device/gemm.h | 771 + .../include/mctlass/gemm/device/gemm_array.h | 737 + .../mctlass/gemm/device/gemm_batched.h | 703 + .../mctlass/gemm/device/gemm_complex.h | 717 + .../mctlass/gemm/device/gemm_grouped.h | 61 + .../device/gemm_layernorm_mainloop_fusion.h | 385 + .../include/mctlass/gemm/device/gemm_sparse.h | 514 + .../gemm/device/gemm_sparse_row_broadcast.h | 514 + .../gemm/device/gemm_splitk_parallel.h | 637 + .../mctlass/gemm/device/gemm_universal.h | 441 + .../gemm/device/gemm_universal_adapter.h | 554 + .../mctlass/gemm/device/gemm_universal_base.h | 419 + .../gemm_universal_streamk_with_broadcast.h | 386 + .../device/gemm_universal_with_broadcast.h | 386 + .../gemm/device/gemm_with_k_reduction.h | 415 + .../include/mctlass/gemm/device/gemv.h | 186 + .../include/mctlass/gemm/device/rank_2k.h | 547 + .../mctlass/gemm/device/rank_2k_grouped.h | 63 + .../include/mctlass/gemm/device/rank_k.h | 509 + .../include/mctlass/gemm/device/symm.h | 602 + .../include/mctlass/gemm/device/trmm.h | 758 + .../include/mctlass/gemm/dispatch_policy.hpp | 170 + csrc/mctlass/include/mctlass/gemm/gemm.h | 623 + .../mctlass/gemm/kernel/default_ell_gemm.h | 837 + .../mctlass/gemm/kernel/default_gemm.h | 1111 ++ .../gemm/kernel/default_gemm_complex.h | 404 + .../gemm/kernel/default_gemm_grouped.h | 384 + ...ult_gemm_grouped_softmax_mainloop_fusion.h | 164 + .../default_gemm_layernorm_mainloop_fusion.h | 137 + .../default_gemm_planar_complex_universal.h | 352 + .../mctlass/gemm/kernel/default_gemm_sparse.h | 191 + .../default_gemm_sparse_row_broadcast.h | 191 + .../kernel/default_gemm_splitk_parallel.h | 136 + .../default_gemm_streamk_with_broadcast.h | 146 + .../gemm/kernel/default_gemm_universal.h | 396 + .../gemm/kernel/default_gemm_with_broadcast.h | 243 + .../kernel/default_gemm_with_k_reduction.h | 150 + .../gemm/kernel/default_gemm_with_reduction.h | 246 + .../mctlass/gemm/kernel/default_gemv.h | 132 + .../mctlass/gemm/kernel/default_rank_2k.h | 285 + .../gemm/kernel/default_rank_2k_complex.h | 334 + .../gemm/kernel/default_rank_2k_grouped.h | 355 + .../gemm/kernel/default_rank_2k_universal.h | 346 + .../mctlass/gemm/kernel/default_rank_k.h | 247 + .../gemm/kernel/default_rank_k_complex.h | 429 + .../gemm/kernel/default_rank_k_universal.h | 305 + .../mctlass/gemm/kernel/default_symm.h | 322 + .../gemm/kernel/default_symm_complex.h | 509 + .../gemm/kernel/default_symm_universal.h | 342 + .../mctlass/gemm/kernel/default_trmm.h | 269 + .../gemm/kernel/default_trmm_complex.h | 265 + .../gemm/kernel/default_trmm_universal.h | 359 + .../include/mctlass/gemm/kernel/ell_gemm.h | 830 + .../include/mctlass/gemm/kernel/gemm.h | 381 + .../include/mctlass/gemm/kernel/gemm_array.h | 264 + .../mctlass/gemm/kernel/gemm_batched.h | 279 + .../mctlass/gemm/kernel/gemm_grouped.h | 481 + .../kernel/gemm_grouped_problem_visitor.h | 122 + .../gemm_grouped_softmax_mainloop_fusion.h | 510 + .../kernel/gemm_layernorm_mainloop_fusion.h | 785 + .../include/mctlass/gemm/kernel/gemm_params.h | 199 + .../mctlass/gemm/kernel/gemm_pipelined.h | 158 + .../mctlass/gemm/kernel/gemm_planar_complex.h | 723 + .../gemm/kernel/gemm_planar_complex_array.h | 617 + .../gemm/kernel/gemm_splitk_parallel.h | 253 + .../kernel/gemm_streamk_with_fused_epilogue.h | 2405 +++ .../gemm/kernel/gemm_transpose_operands.h | 124 + .../mctlass/gemm/kernel/gemm_universal.h | 698 + .../mctlass/gemm/kernel/gemm_universal.hpp | 68 + .../gemm/kernel/gemm_universal_streamk.h | 1175 ++ .../gemm/kernel/gemm_with_fused_epilogue.h | 1500 ++ .../gemm/kernel/gemm_with_k_reduction.h | 700 + .../include/mctlass/gemm/kernel/gemv.h | 638 + .../gemm/kernel/gemv_batched_strided.h | 241 + .../gemm/kernel/grouped_problem_visitor.h | 470 + .../gemm/kernel/params_universal_base.h | 245 + .../mctlass/gemm/kernel/rank_2k_grouped.h | 697 + .../kernel/rank_2k_grouped_problem_visitor.h | 376 + .../gemm/kernel/rank_2k_transpose_operands.h | 129 + .../mctlass/gemm/kernel/rank_2k_universal.h | 779 + .../mctlass/gemm/kernel/rank_k_universal.h | 566 + .../include/mctlass/gemm/kernel/sm70_gemm.hpp | 250 + .../mctlass/gemm/kernel/sm90_gemm_tma.hpp | 328 + .../kernel/sm90_gemm_tma_warpspecialized.hpp | 441 + ...0_gemm_tma_warpspecialized_cooperative.hpp | 504 + ...sm90_gemm_tma_warpspecialized_pingpong.hpp | 536 + .../gemm/kernel/sm90_tile_scheduler.hpp | 182 + .../include/mctlass/gemm/kernel/sparse_gemm.h | 400 + .../gemm/kernel/sparse_gemm_row_broadcast.h | 400 + .../mctlass/gemm/kernel/symm_universal.h | 699 + .../mctlass/gemm/kernel/trmm_universal.h | 599 + .../mctlass/include/mctlass/gemm/thread/mma.h | 90 + .../include/mctlass/gemm/thread/mma_sm50.h | 536 + .../include/mctlass/gemm/thread/mma_sm60.h | 1178 ++ .../include/mctlass/gemm/thread/mma_sm61.h | 284 + .../gemm/threadblock/default_ell_mma.h | 734 + .../gemm/threadblock/default_gemv_core.h | 151 + .../mctlass/gemm/threadblock/default_mma.h | 823 + .../gemm/threadblock/default_mma_core.h | 116 + .../gemm/threadblock/default_mma_core_simt.h | 1723 ++ .../gemm/threadblock/default_mma_core_sm70.h | 682 + .../gemm/threadblock/default_mma_core_sm75.h | 1284 ++ .../gemm/threadblock/default_mma_core_sm80.h | 2916 ++++ .../default_mma_core_sparse_sm80.h | 834 + .../default_mma_core_with_access_size.h | 328 + .../default_mma_core_with_reduction.h | 167 + .../gemm/threadblock/default_mma_core_wmma.h | 712 + .../default_mma_layernorm_mainloop_fusion.h | 178 + .../default_mma_planar_complex_multistage.h | 136 + .../default_mma_planar_complex_pipelined.h | 130 + .../default_mma_softmax_mainloop_fusion.h | 160 + .../threadblock/default_mma_with_reduction.h | 141 + .../default_multistage_mma_complex.h | 159 + .../default_multistage_mma_complex_core.h | 119 + ...default_multistage_mma_complex_core_sm80.h | 1808 ++ .../default_multistage_trmm_complex.h | 556 + .../gemm/threadblock/default_sparse_mma.h | 196 + .../mctlass/gemm/threadblock/default_trmm.h | 445 + .../gemm/threadblock/ell_mma_multistage.h | 648 + .../gemm/threadblock/ell_mma_pipelined.h | 376 + .../include/mctlass/gemm/threadblock/gemv.h | 147 + .../mctlass/gemm/threadblock/index_remat.h | 107 + .../mctlass/gemm/threadblock/mma_base.h | 236 + .../gemm/threadblock/mma_blas3_multistage.h | 707 + ...mma_layernorm_mainloop_fusion_multistage.h | 863 + .../mctlass/gemm/threadblock/mma_multistage.h | 745 + .../mctlass/gemm/threadblock/mma_pipelined.h | 439 + .../threadblock/mma_planar_complex_base.h | 208 + .../mma_planar_complex_multistage.h | 646 + .../mma_planar_complex_pipelined.h | 424 + .../gemm/threadblock/mma_singlestage.h | 265 + .../mma_softmax_mainloop_fusion_multistage.h | 756 + .../gemm/threadblock/mma_sparse_base.h | 273 + .../gemm/threadblock/mma_sparse_multistage.h | 668 + .../mma_with_reduction_multistage.h | 539 + .../gemm/threadblock/threadblock_swizzle.h | 459 + .../threadblock/threadblock_swizzle_streamk.h | 801 + .../gemm/warp/default_mma_complex_tensor_op.h | 612 + .../gemm/warp/default_mma_sparse_tensor_op.h | 165 + .../mctlass/gemm/warp/default_mma_tensor_op.h | 123 + .../gemm/warp/default_mma_tensor_op_sm80.h | 238 + .../default_mma_with_reduction_tensor_op.h | 92 + .../gemm/warp/default_mma_wmma_tensor_op.h | 130 + .../warp/layernorm_scale_bias_transform.h | 141 + csrc/mctlass/include/mctlass/gemm/warp/mma.h | 61 + .../mctlass/gemm/warp/mma_complex_tensor_op.h | 1171 ++ .../warp/mma_complex_tensor_op_fast_f32.h | 663 + ...mma_complex_tensor_op_tile_iterator_sm80.h | 2661 +++ .../warp/mma_gaussian_complex_tensor_op.h | 642 + ...ian_complex_tensor_op_tile_iterator_sm80.h | 395 + .../mctlass/gemm/warp/mma_planar_complex.h | 182 + .../include/mctlass/gemm/warp/mma_simt.h | 263 + .../mctlass/gemm/warp/mma_simt_policy.h | 69 + .../gemm/warp/mma_simt_tile_iterator.h | 1892 +++ .../mctlass/gemm/warp/mma_sparse_tensor_op.h | 341 + .../include/mctlass/gemm/warp/mma_tensor_op.h | 433 + .../gemm/warp/mma_tensor_op_fast_f32.h | 471 + .../warp/mma_tensor_op_fragment_iterator.h | 528 + .../mctlass/gemm/warp/mma_tensor_op_policy.h | 65 + .../mctlass/gemm/warp/mma_tensor_op_sm70.h | 280 + .../warp/mma_tensor_op_tile_access_iterator.h | 362 + .../gemm/warp/mma_tensor_op_tile_iterator.h | 10135 +++++++++++ .../warp/mma_tensor_op_tile_iterator_sm70.h | 3094 ++++ .../warp/mma_tensor_op_tile_iterator_sm80.h | 2946 ++++ .../warp/mma_tensor_op_tile_iterator_sparse.h | 380 + .../warp/mma_tensor_op_tile_iterator_wmma.h | 800 + .../mctlass/gemm/warp/mma_tensor_op_wmma.h | 223 + .../gemm/warp/mma_with_reduction_tensor_op.h | 453 + .../gemm/warp/scale_bias_tile_iterator.h | 572 + .../gemm/warp/softmax_scale_bias_transform.h | 117 + .../gemm/warp/tile_iterator_planar_complex.h | 250 + csrc/mctlass/include/mctlass/half.h | 918 + .../mctlass/include/mctlass/integer_subbyte.h | 241 + .../include/mctlass/kernel_hardware_info.hpp | 73 + csrc/mctlass/include/mctlass/kernel_launch.h | 73 + csrc/mctlass/include/mctlass/layout/layout.h | 64 + csrc/mctlass/include/mctlass/layout/matrix.h | 1365 ++ csrc/mctlass/include/mctlass/layout/permute.h | 827 + .../include/mctlass/layout/pitch_linear.h | 148 + csrc/mctlass/include/mctlass/layout/tensor.h | 636 + .../layout/tensor_op_multiplicand_sm70.h | 1044 ++ .../layout/tensor_op_multiplicand_sm75.h | 1573 ++ .../layout/tensor_op_multiplicand_sm80.h | 1139 ++ csrc/mctlass/include/mctlass/layout/vector.h | 104 + csrc/mctlass/include/mctlass/matrix.h | 14129 ++++++++++++++++ csrc/mctlass/include/mctlass/matrix_coord.h | 164 + csrc/mctlass/include/mctlass/matrix_shape.h | 65 + csrc/mctlass/include/mctlass/mctlass.h | 245 + .../include/mctlass/numeric_conversion.h | 2649 +++ csrc/mctlass/include/mctlass/numeric_types.h | 94 + .../include/mctlass/pipeline/pipeline.hpp | 36 + .../mctlass/pipeline/sm90_pipeline.hpp | 995 ++ .../include/mctlass/pitch_linear_coord.h | 181 + .../include/mctlass/platform/platform.h | 883 + .../include/mctlass/predicate_vector.h | 524 + csrc/mctlass/include/mctlass/quaternion.h | 752 + csrc/mctlass/include/mctlass/real.h | 61 + .../mctlass/reduction/device/reduce_split_k.h | 223 + .../mctlass/reduction/device/tensor_reduce.h | 264 + .../device/tensor_reduce_affine_contiguous.h | 373 + .../device/tensor_reduce_affine_strided.h | 361 + .../reduction/kernel/reduce_softmax_final.h | 267 + .../mctlass/reduction/kernel/reduce_split_k.h | 248 + .../kernel/tensor_reduce_affine_contiguous.h | 606 + .../kernel/tensor_reduce_affine_strided.h | 641 + .../include/mctlass/reduction/thread/reduce.h | 234 + .../reduction/thread/reduction_operators.h | 235 + .../mctlass/reduction/threadblock_swizzle.h | 67 + .../include/mctlass/relatively_equal.h | 219 + csrc/mctlass/include/mctlass/semaphore.h | 125 + .../include/mctlass/subbyte_reference.h | 633 + csrc/mctlass/include/mctlass/tensor_coord.h | 326 + csrc/mctlass/include/mctlass/tensor_ref.h | 418 + .../mctlass/tensor_ref_planar_complex.h | 374 + csrc/mctlass/include/mctlass/tensor_view.h | 297 + .../mctlass/tensor_view_planar_complex.h | 301 + csrc/mctlass/include/mctlass/tfloat32.h | 478 + csrc/mctlass/include/mctlass/thread/matrix.h | 198 + csrc/mctlass/include/mctlass/trace.h | 59 + .../collective/sm90_wgmma_transpose.hpp | 336 + .../transform/pitch_linear_thread_map.h | 931 + .../mctlass/transform/thread/transpose.h | 164 + .../mctlass/transform/thread/unary_op.h | 105 + .../transform/threadblock/ell_iterator.h | 199 + .../ell_predicated_tile_access_iterator.h | 1350 ++ .../ell_predicated_tile_iterator.h | 1315 ++ ...icated_scale_bias_vector_access_iterator.h | 375 + .../predicated_scale_bias_vector_iterator.h | 335 + .../predicated_tile_access_iterator.h | 2117 +++ ...icated_tile_access_iterator_2dthreadtile.h | 835 + .../predicated_tile_access_iterator_params.h | 289 + ...d_tile_access_iterator_triangular_matrix.h | 892 + .../threadblock/predicated_tile_iterator.h | 1888 +++ .../predicated_tile_iterator_2dthreadtile.h | 787 + ...edicated_tile_iterator_triangular_matrix.h | 818 + .../predicated_vector_access_iterator.h | 417 + ...egular_scale_bias_vector_access_iterator.h | 253 + .../regular_tile_access_iterator.h | 58 + ...egular_tile_access_iterator_pitch_linear.h | 408 + ...access_iterator_pitch_linear_direct_conv.h | 587 + .../regular_tile_access_iterator_tensor_op.h | 820 + ...ular_tile_access_iterator_tensor_op_sm80.h | 1532 ++ .../threadblock/regular_tile_iterator.h | 62 + .../regular_tile_iterator_pitch_linear.h | 552 + ..._tile_iterator_pitch_linear_2dthreadtile.h | 509 + .../regular_tile_iterator_tensor_op.h | 1107 ++ .../regular_tile_iterator_tensor_op_sm70.h | 1460 ++ .../transform/threadblock/vector_iterator.h | 149 + .../transform/warp/vector_fragment_iterator.h | 283 + csrc/mctlass/include/mctlass/uint128.h | 266 + csrc/mctlass/include/mctlass/wmma_array.h | 133 + csrc/utils/arch.h | 6 + csrc/utils/block_info.h | 4 +- csrc/utils/hdim_switch.h | 94 - csrc/utils/host_utils.h | 22 + csrc/utils/logger.cpp | 170 + csrc/utils/logger.h | 263 + csrc/utils/philox.cuh | 54 - csrc/utils/print_parameter.cpp | 218 + csrc/utils/print_parameter.h | 15 + csrc/utils/process_str.cpp | 31 + csrc/utils/process_str.h | 10 + csrc/utils/static_switch.h | 98 +- csrc/utils/utils.h | 637 +- flash_mla/__init__.py | 3 +- flash_mla/flash_mla_interface.py | 204 +- setup.py | 372 +- tests/test_flash_mla.py | 135 - 617 files changed, 296690 insertions(+), 4379 deletions(-) delete mode 100644 benchmark/bench_flash_mla.py delete mode 100644 benchmark/visualize.py delete mode 100644 csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_bf16_True_True_sm80.cu delete mode 100644 csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_bf16_True_True_sm80.cu delete mode 100644 csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_fp16_True_True_sm80.cu create mode 100644 csrc/flash_dispatch/flash_fwd_run_template.h delete mode 100644 csrc/flash_dispatch/flash_run_fwd_template_impl.h delete mode 100644 csrc/flash_kernel/feature/alibi.h delete mode 100644 csrc/flash_kernel/feature/attn_mask.h delete mode 100644 csrc/flash_kernel/feature/dropout.h delete mode 100644 csrc/flash_kernel/feature/rotary.h create mode 100644 csrc/flash_kernel/flash_dense_mla_decode_kernel.h delete mode 100644 csrc/flash_kernel/flash_fwd_split_kernel_k64_V1x8.h rename csrc/flash_kernel/{flash_fwd_split_kernel.h => flash_fwd_splitkv_mla_combine_kernel.h} (69%) create mode 100644 csrc/flash_kernel/flash_sparse_mla_decode_kernel.h create mode 100644 csrc/flash_kernel/xcore1000/flash_fwd_mla_kernel_k64_16x16_4waves_stage1_xcore1000.h rename csrc/flash_kernel/{flash_fwd_split_kernel_k64_16x16_4waves.h => xcore1000/flash_fwd_mla_kernel_k64_16x16_4waves_xcore1000.h} (76%) rename csrc/flash_kernel/{flash_fwd_split_kernel_k64_32x16_4waves.h => xcore1000/flash_fwd_mla_kernel_k64_32x16_4waves_xcore1000.h} (80%) rename csrc/flash_kernel/{flash_fwd_split_kernel_k64_64x16_8waves.h => xcore1000/flash_fwd_mla_kernel_k64_64x16_8waves_xcore1000.h} (76%) create mode 100644 csrc/flash_kernel/xcore1000/flash_fwd_sparse_mla_kernel_k64_64x16_8waves_xcore1000.h create mode 100644 csrc/flash_kernel/xcore1000/sparse_prefill_kernel_64x16_8waves_xcore1000.h create mode 100644 csrc/flash_kernel/xcore1500/flash_fwd_mla_kernel_k64_64x16_8waves_xcore1500.h create mode 100644 csrc/flash_kernel/xcore1500/flash_fwd_mla_kernel_k64_64x32_8waves_xcore1500.h delete mode 100644 csrc/flash_run/run_mha_fwd.cpp rename csrc/flash_run/{run_mha.h => run_mla.h} (52%) create mode 100644 csrc/flash_run/run_mla_fwd.cpp create mode 100644 csrc/instantiations/flash_fwd_mla_metadata.cu create mode 100644 csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_bf16_True_True_1stage_xcore1000.cu rename csrc/{flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_bf16_True_True_split_sm80.cu => instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_bf16_True_True_split_xcore1000.cu} (76%) rename csrc/{flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_fp16_True_True_sm80.cu => instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_fp16_True_True_1stage_xcore1000.cu} (69%) rename csrc/{flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_fp16_True_True_split_sm80.cu => instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_fp16_True_True_split_xcore1000.cu} (77%) rename csrc/{flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_bf16_True_True_split_sm80.cu => instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m32n16_bf16_True_True_split_xcore1000.cu} (76%) rename csrc/{flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_fp16_True_True_split_sm80.cu => instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m32n16_fp16_True_True_split_xcore1000.cu} (77%) rename csrc/{flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_bf16_True_True_split_sm80.cu => instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu} (76%) rename csrc/{flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_fp16_True_True_split_sm80.cu => instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1000.cu} (77%) create mode 100644 csrc/instantiations/xcore1000/flash_fwd_prefill_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu rename csrc/{flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_fp16_True_True_sm80.cu => instantiations/xcore1000/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu} (71%) rename csrc/{flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_bf16_True_True_sm80.cu => instantiations/xcore1000/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1000.cu} (72%) create mode 100644 csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu create mode 100644 csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1500.cu create mode 100644 csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n32_bf16_True_True_split_xcore1500.cu create mode 100644 csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n32_fp16_True_True_split_xcore1500.cu create mode 100644 csrc/instantiations/xcore1500/flash_fwd_prefill_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu create mode 100644 csrc/instantiations/xcore1500/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu create mode 100644 csrc/instantiations/xcore1500/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1500.cu create mode 100644 csrc/mctlass/.gitignore create mode 100644 csrc/mctlass/include/cute/algorithm/axpby.hpp create mode 100644 csrc/mctlass/include/cute/algorithm/clear.hpp create mode 100644 csrc/mctlass/include/cute/algorithm/copy.hpp create mode 100644 csrc/mctlass/include/cute/algorithm/fill.hpp create mode 100644 csrc/mctlass/include/cute/algorithm/functional.hpp create mode 100644 csrc/mctlass/include/cute/algorithm/gemm.hpp create mode 100644 csrc/mctlass/include/cute/algorithm/prefer.hpp create mode 100644 csrc/mctlass/include/cute/algorithm/tensor_algorithms.hpp create mode 100644 csrc/mctlass/include/cute/algorithm/tuple_algorithms.hpp create mode 100644 csrc/mctlass/include/cute/arch/cluster_sm90.hpp create mode 100644 csrc/mctlass/include/cute/arch/copy.hpp create mode 100644 csrc/mctlass/include/cute/arch/copy_sm75.hpp create mode 100644 csrc/mctlass/include/cute/arch/copy_sm80.hpp create mode 100644 csrc/mctlass/include/cute/arch/copy_sm90.hpp create mode 100644 csrc/mctlass/include/cute/arch/copy_sm90_desc.hpp create mode 100644 csrc/mctlass/include/cute/arch/copy_sm90_tma.hpp create mode 100644 csrc/mctlass/include/cute/arch/mma.hpp create mode 100644 csrc/mctlass/include/cute/arch/mma_sm75.hpp create mode 100644 csrc/mctlass/include/cute/arch/mma_sm80.hpp create mode 100644 csrc/mctlass/include/cute/arch/mma_sm90.hpp create mode 100644 csrc/mctlass/include/cute/arch/mma_sm90_desc.hpp create mode 100644 csrc/mctlass/include/cute/arch/mma_sm90_gmma.hpp create mode 100644 csrc/mctlass/include/cute/arch/util.hpp create mode 100644 csrc/mctlass/include/cute/atom/copy_atom.hpp create mode 100644 csrc/mctlass/include/cute/atom/copy_traits.hpp create mode 100644 csrc/mctlass/include/cute/atom/copy_traits_sm75.hpp create mode 100644 csrc/mctlass/include/cute/atom/copy_traits_sm80.hpp create mode 100644 csrc/mctlass/include/cute/atom/copy_traits_sm90.hpp create mode 100644 csrc/mctlass/include/cute/atom/copy_traits_sm90_tma.hpp create mode 100644 csrc/mctlass/include/cute/atom/mma_atom.hpp create mode 100644 csrc/mctlass/include/cute/atom/mma_traits.hpp create mode 100644 csrc/mctlass/include/cute/atom/mma_traits_sm75.hpp create mode 100644 csrc/mctlass/include/cute/atom/mma_traits_sm80.hpp create mode 100644 csrc/mctlass/include/cute/atom/mma_traits_sm90.hpp create mode 100644 csrc/mctlass/include/cute/atom/mma_traits_sm90_gmma.hpp create mode 100644 csrc/mctlass/include/cute/config.hpp create mode 100644 csrc/mctlass/include/cute/container/alignment.hpp create mode 100644 csrc/mctlass/include/cute/container/array.hpp create mode 100644 csrc/mctlass/include/cute/container/array_aligned.hpp create mode 100644 csrc/mctlass/include/cute/container/array_subbyte.hpp create mode 100644 csrc/mctlass/include/cute/container/bit_field.hpp create mode 100644 csrc/mctlass/include/cute/container/cuda_types.hpp create mode 100644 csrc/mctlass/include/cute/container/tuple.hpp create mode 100644 csrc/mctlass/include/cute/container/type_list.hpp create mode 100644 csrc/mctlass/include/cute/int_tuple.hpp create mode 100644 csrc/mctlass/include/cute/layout.hpp create mode 100644 csrc/mctlass/include/cute/numeric/arithmetic_tuple.hpp create mode 100644 csrc/mctlass/include/cute/numeric/bfloat.hpp create mode 100644 csrc/mctlass/include/cute/numeric/complex.hpp create mode 100644 csrc/mctlass/include/cute/numeric/float8.hpp create mode 100644 csrc/mctlass/include/cute/numeric/half.hpp create mode 100644 csrc/mctlass/include/cute/numeric/int.hpp create mode 100644 csrc/mctlass/include/cute/numeric/integer_sequence.hpp create mode 100644 csrc/mctlass/include/cute/numeric/integer_subbyte.hpp create mode 100644 csrc/mctlass/include/cute/numeric/integral_constant.hpp create mode 100644 csrc/mctlass/include/cute/numeric/math.hpp create mode 100644 csrc/mctlass/include/cute/numeric/real.hpp create mode 100644 csrc/mctlass/include/cute/numeric/tfloat.hpp create mode 100644 csrc/mctlass/include/cute/numeric/uint128.hpp create mode 100644 csrc/mctlass/include/cute/pointer.hpp create mode 100644 csrc/mctlass/include/cute/stride.hpp create mode 100644 csrc/mctlass/include/cute/swizzle.hpp create mode 100644 csrc/mctlass/include/cute/swizzle_layout.hpp create mode 100644 csrc/mctlass/include/cute/swizzle_ptr.hpp create mode 100644 csrc/mctlass/include/cute/tensor.hpp create mode 100644 csrc/mctlass/include/cute/tensor_predicate.hpp create mode 100644 csrc/mctlass/include/cute/tile.hpp create mode 100644 csrc/mctlass/include/cute/underscore.hpp create mode 100644 csrc/mctlass/include/cute/util/debug.hpp create mode 100644 csrc/mctlass/include/cute/util/print.hpp create mode 100644 csrc/mctlass/include/cute/util/type_traits.hpp create mode 100644 csrc/mctlass/include/mctlass/aligned_buffer.h create mode 100644 csrc/mctlass/include/mctlass/arch/arch.h create mode 100644 csrc/mctlass/include/mctlass/arch/barrier.h create mode 100644 csrc/mctlass/include/mctlass/arch/cache_operation.h create mode 100644 csrc/mctlass/include/mctlass/arch/memory.h create mode 100644 csrc/mctlass/include/mctlass/arch/memory_sm75.h create mode 100644 csrc/mctlass/include/mctlass/arch/memory_sm80.h create mode 100644 csrc/mctlass/include/mctlass/arch/mma.h create mode 100644 csrc/mctlass/include/mctlass/arch/mma_sm50.h create mode 100644 csrc/mctlass/include/mctlass/arch/mma_sm60.h create mode 100644 csrc/mctlass/include/mctlass/arch/mma_sm61.h create mode 100644 csrc/mctlass/include/mctlass/arch/mma_sm70.h create mode 100644 csrc/mctlass/include/mctlass/arch/mma_sm75.h create mode 100644 csrc/mctlass/include/mctlass/arch/mma_sm80.h create mode 100644 csrc/mctlass/include/mctlass/arch/mma_sm90.h create mode 100644 csrc/mctlass/include/mctlass/arch/mma_sparse_sm80.h create mode 100644 csrc/mctlass/include/mctlass/arch/reg_reconfig.h create mode 100644 csrc/mctlass/include/mctlass/arch/simd.h create mode 100644 csrc/mctlass/include/mctlass/arch/simd_sm60.h create mode 100644 csrc/mctlass/include/mctlass/arch/simd_sm61.h create mode 100644 csrc/mctlass/include/mctlass/arch/wmma.h create mode 100644 csrc/mctlass/include/mctlass/arch/wmma_sm70.h create mode 100644 csrc/mctlass/include/mctlass/arch/wmma_sm72.h create mode 100644 csrc/mctlass/include/mctlass/arch/wmma_sm75.h create mode 100644 csrc/mctlass/include/mctlass/array.h create mode 100644 csrc/mctlass/include/mctlass/array_planar_complex.h create mode 100644 csrc/mctlass/include/mctlass/array_subbyte.h create mode 100644 csrc/mctlass/include/mctlass/barrier.h create mode 100644 csrc/mctlass/include/mctlass/bfloat16.h create mode 100644 csrc/mctlass/include/mctlass/blas3.h create mode 100644 csrc/mctlass/include/mctlass/block_striped.h create mode 100644 csrc/mctlass/include/mctlass/cluster_launch.hpp create mode 100644 csrc/mctlass/include/mctlass/complex.h create mode 100644 csrc/mctlass/include/mctlass/constants.h create mode 100644 csrc/mctlass/include/mctlass/conv/conv2d_problem_size.h create mode 100644 csrc/mctlass/include/mctlass/conv/conv3d_problem_size.h create mode 100644 csrc/mctlass/include/mctlass/conv/convolution.h create mode 100644 csrc/mctlass/include/mctlass/conv/device/direct_convolution.h create mode 100644 csrc/mctlass/include/mctlass/conv/device/implicit_gemm_convolution.h create mode 100644 csrc/mctlass/include/mctlass/conv/device/implicit_gemm_convolution_fusion.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_conv2d.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_dgrad.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop_fusion.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop_with_broadcast.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop_with_reduction.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_group_fprop.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_wgrad.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_wgrad_fusion.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_dgrad.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_fprop.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_fprop_fusion.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_wgrad.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/default_depthwise_fprop.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/direct_convolution.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution_fusion.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution_strided_dgrad.h create mode 100644 csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution_with_fused_epilogue.h create mode 100644 csrc/mctlass/include/mctlass/conv/thread/depthwise_mma.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_filter_tile_access_iterator_analytic.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_filter_tile_access_iterator_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_output_gradient_tile_access_iterator_analytic.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_output_gradient_tile_access_iterator_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_analytic.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_few_channels.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_fixed_channels.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_few_channels.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_fixed_channels.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_params.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_tile_iterator.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_analytic.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_analytic.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_filter_tile_access_iterator_analytic.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_filter_tile_access_iterator_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_output_gradient_tile_access_iterator_analytic.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_output_gradient_tile_access_iterator_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_analytic.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_analytic.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv3d_params.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_analytic.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_analytic.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/depthwise_direct_conv_params.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_activation_tile_access_iterator_direct_conv_fixed_stride_dilation.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_activation_tile_access_iterator_direct_conv_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_direct_conv_multistage.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_filter_tile_access_iterator_direct_conv_optimized.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_pipelined.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/depthwise_mma_base.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/depthwise_mma_core_with_lane_access_size.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_fprop_fusion_multistage.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_multistage.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_pipelined.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_wgrad_fusion_multistage.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/predicated_scale_bias_vector_access_iterator.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/predicated_scale_bias_vector_iterator.h create mode 100644 csrc/mctlass/include/mctlass/conv/threadblock/threadblock_swizzle.h create mode 100644 csrc/mctlass/include/mctlass/conv/warp/mma_depthwise_simt.h create mode 100644 csrc/mctlass/include/mctlass/conv/warp/mma_depthwise_simt_tile_iterator.h create mode 100644 csrc/mctlass/include/mctlass/conv/warp/scale_bias_relu_transform.h create mode 100644 csrc/mctlass/include/mctlass/coord.h create mode 100644 csrc/mctlass/include/mctlass/core_io.h create mode 100644 csrc/mctlass/include/mctlass/detail/dependent_false.hpp create mode 100644 csrc/mctlass/include/mctlass/device_kernel.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/collective/builders/sm90_builder.inl create mode 100644 csrc/mctlass/include/mctlass/epilogue/collective/collective_builder.hpp create mode 100644 csrc/mctlass/include/mctlass/epilogue/collective/collective_epilogue.hpp create mode 100644 csrc/mctlass/include/mctlass/epilogue/collective/default_epilogue.hpp create mode 100644 csrc/mctlass/include/mctlass/epilogue/collective/detail.hpp create mode 100644 csrc/mctlass/include/mctlass/epilogue/collective/epilogue_tensor_broadcast.hpp create mode 100644 csrc/mctlass/include/mctlass/epilogue/collective/sm70_epilogue_vectorized.hpp create mode 100644 csrc/mctlass/include/mctlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp create mode 100644 csrc/mctlass/include/mctlass/epilogue/collective/sm90_epilogue_tma_warpspecialized_bias_elementwise.hpp create mode 100644 csrc/mctlass/include/mctlass/epilogue/dispatch_policy.hpp create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/activation.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/conversion_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/detail.hpp create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_bias_elementwise.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_bias_relu.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_clamp.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_dgelu.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_drelu.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_gelu.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_generic.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_hardswish.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_leaky_relu.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_params.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_planar_complex.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_relu.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_relu0.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_residual_block.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_sigmoid.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_silu.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_tensor_broadcast.hpp create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_with_elementwise.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/reduction_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/thread/scale_type.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_complex_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_complex_tensor_op_blas3.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_direct_store.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_planar_complex.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_simt.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_tensor_op_blas3.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_tensor_op_row_broadcast.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_with_broadcast.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_with_reduction.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_simt.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_volta_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_wmma_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/direct_store_epilogue_iterator.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_base.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_base_streamk.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_depthwise.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_direct_store.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_gemm_k_reduction.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_planar_complex.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_smem_accumulator.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_streamk_with_broadcast.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_visitor_with_softmax.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_with_broadcast.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_with_reduction.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_with_visitor.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_workspace.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/interleaved_epilogue.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/output_iterator_parameter.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/output_tile_thread_map.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_affine.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_affine_layout_params.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_blas3.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_direct_conv.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_params.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_predicates.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_row_broadcast.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_strided_dgrad.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/shared_load_iterator.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/shared_load_iterator_mixed.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/threadblock/shared_load_iterator_pitch_liner.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_complex_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_gaussian_complex_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_simt.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_volta_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_wmma_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/simt_policy.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/tensor_op_policy.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_simt.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_tensor_op_mixed.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_volta_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_wmma_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/volta_tensor_op_policy.h create mode 100644 csrc/mctlass/include/mctlass/epilogue/warp/wmma_tensor_op_policy.h create mode 100644 csrc/mctlass/include/mctlass/fast_math.h create mode 100644 csrc/mctlass/include/mctlass/float8.h create mode 100644 csrc/mctlass/include/mctlass/floating_point_nvrtc.h create mode 100644 csrc/mctlass/include/mctlass/functional.h create mode 100644 csrc/mctlass/include/mctlass/gemm/collective/builders/sm90_gmma_builder.inl create mode 100644 csrc/mctlass/include/mctlass/gemm/collective/collective_builder.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/collective/collective_mma.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/collective/sm70_mma_twostage.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/collective/sm80_mma_multistage.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_multistage_gmma_ss.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/device/base_grouped.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/default_gemm_configuration.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/ell_gemm.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_array.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_batched.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_complex.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_grouped.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_layernorm_mainloop_fusion.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_sparse.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_sparse_row_broadcast.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_splitk_parallel.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_universal.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_universal_adapter.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_universal_base.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_universal_streamk_with_broadcast.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_universal_with_broadcast.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemm_with_k_reduction.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/gemv.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/rank_2k.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/rank_2k_grouped.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/rank_k.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/symm.h create mode 100644 csrc/mctlass/include/mctlass/gemm/device/trmm.h create mode 100644 csrc/mctlass/include/mctlass/gemm/dispatch_policy.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/gemm.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_ell_gemm.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_complex.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_grouped.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_grouped_softmax_mainloop_fusion.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_layernorm_mainloop_fusion.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_planar_complex_universal.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_sparse.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_sparse_row_broadcast.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_splitk_parallel.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_streamk_with_broadcast.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_universal.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_with_broadcast.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_with_k_reduction.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_with_reduction.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_gemv.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k_complex.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k_grouped.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k_universal.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_rank_k.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_rank_k_complex.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_rank_k_universal.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_symm.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_symm_complex.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_symm_universal.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_trmm.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_trmm_complex.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/default_trmm_universal.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/ell_gemm.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_array.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_batched.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_grouped.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_grouped_problem_visitor.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_grouped_softmax_mainloop_fusion.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_layernorm_mainloop_fusion.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_params.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_pipelined.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_planar_complex.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_planar_complex_array.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_splitk_parallel.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_streamk_with_fused_epilogue.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_transpose_operands.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_universal.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_universal.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_universal_streamk.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_with_fused_epilogue.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemm_with_k_reduction.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemv.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/gemv_batched_strided.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/grouped_problem_visitor.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/params_universal_base.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_grouped.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_grouped_problem_visitor.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_transpose_operands.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_universal.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/rank_k_universal.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/sm70_gemm.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/sm90_tile_scheduler.hpp create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/sparse_gemm.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/sparse_gemm_row_broadcast.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/symm_universal.h create mode 100644 csrc/mctlass/include/mctlass/gemm/kernel/trmm_universal.h create mode 100644 csrc/mctlass/include/mctlass/gemm/thread/mma.h create mode 100644 csrc/mctlass/include/mctlass/gemm/thread/mma_sm50.h create mode 100644 csrc/mctlass/include/mctlass/gemm/thread/mma_sm60.h create mode 100644 csrc/mctlass/include/mctlass/gemm/thread/mma_sm61.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_ell_mma.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_gemv_core.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_simt.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sm70.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sm75.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sm80.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sparse_sm80.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_with_access_size.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_with_reduction.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_wmma.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_layernorm_mainloop_fusion.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_planar_complex_multistage.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_planar_complex_pipelined.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_softmax_mainloop_fusion.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_with_reduction.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_mma_complex.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_mma_complex_core.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_trmm_complex.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_sparse_mma.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/default_trmm.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/ell_mma_multistage.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/ell_mma_pipelined.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/gemv.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/index_remat.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/mma_base.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/mma_blas3_multistage.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/mma_layernorm_mainloop_fusion_multistage.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/mma_multistage.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/mma_pipelined.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/mma_planar_complex_base.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/mma_planar_complex_multistage.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/mma_planar_complex_pipelined.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/mma_singlestage.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/mma_softmax_mainloop_fusion_multistage.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/mma_sparse_base.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/mma_sparse_multistage.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/mma_with_reduction_multistage.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/threadblock_swizzle.h create mode 100644 csrc/mctlass/include/mctlass/gemm/threadblock/threadblock_swizzle_streamk.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/default_mma_complex_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/default_mma_sparse_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/default_mma_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/default_mma_tensor_op_sm80.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/default_mma_with_reduction_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/default_mma_wmma_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/layernorm_scale_bias_transform.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_complex_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_complex_tensor_op_fast_f32.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_complex_tensor_op_tile_iterator_sm80.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_gaussian_complex_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_gaussian_complex_tensor_op_tile_iterator_sm80.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_planar_complex.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_simt.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_simt_policy.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_simt_tile_iterator.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_sparse_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_fast_f32.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_fragment_iterator.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_policy.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_sm70.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_access_iterator.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm70.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_sparse.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_wmma.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_wmma.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/mma_with_reduction_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/scale_bias_tile_iterator.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/softmax_scale_bias_transform.h create mode 100644 csrc/mctlass/include/mctlass/gemm/warp/tile_iterator_planar_complex.h create mode 100644 csrc/mctlass/include/mctlass/half.h create mode 100644 csrc/mctlass/include/mctlass/integer_subbyte.h create mode 100644 csrc/mctlass/include/mctlass/kernel_hardware_info.hpp create mode 100644 csrc/mctlass/include/mctlass/kernel_launch.h create mode 100644 csrc/mctlass/include/mctlass/layout/layout.h create mode 100644 csrc/mctlass/include/mctlass/layout/matrix.h create mode 100644 csrc/mctlass/include/mctlass/layout/permute.h create mode 100644 csrc/mctlass/include/mctlass/layout/pitch_linear.h create mode 100644 csrc/mctlass/include/mctlass/layout/tensor.h create mode 100644 csrc/mctlass/include/mctlass/layout/tensor_op_multiplicand_sm70.h create mode 100644 csrc/mctlass/include/mctlass/layout/tensor_op_multiplicand_sm75.h create mode 100644 csrc/mctlass/include/mctlass/layout/tensor_op_multiplicand_sm80.h create mode 100644 csrc/mctlass/include/mctlass/layout/vector.h create mode 100644 csrc/mctlass/include/mctlass/matrix.h create mode 100644 csrc/mctlass/include/mctlass/matrix_coord.h create mode 100644 csrc/mctlass/include/mctlass/matrix_shape.h create mode 100644 csrc/mctlass/include/mctlass/mctlass.h create mode 100644 csrc/mctlass/include/mctlass/numeric_conversion.h create mode 100644 csrc/mctlass/include/mctlass/numeric_types.h create mode 100644 csrc/mctlass/include/mctlass/pipeline/pipeline.hpp create mode 100644 csrc/mctlass/include/mctlass/pipeline/sm90_pipeline.hpp create mode 100644 csrc/mctlass/include/mctlass/pitch_linear_coord.h create mode 100644 csrc/mctlass/include/mctlass/platform/platform.h create mode 100644 csrc/mctlass/include/mctlass/predicate_vector.h create mode 100644 csrc/mctlass/include/mctlass/quaternion.h create mode 100644 csrc/mctlass/include/mctlass/real.h create mode 100644 csrc/mctlass/include/mctlass/reduction/device/reduce_split_k.h create mode 100644 csrc/mctlass/include/mctlass/reduction/device/tensor_reduce.h create mode 100644 csrc/mctlass/include/mctlass/reduction/device/tensor_reduce_affine_contiguous.h create mode 100644 csrc/mctlass/include/mctlass/reduction/device/tensor_reduce_affine_strided.h create mode 100644 csrc/mctlass/include/mctlass/reduction/kernel/reduce_softmax_final.h create mode 100644 csrc/mctlass/include/mctlass/reduction/kernel/reduce_split_k.h create mode 100644 csrc/mctlass/include/mctlass/reduction/kernel/tensor_reduce_affine_contiguous.h create mode 100644 csrc/mctlass/include/mctlass/reduction/kernel/tensor_reduce_affine_strided.h create mode 100644 csrc/mctlass/include/mctlass/reduction/thread/reduce.h create mode 100644 csrc/mctlass/include/mctlass/reduction/thread/reduction_operators.h create mode 100644 csrc/mctlass/include/mctlass/reduction/threadblock_swizzle.h create mode 100644 csrc/mctlass/include/mctlass/relatively_equal.h create mode 100644 csrc/mctlass/include/mctlass/semaphore.h create mode 100644 csrc/mctlass/include/mctlass/subbyte_reference.h create mode 100644 csrc/mctlass/include/mctlass/tensor_coord.h create mode 100644 csrc/mctlass/include/mctlass/tensor_ref.h create mode 100644 csrc/mctlass/include/mctlass/tensor_ref_planar_complex.h create mode 100644 csrc/mctlass/include/mctlass/tensor_view.h create mode 100644 csrc/mctlass/include/mctlass/tensor_view_planar_complex.h create mode 100644 csrc/mctlass/include/mctlass/tfloat32.h create mode 100644 csrc/mctlass/include/mctlass/thread/matrix.h create mode 100644 csrc/mctlass/include/mctlass/trace.h create mode 100644 csrc/mctlass/include/mctlass/transform/collective/sm90_wgmma_transpose.hpp create mode 100644 csrc/mctlass/include/mctlass/transform/pitch_linear_thread_map.h create mode 100644 csrc/mctlass/include/mctlass/transform/thread/transpose.h create mode 100644 csrc/mctlass/include/mctlass/transform/thread/unary_op.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/ell_iterator.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/ell_predicated_tile_access_iterator.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/ell_predicated_tile_iterator.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/predicated_scale_bias_vector_access_iterator.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/predicated_scale_bias_vector_iterator.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator_2dthreadtile.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator_params.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator_triangular_matrix.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_iterator.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_iterator_2dthreadtile.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_iterator_triangular_matrix.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/predicated_vector_access_iterator.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/regular_scale_bias_vector_access_iterator.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear_direct_conv.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op_sm80.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_pitch_linear_2dthreadtile.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_tensor_op.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_tensor_op_sm70.h create mode 100644 csrc/mctlass/include/mctlass/transform/threadblock/vector_iterator.h create mode 100644 csrc/mctlass/include/mctlass/transform/warp/vector_fragment_iterator.h create mode 100644 csrc/mctlass/include/mctlass/uint128.h create mode 100644 csrc/mctlass/include/mctlass/wmma_array.h create mode 100644 csrc/utils/arch.h delete mode 100644 csrc/utils/hdim_switch.h create mode 100644 csrc/utils/host_utils.h create mode 100644 csrc/utils/logger.cpp create mode 100644 csrc/utils/logger.h delete mode 100644 csrc/utils/philox.cuh create mode 100644 csrc/utils/print_parameter.cpp create mode 100644 csrc/utils/print_parameter.h create mode 100644 csrc/utils/process_str.cpp create mode 100644 csrc/utils/process_str.h delete mode 100644 tests/test_flash_mla.py diff --git a/LICENSE b/LICENSE index 87f733d..6ecdcb5 100644 --- a/LICENSE +++ b/LICENSE @@ -1,8 +1,8 @@ MetaX-MACA/FlashMLA是deepseek-ai/FlashMLA算法在MXMACA软件栈及FlashAttention-2（2.6.3版）的实现。MetaX-MACA/FlashMLA（以下简称“本软件”）适用MIT许可证。本软件亦包含第三方开源组件，后者采用的开源许可证将在下文列出。 MetaX-MACA/FlashMLA is the implementation of the deepseek-ai/FlashMLA algorithm on the MXMACA software stack and FlashAttention-2 (version 2.6.3). MetaX-MACA/FlashMLA (“This software”) is licensed under MIT. This software also contains third-party open source components, the open source licenses of which are listed below. -Copyright © 2025 MetaX Integrated Circuits (Shanghai) Co., Ltd. -版权所有©2025 沐曦集成电路（上海）股份有限公司。 +Copyright © 2025-2026 MetaX Integrated Circuits (Shanghai) Co., Ltd. +版权所有©2026 沐曦集成电路（上海）股份有限公司。 MIT License Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: @@ -18,7 +18,7 @@ This software also contains code from deepseek-ai /FlashMLA（https://github.com deepseek-ai /FlashMLA MIT License -Copyright (c) 2025 DeepSeek +Copyright (c) 2025-2026 DeepSeek Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal @@ -101,4 +101,3 @@ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. - diff --git a/benchmark/bench_flash_mla.py b/benchmark/bench_flash_mla.py deleted file mode 100644 index 2b59f8e..0000000 --- a/benchmark/bench_flash_mla.py +++ /dev/null @@ -1,520 +0,0 @@ -# MLA Triton kernel is from: https://github.com/monellz/vllm/commit/feebaa7c063be6bfb590a876741aeef1c5f58cf8#diff-7b2e1c9032522f7266051b9887246a65753871dfb3625a258fee40109fe6e87a -import argparse -import math -import random - -import flashinfer -import torch -import triton -import triton.language as tl - -# pip install flashinfer-python -from flash_mla import flash_mla_with_kvcache, get_mla_metadata - - -def scaled_dot_product_attention(query, key, value, h_q, h_kv, is_causal=False): - query = query.float() - key = key.float() - value = value.float() - key = key.repeat_interleave(h_q // h_kv, dim=0) - value = value.repeat_interleave(h_q // h_kv, dim=0) - attn_weight = query @ key.transpose(-2, -1) / math.sqrt(query.size(-1)) - if is_causal: - s_q = query.shape[-2] - s_k = key.shape[-2] - attn_bias = torch.zeros(s_q, s_k, dtype=query.dtype) - temp_mask = torch.ones(s_q, s_k, dtype=torch.bool).tril(diagonal=s_k - s_q) - attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf")) - attn_bias.to(query.dtype) - attn_weight += attn_bias - lse = attn_weight.logsumexp(dim=-1) - attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32) - return attn_weight @ value, lse - - -@torch.inference_mode() -def run_torch_mla(q, block_table, blocked_k, max_seqlen_pad, block_size, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype): - for i in range(b): - blocked_k.view(b, max_seqlen_pad, h_kv, d)[i, cache_seqlens[i].item():] = float("nan") - blocked_v = blocked_k[..., :dv] - - def ref_mla(): - out = torch.empty(b, s_q, h_q, dv, dtype=torch.float32) - lse = torch.empty(b, h_q, s_q, dtype=torch.float32) - for i in range(b): - begin = i * max_seqlen_pad - end = begin + cache_seqlens[i] - O, LSE = scaled_dot_product_attention( - q[i].transpose(0, 1), - blocked_k.view(-1, h_kv, d)[begin:end].transpose(0, 1), - blocked_v.view(-1, h_kv, dv)[begin:end].transpose(0, 1), - h_q, h_kv, - is_causal=causal, - ) - out[i] = O.transpose(0, 1) - lse[i] = LSE - return out, lse - - out_torch, lse_torch = ref_mla() - t = triton.testing.do_bench(ref_mla) - return out_torch, lse_torch, t - -@torch.inference_mode() -def run_flash_mla(q, block_table, blocked_k, max_seqlen_pad, block_size, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype): - for i in range(b): - blocked_k.view(b, max_seqlen_pad, h_kv, d)[i, cache_seqlens[i].item():] = float("nan") - blocked_v = blocked_k[..., :dv] - - tile_scheduler_metadata, num_splits = get_mla_metadata(cache_seqlens, s_q * h_q // h_kv, h_kv) - - def flash_mla(): - return flash_mla_with_kvcache( - q, blocked_k, block_table, cache_seqlens, dv, - tile_scheduler_metadata, num_splits, causal=causal, - ) - - out_flash, lse_flash = flash_mla() - t = triton.testing.do_bench(flash_mla) - return out_flash, lse_flash, t - - -@torch.inference_mode() -def run_flash_infer(q, block_table, blocked_k, max_seqlen_pad, block_size, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype): - - for i in range(b): - blocked_k.view(b, max_seqlen_pad, h_kv, d)[i, cache_seqlens[i].item():] = float("nan") - - assert d > dv, "mla with rope dim should be larger than no rope dim" - q_nope, q_pe = q[..., :dv].contiguous(), q[..., dv:].contiguous() - blocked_k_nope, blocked_k_pe = blocked_k[..., :dv].contiguous(), blocked_k[..., dv:].contiguous() - - - kv_indptr = [0] - kv_indices = [] - for i in range(b): - seq_len = cache_seqlens[i] - assert seq_len > 0 - num_blocks = (seq_len + block_size - 1) // block_size - kv_indices.extend(block_table[i, :num_blocks]) - kv_indptr.append(kv_indptr[-1] + num_blocks) - for seq_len in cache_seqlens[1:]: - kv_indptr.append((seq_len + block_size - 1) // block_size + kv_indptr[-1]) - - q_indptr = torch.arange(0, b + 1).int() * s_q - kv_indptr = torch.tensor(kv_indptr, dtype=torch.int32) - kv_indices = torch.tensor(kv_indices, dtype=torch.int32) - - mla_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper( - torch.empty(128 * 1024 * 1024, dtype=torch.int8), - backend="fa3" - ) - mla_wrapper.plan( - q_indptr, - kv_indptr, - kv_indices, - cache_seqlens, - h_q, - dv, - d-dv, - block_size, - causal, - 1 / math.sqrt(d), - q.dtype, - blocked_k.dtype, - ) - - def flash_infer(): - output, lse = mla_wrapper.run(q_nope.view(-1, h_q, dv), q_pe.view(-1, h_q, d-dv), blocked_k_nope, blocked_k_pe, return_lse=True) - return output.view(b, -1, h_q, dv), lse.view(b, h_q, 1) - - out_flash, lse_flash = flash_infer() - t = triton.testing.do_bench(flash_infer) - return out_flash, lse_flash, t - - -@triton.jit -def _mla_attn_kernel( - Q_nope, - Q_pe, - Kv_c_cache, - K_pe_cache, - Req_to_tokens, - B_seq_len, - O, - sm_scale, - stride_q_nope_bs, - stride_q_nope_h, - stride_q_pe_bs, - stride_q_pe_h, - stride_kv_c_bs, - stride_k_pe_bs, - stride_req_to_tokens_bs, - stride_o_b, - stride_o_h, - stride_o_s, - BLOCK_H: tl.constexpr, - BLOCK_N: tl.constexpr, - NUM_KV_SPLITS: tl.constexpr, - PAGE_SIZE: tl.constexpr, - HEAD_DIM_CKV: tl.constexpr, - HEAD_DIM_KPE: tl.constexpr, -): - cur_batch = tl.program_id(1) - cur_head_id = tl.program_id(0) - split_kv_id = tl.program_id(2) - - cur_batch_seq_len = tl.load(B_seq_len + cur_batch) - - offs_d_ckv = tl.arange(0, HEAD_DIM_CKV) - cur_head = cur_head_id * BLOCK_H + tl.arange(0, BLOCK_H) - offs_q_nope = cur_batch * stride_q_nope_bs + cur_head[:, None] * stride_q_nope_h + offs_d_ckv[None, :] - q_nope = tl.load(Q_nope + offs_q_nope) - - offs_d_kpe = tl.arange(0, HEAD_DIM_KPE) - offs_q_pe = cur_batch * stride_q_pe_bs + cur_head[:, None] * stride_q_pe_h + offs_d_kpe[None, :] - q_pe = tl.load(Q_pe + offs_q_pe) - - e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - float("inf") - e_sum = tl.zeros([BLOCK_H], dtype=tl.float32) - acc = tl.zeros([BLOCK_H, HEAD_DIM_CKV], dtype=tl.float32) - - kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) - split_kv_start = kv_len_per_split * split_kv_id - split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) - - for start_n in range(split_kv_start, split_kv_end, BLOCK_N): - offs_n = start_n + tl.arange(0, BLOCK_N) - kv_page_number = tl.load( - Req_to_tokens + stride_req_to_tokens_bs * cur_batch + offs_n // PAGE_SIZE, - mask=offs_n < split_kv_end, - other=0, - ) - kv_loc = kv_page_number * PAGE_SIZE + offs_n % PAGE_SIZE - offs_k_c = kv_loc[None, :] * stride_kv_c_bs + offs_d_ckv[:, None] - k_c = tl.load(Kv_c_cache + offs_k_c, mask=offs_n[None, :] < split_kv_end, other=0.0) - - qk = tl.dot(q_nope, k_c.to(q_nope.dtype)) - - offs_k_pe = kv_loc[None, :] * stride_k_pe_bs + offs_d_kpe[:, None] - k_pe = tl.load(K_pe_cache + offs_k_pe, mask=offs_n[None, :] < split_kv_end, other=0.0) - - qk += tl.dot(q_pe, k_pe.to(q_pe.dtype)) - qk *= sm_scale - - qk = tl.where(offs_n[None, :] < split_kv_end, qk, float("-inf")) - - v_c = tl.trans(k_c) - - n_e_max = tl.maximum(tl.max(qk, 1), e_max) - re_scale = tl.exp(e_max - n_e_max) - p = tl.exp(qk - n_e_max[:, None]) - acc *= re_scale[:, None] - acc += tl.dot(p.to(v_c.dtype), v_c) - - e_sum = e_sum * re_scale + tl.sum(p, 1) - e_max = n_e_max - offs_o = cur_batch * stride_o_b + cur_head[:, None] * stride_o_h + split_kv_id * stride_o_s + offs_d_ckv[None, :] - tl.store(O + offs_o, acc / e_sum[:, None]) - offs_o_1 = cur_batch * stride_o_b + cur_head * stride_o_h + split_kv_id * stride_o_s + HEAD_DIM_CKV - tl.store(O + offs_o_1, e_max + tl.log(e_sum)) - - -def _mla_attn( - q_nope, - q_pe, - kv_c_cache, - k_pe_cache, - attn_logits, - req_to_tokens, - b_seq_len, - num_kv_splits, - sm_scale, - page_size, -): - batch_size, head_num = q_nope.shape[0], q_nope.shape[1] - head_dim_ckv = q_nope.shape[-1] - head_dim_kpe = q_pe.shape[-1] - - BLOCK_H = 16 - BLOCK_N = 64 - grid = ( - triton.cdiv(head_num, BLOCK_H), - batch_size, - num_kv_splits, - ) - _mla_attn_kernel[grid]( - q_nope, - q_pe, - kv_c_cache, - k_pe_cache, - req_to_tokens, - b_seq_len, - attn_logits, - sm_scale, - # stride - q_nope.stride(0), - q_nope.stride(1), - q_pe.stride(0), - q_pe.stride(1), - kv_c_cache.stride(-2), - k_pe_cache.stride(-2), - req_to_tokens.stride(0), - attn_logits.stride(0), - attn_logits.stride(1), - attn_logits.stride(2), - BLOCK_H=BLOCK_H, - BLOCK_N=BLOCK_N, - NUM_KV_SPLITS=num_kv_splits, - PAGE_SIZE=page_size, - HEAD_DIM_CKV=head_dim_ckv, - HEAD_DIM_KPE=head_dim_kpe, - ) - -@triton.jit -def _mla_softmax_reducev_kernel( - Logits, - B_seq_len, - O, - stride_l_b, - stride_l_h, - stride_l_s, - stride_o_b, - stride_o_h, - NUM_KV_SPLITS: tl.constexpr, - HEAD_DIM_CKV: tl.constexpr, -): - cur_batch = tl.program_id(0) - cur_head = tl.program_id(1) - cur_batch_seq_len = tl.load(B_seq_len + cur_batch) - - offs_d_ckv = tl.arange(0, HEAD_DIM_CKV) - - e_sum = 0.0 - e_max = -float("inf") - acc = tl.zeros([HEAD_DIM_CKV], dtype=tl.float32) - - offs_l = cur_batch * stride_l_b + cur_head * stride_l_h + offs_d_ckv - offs_l_1 = cur_batch * stride_l_b + cur_head * stride_l_h + HEAD_DIM_CKV - - for split_kv_id in range(0, NUM_KV_SPLITS): - kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) - split_kv_start = kv_len_per_split * split_kv_id - split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) - - if split_kv_end > split_kv_start: - logits = tl.load(Logits + offs_l + split_kv_id * stride_l_s) - logits_1 = tl.load(Logits + offs_l_1 + split_kv_id * stride_l_s) - - n_e_max = tl.maximum(logits_1, e_max) - old_scale = tl.exp(e_max - n_e_max) - acc *= old_scale - exp_logic = tl.exp(logits_1 - n_e_max) - acc += exp_logic * logits - - e_sum = e_sum * old_scale + exp_logic - e_max = n_e_max - - tl.store( - O + cur_batch * stride_o_b + cur_head * stride_o_h + offs_d_ckv, - acc / e_sum, - ) - - -def _mla_softmax_reducev( - logits, - o, - b_seq_len, - num_kv_splits, -): - batch_size, head_num, head_dim_ckv = o.shape[0], o.shape[1], o.shape[2] - grid = (batch_size, head_num) - _mla_softmax_reducev_kernel[grid]( - logits, - b_seq_len, - o, - logits.stride(0), - logits.stride(1), - logits.stride(2), - o.stride(0), - o.stride(1), - NUM_KV_SPLITS=num_kv_splits, - HEAD_DIM_CKV=head_dim_ckv, - num_warps=4, - num_stages=2, - ) - -def mla_decode_triton( - q_nope, - q_pe, - kv_c_cache, - k_pe_cache, - o, - req_to_tokens, - b_seq_len, - attn_logits, - num_kv_splits, - sm_scale, - page_size, -): - assert num_kv_splits == attn_logits.shape[2] - _mla_attn( - q_nope, - q_pe, - kv_c_cache, - k_pe_cache, - attn_logits, - req_to_tokens, - b_seq_len, - num_kv_splits, - sm_scale, - page_size, - ) - _mla_softmax_reducev( - attn_logits, - o, - b_seq_len, - num_kv_splits, - ) - - -@torch.inference_mode() -def run_flash_mla_triton(q, block_table, blocked_k, max_seqlen_pad, block_size, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype): - - for i in range(b): - blocked_k.view(b, max_seqlen_pad, h_kv, d)[i, cache_seqlens[i].item():] = float("nan") - blocked_v = blocked_k[..., :dv] - - assert d > dv, "mla with rope dim should be larger than no rope dim" - q_nope, q_pe = q[..., :dv].contiguous(), q[..., dv:].contiguous() - blocked_k_nope, blocked_k_pe = blocked_k[..., :dv].contiguous(), blocked_k[..., dv:].contiguous() - - def flash_mla_triton(): - num_kv_splits = 32 - o = torch.empty([b * s_q, h_q, dv]) - attn_logits = torch.empty([b * s_q, h_q, num_kv_splits, dv + 1]) - mla_decode_triton(q_nope.view(-1, h_q, dv), q_pe.view(-1, h_q, d-dv), blocked_k_nope.view(-1, dv), blocked_k_pe.view(-1, d-dv), o, block_table, cache_seqlens, attn_logits, num_kv_splits, 1 / math.sqrt(d), block_size) - return o.view([b, s_q, h_q, dv]) - - out_flash = flash_mla_triton() - t = triton.testing.do_bench(flash_mla_triton) - return out_flash, None, t - - -FUNC_TABLE = { - "torch": run_torch_mla, - "flash_mla": run_flash_mla, - "flash_infer": run_flash_infer, - "flash_mla_triton": run_flash_mla_triton, -} - -def compare_ab(baseline, target, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype): - print(f"comparing {baseline} vs {target}: {b=}, {s_q=}, mean_seqlens={cache_seqlens.float().mean()}, {h_q=}, {h_kv=}, {d=}, {dv=}, {causal=}, {dtype=}") - device = torch.device("cuda:0") - torch.set_default_dtype(dtype) - torch.set_default_device(device) - torch.cuda.set_device(device) - torch.manual_seed(0) - random.seed(0) - assert baseline in FUNC_TABLE - assert target in FUNC_TABLE - baseline_func = FUNC_TABLE[baseline] - target_func = FUNC_TABLE[target] - - total_seqlens = cache_seqlens.sum().item() - mean_seqlens = cache_seqlens.float().mean().int().item() - max_seqlen = cache_seqlens.max().item() - max_seqlen_pad = triton.cdiv(max_seqlen, 256) * 256 - # print(f"{total_seqlens=}, {mean_seqlens=}, {max_seqlen=}") - - q = torch.randn(b, s_q, h_q, d) - block_size = 64 - block_table = torch.arange(b * max_seqlen_pad // block_size, dtype=torch.int32).view(b, max_seqlen_pad // block_size) - blocked_k = torch.randn(block_table.numel(), block_size, h_kv, d) - - out_a, lse_a, perf_a = baseline_func(q, block_table, blocked_k, max_seqlen_pad, block_size, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype) - out_b, lse_b, perf_b = target_func(q, block_table, blocked_k, max_seqlen_pad, block_size, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype) - - torch.testing.assert_close(out_b.float(), out_a.float(), atol=1e-2, rtol=1e-2), "out" - if target not in ["flash_infer", "flash_mla_triton"]: - # flash_infer has a different lse return value - # flash_mla_triton doesn't return lse - torch.testing.assert_close(lse_b.float(), lse_a.float(), atol=1e-2, rtol=1e-2), "lse" - - FLOPS = s_q * total_seqlens * h_q * (d + dv) * 2 - bytes = (total_seqlens * h_kv * d + b * s_q * h_q * d + b * s_q * h_q * dv) * (torch.finfo(dtype).bits // 8) - print(f"perf {baseline}: {perf_a:.3f} ms, {FLOPS / 10 ** 9 / perf_a:.0f} TFLOPS, {bytes / 10 ** 6 / perf_a:.0f} GB/s") - print(f"perf {target}: {perf_b:.3f} ms, {FLOPS / 10 ** 9 / perf_b:.0f} TFLOPS, {bytes / 10 ** 6 / perf_b:.0f} GB/s") - return bytes / 10 ** 6 / perf_a, bytes / 10 ** 6 / perf_b - - -def compare_a(target, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype): - print(f"{target}: {b=}, {s_q=}, mean_seqlens={cache_seqlens.float().mean()}, {h_q=}, {h_kv=}, {d=}, {dv=}, {causal=}, {dtype=}") - torch.set_default_dtype(dtype) - device = torch.device("cuda:0") - torch.set_default_device(device) - torch.cuda.set_device(device) - torch.manual_seed(0) - random.seed(0) - assert target in FUNC_TABLE - target_func = FUNC_TABLE[target] - - total_seqlens = cache_seqlens.sum().item() - mean_seqlens = cache_seqlens.float().mean().int().item() - max_seqlen = cache_seqlens.max().item() - max_seqlen_pad = triton.cdiv(max_seqlen, 256) * 256 - # print(f"{total_seqlens=}, {mean_seqlens=}, {max_seqlen=}") - - q = torch.randn(b, s_q, h_q, d) - block_size = 64 - block_table = torch.arange(b * max_seqlen_pad // block_size, dtype=torch.int32).view(b, max_seqlen_pad // block_size) - blocked_k = torch.randn(block_table.numel(), block_size, h_kv, d) - - out_b, lse_b, perf_b = target_func(q, block_table, blocked_k, max_seqlen_pad, block_size, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype) - - FLOPS = s_q * total_seqlens * h_q * (d + dv) * 2 - bytes = (total_seqlens * h_kv * d + b * s_q * h_q * d + b * s_q * h_q * dv) * (torch.finfo(dtype).bits // 8) - print(f"perf {target}: {perf_b:.3f} ms, {FLOPS / 10 ** 9 / perf_b:.0f} TFLOPS, {bytes / 10 ** 6 / perf_b:.0f} GB/s") - return bytes / 10 ** 6 / perf_b - - -available_targets = [ - "torch", - "flash_mla", - "flash_infer", - "flash_mla_triton", -] - -shape_configs = [ - {"b": batch, "s_q": 1, "cache_seqlens": torch.tensor([seqlen + 2 * i for i in range(batch)], dtype=torch.int32, device="cuda"), "h_q": head, "h_kv": 1, "d": 512+64, "dv": 512, "causal": True, "dtype": torch.bfloat16} - for batch in [128] for seqlen in [1024, 2048, 4096, 8192, 8192*2, 8192*4] for head in [128] -] - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--baseline", type=str, default="torch") - parser.add_argument("--target", type=str, default="flash_mla") - parser.add_argument("--all", action="store_true") - parser.add_argument("--one", action="store_true") - parser.add_argument("--compare", action="store_true") - args = parser.parse_args() - return args - - -if __name__ == "__main__": - args = get_args() - benchmark_type = "all" if args.all else f"{args.baseline}_vs_{args.target}" if args.compare else args.target - with open(f"{benchmark_type}_perf.csv", "w") as fout: - fout.write("name,batch,seqlen,head,bw\n") - for shape in shape_configs: - if args.all: - for target in available_targets: - perf = compare_a(target, shape["b"], shape["s_q"], shape["cache_seqlens"], shape["h_q"], shape["h_kv"], shape["d"], shape["dv"], shape["causal"], shape["dtype"]) - fout.write(f'{target},{shape["b"]},{shape["cache_seqlens"].float().mean().cpu().item():.0f},{shape["h_q"]},{perf:.0f}\n') - elif args.compare: - perfa, prefb = compare_ab(args.baseline, args.target, shape["b"], shape["s_q"], shape["cache_seqlens"], shape["h_q"], shape["h_kv"], shape["d"], shape["dv"], shape["causal"], shape["dtype"]) - fout.write(f'{args.baseline},{shape["b"]},{shape["cache_seqlens"].float().mean().cpu().item():.0f},{shape["h_q"]},{perfa:.0f}\n') - fout.write(f'{args.target},{shape["b"]},{shape["cache_seqlens"].float().mean().cpu().item():.0f},{shape["h_q"]},{prefb:.0f}\n') - elif args.one: - perf = compare_a(args.target, shape["b"], shape["s_q"], shape["cache_seqlens"], shape["h_q"], shape["h_kv"], shape["d"], shape["dv"], shape["causal"], shape["dtype"]) - fout.write(f'{args.target},{shape["b"]},{shape["cache_seqlens"].float().mean().cpu().item():.0f},{shape["h_q"]},{perf:.0f}\n') \ No newline at end of file diff --git a/benchmark/visualize.py b/benchmark/visualize.py deleted file mode 100644 index c1fb37e..0000000 --- a/benchmark/visualize.py +++ /dev/null @@ -1,29 +0,0 @@ -import argparse - -import matplotlib.pyplot as plt -import pandas as pd - - -def parse_args(): - parser = argparse.ArgumentParser(description='Visualize benchmark results') - parser.add_argument('--file', type=str, default='all_perf.csv', - help='Path to the CSV file with benchmark results (default: all_perf.csv)') - return parser.parse_args() - -args = parse_args() -file_path = args.file - -df = pd.read_csv(file_path) - -names = df['name'].unique() - -for name in names: - subset = df[df['name'] == name] - plt.plot(subset['seqlen'], subset['bw'], label=name) - -plt.title('bandwidth') -plt.xlabel('seqlen') -plt.ylabel('bw (GB/s)') -plt.legend() - -plt.savefig(f'{file_path.split(".")[0].split("/")[-1]}_bandwidth_vs_seqlen.png') \ No newline at end of file diff --git a/csrc/flash_api/flash_api.cpp b/csrc/flash_api/flash_api.cpp index ff5a17f..ff284b6 100644 --- a/csrc/flash_api/flash_api.cpp +++ b/csrc/flash_api/flash_api.cpp @@ -10,122 +10,90 @@ #include "flash_mla.h" #include "static_switch.h" -#include "run_mha.h" +#include "run_mla.h" +#include "host_utils.h" #define CHECK_DEVICE(x) TORCH_CHECK(x.is_cuda(), #x " must be on CUDA") #define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")") #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") -// Find the number of splits that maximizes the occupancy. For example, if we have -// batch * n_heads = 48 and we have 108 SMs, having 2 splits (efficiency = 0.89) is -// better than having 3 splits (efficiency = 0.67). However, we also don't want too many -// splits as that would incur more HBM reads/writes. -// So we find the best efficiency, then find the smallest number of splits that gets 85% -// of the best efficiency. -int num_splits_heuristic(int batch_nheads_mblocks, int num_SMs, int num_n_blocks, int max_splits, float balance_weight) { - // If we have enough to almost fill the SMs, then just use 1 split - // if (batch_nheads_mblocks >= 0.9f * num_SMs) { return 1; } - max_splits = std::min({max_splits, num_SMs, num_n_blocks}); - float max_efficiency = 0.f; - std::vector efficiency; - efficiency.reserve(max_splits); - auto ceildiv = [](int a, int b) { return (a + b - 1) / b; }; - // Some splits are not eligible. For example, if we have 64 blocks and choose 11 splits, - // we'll have 6 * 10 + 4 blocks. If we choose 12 splits, we'll have 6 * 11 + (-2) blocks - // (i.e. it's 11 splits anyway). - // So we check if the number of blocks per split is the same as the previous num_splits. - auto is_split_eligible = [&ceildiv, &num_n_blocks](int num_splits) { - return num_splits == 1 || ceildiv(num_n_blocks, num_splits) != ceildiv(num_n_blocks, num_splits - 1); - }; - for (int num_splits = 1; num_splits <= max_splits; num_splits++) { - if (!is_split_eligible(num_splits)) { - efficiency.push_back(0.f); +inline int int64_stride_to_int(int64_t orig_stride) { + if (orig_stride > std::numeric_limits::max()) { + TORCH_CHECK(false, "[Sparse TopK Attention] Stride exceeds int32 limit: ", orig_stride); + } + return static_cast(orig_stride); +} + +// Note: should match the kernel dispatch tile size +inline std::pair get_tile_size(int arch, int seqlen_q, bool is_sparse_attn) { + int block_m = 0; + int block_n = 0; + if (is_sparse_attn) { + // xcore1500 use the same kernel with xcore1000 in sparse decode now + block_m = 64, block_n = 16; + } else { + if (arch >= 1500) { + // only support blockM=64 in xcore1500 dense decode now + block_m = 64, block_n = 32; } else { - float n_waves = float(batch_nheads_mblocks * num_splits) / num_SMs; - float eff = n_waves / ceil(n_waves); - // printf("num_splits = %d, eff = %f\n", num_splits, eff); - if (eff > max_efficiency) { max_efficiency = eff; } - efficiency.push_back(eff); + if (seqlen_q >= 64) { + block_m = 64, block_n = 16; + } else if (seqlen_q >= 32) { + block_m = 32, block_n = 16; + } else { + block_m = 16, block_n = 16; + } } } - for (int num_splits = 1; num_splits <= max_splits; num_splits++) { - if (!is_split_eligible(num_splits)) { continue; } - if (efficiency[num_splits - 1] >= balance_weight * max_efficiency) { - // printf("num_splits chosen = %d\n", num_splits); - return num_splits; - } - } - return 1; + return {block_m, block_n}; } -void compute_params_numsplits(mcFlashAttn::Flash_fwd_mla_params ¶ms, const int num_splits){ - auto num_heads = params.h; - auto batch_size = params.b; - auto max_seqlen_k = params.seqlen_k; - auto max_seqlen_q = params.seqlen_q; - auto dprops = at::cuda::getCurrentDeviceProperties(); +struct DecodingAttnImplMeta { + int num_sm_parts; + int fixed_overhead_num_blocks; + int k_block_size; +}; - const int block_n = 16; - const int num_n_blocks = (max_seqlen_k + block_n - 1) / block_n; - const int block_m = max_seqlen_q >= 64 ? 64 : 32; - const int num_m_blocks = (max_seqlen_q + block_m - 1) / block_m; - params.num_splits = num_splits; - - if (num_splits < 1) { - const int AP_nums = dprops->multiProcessorCount; - int block_nums_per_AP = 1; - // TODO: fine tune balance_weight later - float balance_weight = batch_size == 128 ? 0.95 : 0.9; - params.num_splits = num_splits_heuristic(batch_size * num_heads * num_m_blocks, AP_nums * block_nums_per_AP, - num_n_blocks, 128, balance_weight); - } -} - -std::vector -get_mla_metadata( - at::Tensor &seqlens_k, - const int num_heads_per_head_k, - const int num_heads_k +DecodingAttnImplMeta get_attn_impl_meta( + int arch, + int sm_count, + int num_q_tokens_per_head_k, + int h_k, + int block_m, + int block_n, + std::optional h_q_, + bool is_fp8_kvcache, + bool is_sparse_attn ) { - // This should match the logic in the MLA kernel. - static constexpr int block_size_m = 64; - static constexpr int block_size_n = 64; - static constexpr int fixed_overhead_num_blocks = 5; - - CHECK_DEVICE(seqlens_k); - TORCH_CHECK(seqlens_k.is_contiguous()); - TORCH_CHECK(seqlens_k.dtype() == torch::kInt32); - - int batch_size = seqlens_k.size(0); - int *seqlens_k_ptr = seqlens_k.data_ptr(); - auto options = seqlens_k.options(); - - auto dprops = at::cuda::getCurrentDeviceProperties(); - int sm_count = dprops->multiProcessorCount; - int num_sm_parts = sm_count / num_heads_k / mctlass::ceil_div(num_heads_per_head_k, block_size_m); - - auto tile_scheduler_metadata = torch::empty({num_sm_parts, TileSchedulerMetaDataSize}, options); - auto num_splits = torch::empty({batch_size + 1}, options); - int *tile_scheduler_metadata_ptr = tile_scheduler_metadata.data_ptr(); - int *num_splits_ptr = num_splits.data_ptr(); - - at::cuda::CUDAGuard device_guard{(char)seqlens_k.get_device()}; - auto stream = at::cuda::getCurrentCUDAStream().stream(); - Mla_metadata_params params = {}; - params.seqlens_k_ptr = seqlens_k_ptr; - params.tile_scheduler_metadata_ptr = tile_scheduler_metadata_ptr; - params.num_splits_ptr = num_splits_ptr; - params.batch_size = batch_size; - params.block_size_n = block_size_n; - params.fixed_overhead_num_blocks = fixed_overhead_num_blocks; - params.num_sm_parts = num_sm_parts; - // get_mla_metadata_func(params, stream); - - return {tile_scheduler_metadata, num_splits}; + if (is_sparse_attn) { + if (is_fp8_kvcache) { + TORCH_CHECK(false, "Sparse fp8 MLA is not supported."); + } else { + // Sparse BF16 MLA + TORCH_CHECK(h_q_.has_value()); + int h_q = h_q_.value(); + TORCH_CHECK(h_q % h_k == 0, "h_k must be divisible by h_q."); + int s_q = num_q_tokens_per_head_k * h_k / h_q; + // BF16/FP16 + Sparse MLA + return { + std::max((sm_count/2) / h_k / (mctlass::ceil_div(h_q/h_k, 2*64) * s_q), 1), + 5, + block_n // block_n + }; + } + } else { + TORCH_CHECK(!is_fp8_kvcache, "FP8 KV Cache is not supported."); + // Dense BF16/FP8 MLA + return { + std::max(sm_count / h_k / mctlass::ceil_div(num_q_tokens_per_head_k, block_m), 1), + 5, + block_n, + }; + } } std::vector -mha_fwd_kvcache_mla( +fwd_kvcache_mla( at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size const at::Tensor &kcache, // num_blocks x page_block_size x num_heads_k x head_size c10::optional &vcache_, // num_blocks x page_block_size x num_heads_k x head_size_v @@ -135,17 +103,23 @@ mha_fwd_kvcache_mla( const float softmax_scale, bool is_causal, const at::Tensor &tile_scheduler_metadata, // num_sm_parts x TileSchedulerMetaDataSize - const at::Tensor &num_splits // batch_size + 1 + const at::Tensor &num_splits, // batch_size + 1 + bool is_fp8_kvcache, // fp8 kvcache=False + c10::optional &indices, // None, or batch_size x seqlen_q x topk + c10::optional &indices_all_valid_per_q, // batch_size x seqlen_q x 1, per-query flag indicating whether all top-k indices for each query token are valid. + int const cp_world_size, // context parallelism (cp) world size + int const cp_rank, // cp rank + c10::optional &cp_tot_seqused_k_ // b. total seqused_k in cp world ) { - auto dprops = at::cuda::getCurrentDeviceProperties(); - bool is_sm90 = dprops->major == 9 && dprops->minor == 0; - // TORCH_CHECK(is_sm90); + auto dprops = flash::mcGetCurrentDeviceProperties(); + int arch = dprops.major * 100 + dprops.minor; at::Tensor vcache = vcache_.has_value() ? vcache_.value() : kcache; auto q_dtype = q.dtype(); TORCH_CHECK(q_dtype == torch::kBFloat16 || q_dtype == torch::kFloat16); TORCH_CHECK(kcache.dtype() == q_dtype, "query and key must have the same dtype"); + TORCH_CHECK(!is_fp8_kvcache, "flash mla with kvcache api not support fp8 now"); CHECK_DEVICE(q); CHECK_DEVICE(kcache); CHECK_DEVICE(vcache); @@ -157,6 +131,14 @@ mha_fwd_kvcache_mla( TORCH_CHECK(block_table.dtype() == torch::kInt32, "block_table must have dtype torch.int32"); TORCH_CHECK(block_table.stride(-1) == 1, "block_table must have contiguous last dimension"); + bool is_sparse_attn = indices.has_value(); + int topk = is_sparse_attn ? indices->size(-1) : -1; + TORCH_CHECK(!is_sparse_attn || indices->dtype() == torch::kInt32, "indices must have dtype int32"); + TORCH_CHECK(!is_sparse_attn || indices->stride(-1) == 1, "indices must have contiguous last dimension"); + TORCH_CHECK(!is_sparse_attn || indices_all_valid_per_q->dtype() == torch::kBool, "indices_all_valid_per_q must have dtype bool"); + TORCH_CHECK(!is_sparse_attn || indices_all_valid_per_q->stride(-1) == 1, "indices_all_valid_per_q must have contiguous last dimension"); + + const auto sizes = q.sizes(); const int batch_size = sizes[0]; const int seqlen_q_ori = sizes[1]; @@ -176,6 +158,9 @@ mha_fwd_kvcache_mla( const int ngroups = num_heads_ori / num_heads_k; const int seqlen_q = seqlen_q_ori * ngroups; const int num_heads = num_heads_k; + if (is_sparse_attn){ + TORCH_CHECK(num_heads_ori >= 64 || seqlen_q_ori == 1, "sparse decoding head q must greter than 64 when seqlen q > 1"); + } q = q.view({batch_size, seqlen_q_ori, num_heads_k, ngroups, head_size}).transpose(2, 3) .reshape({batch_size, seqlen_q, num_heads, head_size}); @@ -191,6 +176,15 @@ mha_fwd_kvcache_mla( CHECK_CONTIGUOUS(seqlens_k); CHECK_SHAPE(seqlens_k, batch_size); + if (cp_tot_seqused_k_.has_value()) { + auto cp_tot_seqused_k = cp_tot_seqused_k_.value(); + TORCH_CHECK(cp_tot_seqused_k.dtype() == torch::kInt32, "seqused_k must have dtype int32"); + CHECK_DEVICE(cp_tot_seqused_k); CHECK_CONTIGUOUS(cp_tot_seqused_k); + CHECK_SHAPE(cp_tot_seqused_k, batch_size); + } + + + at::cuda::CUDAGuard device_guard{(char)q.get_device()}; auto opts = q.options(); @@ -202,13 +196,15 @@ mha_fwd_kvcache_mla( // Set the sizes. params.b = batch_size; params.seqlen_q = seqlen_q; - params.seqlen_k = seqlens_k.max().cpu().item(); + // params.seqlen_k = seqlens_k.max().cpu().item(); params.cu_seqlens_k = seqlens_k.data_ptr(); params.is_seqlens_k_cumulative = false; // seqlens_k always has value params.h = num_heads; params.h_h_k_ratio = num_heads / num_heads_k; params.ngroups = ngroups; params.is_causal = is_causal; + params.is_sparse_attn = is_sparse_attn; + params.topk = topk; params.d = head_size; params.d_v = head_size_v; params.scale_softmax = softmax_scale; @@ -233,34 +229,49 @@ mha_fwd_kvcache_mla( params.v_head_stride = vcache.stride(-2); params.o_head_stride = out.stride(-2); + // indices ptr + params.indices_ptr = is_sparse_attn ? indices->data_ptr() : nullptr; + params.indices_batch_stride = is_sparse_attn ? indices->stride(0) : 0; + params.indices_row_stride = is_sparse_attn ? indices->stride(1) : 0; + + params.indices_all_valid_per_q_ptr = is_sparse_attn ? indices_all_valid_per_q->data_ptr() : nullptr; + params.indices_all_valid_per_q_batch_stride = is_sparse_attn ? indices_all_valid_per_q->stride(0) : 0; + params.indices_all_valid_per_q_row_stride = is_sparse_attn ? indices_all_valid_per_q->stride(1) : 0; + params.block_table = block_table.data_ptr(); params.block_table_batch_stride = block_table.stride(0); params.page_block_size = page_block_size; + params.arch = arch; + + params.cp_world_size = cp_world_size; + params.cp_rank = cp_rank; + params.cp_tot_seqused_k = cp_tot_seqused_k_.has_value() ? cp_tot_seqused_k_->data_ptr() : nullptr; + TORCH_CHECK(cp_world_size > 0, "cp_world_size must be positive, required by downstream unified code path. Use 1 if CP is not enabled."); + TORCH_CHECK(cp_world_size != 1 || cp_rank == 0, "When context parallelism is disabled, cp_rank must be zero"); + TORCH_CHECK(cp_world_size == 1 || cp_tot_seqused_k_.has_value(), "cp_tot_seqused_k_ must be provided when context parallelism is enabled."); + + TORCH_CHECK(num_splits.dtype() == torch::kInt32, "num_splits must have dtype int32"); + // printf("num_splits%d",num_splits); + CHECK_DEVICE(num_splits); + CHECK_CONTIGUOUS(num_splits); TORCH_CHECK(tile_scheduler_metadata.dtype() == torch::kInt32, "tile_scheduler_metadata must have dtype int32"); TORCH_CHECK(tile_scheduler_metadata.size(1) == TileSchedulerMetaDataSize); CHECK_DEVICE(tile_scheduler_metadata); CHECK_CONTIGUOUS(tile_scheduler_metadata); - // params.tile_scheduler_metadata_ptr = tile_scheduler_metadata.data_ptr(); - // params.num_sm_parts = tile_scheduler_metadata.size(0); - TORCH_CHECK(num_splits.dtype() == torch::kInt32, "num_splits must have dtype int32"); - CHECK_DEVICE(num_splits); - CHECK_CONTIGUOUS(num_splits); - // params.num_splits_ptr = num_splits.data_ptr(); - - const int max_num_splits = 128; - // TODO: enable get_mla_mate_data for load balance - compute_params_numsplits(params, 0); - TORCH_CHECK(params.num_splits <= max_num_splits, "num_splits must less than or equal to 128"); - at::Tensor softmax_lse_accum = torch::empty({params.num_splits, batch_size, num_heads, seqlen_q}, opts.dtype(torch::kFloat32)); - at::Tensor out_accum = torch::empty({params.num_splits, batch_size, num_heads, seqlen_q, head_size_v}, opts.dtype(torch::kFloat32)); + params.tile_scheduler_metadata_ptr = tile_scheduler_metadata.data_ptr(); + params.num_sm_parts = tile_scheduler_metadata.size(0); + params.num_splits_ptr = num_splits.data_ptr(); + at::Tensor softmax_lse_accum = torch::empty({batch_size + params.num_sm_parts, num_heads, seqlen_q}, opts.dtype(at::kFloat)); + at::Tensor out_accum = torch::empty({batch_size + params.num_sm_parts, num_heads, seqlen_q, head_size_v}, opts.dtype(at::kFloat)); params.softmax_lseaccum_ptr = softmax_lse_accum.data_ptr(); params.oaccum_ptr = out_accum.data_ptr(); + auto stream = at::cuda::getCurrentCUDAStream().stream(); TORCH_CHECK(head_size == 576); params.is_bf16 = q_dtype == torch::kBFloat16; - run_mha_fwd(params,stream, /*force_split_kernel*/true); + run_mla_fwd(params, stream); out = out.view({batch_size, seqlen_q_ori, ngroups, num_heads_k, head_size_v}).transpose(2, 3) .reshape({batch_size, seqlen_q_ori, num_heads_ori, head_size_v}); softmax_lse = softmax_lse.view({batch_size, num_heads_k, seqlen_q_ori, ngroups}).transpose(2, 3) @@ -268,9 +279,148 @@ mha_fwd_kvcache_mla( return {out, softmax_lse}; } -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.doc() = "FlashAttention"; - //FlashMLA - m.def("get_mla_metadata", &get_mla_metadata); - m.def("fwd_kvcache_mla", &mha_fwd_kvcache_mla); + +std::vector sparse_prefill_fwd( + const at::Tensor &q, + const at::Tensor &kv, + const at::Tensor &indices, + float sm_scale, + int d_v, + const at::Tensor &indices_all_valid_per_q +) { + auto dprops = flash::mcGetCurrentDeviceProperties(); + int arch = dprops.major * 100 + dprops.minor; + + CHECK_DEVICE(q); + CHECK_DEVICE(kv); + CHECK_DEVICE(indices); + CHECK_DEVICE(indices_all_valid_per_q); + + TORCH_CHECK(q.dtype() == torch::kBFloat16); + TORCH_CHECK(kv.dtype() == torch::kBFloat16); + TORCH_CHECK(indices.dtype() == torch::kInt32); + TORCH_CHECK(indices_all_valid_per_q.dtype() == torch::kBool); + + int s_q = q.size(0); + int s_kv = kv.size(0); + int h_q = q.size(1); + int h_kv = kv.size(1); + int d_qk = q.size(2); + int topk = indices.size(2); + TORCH_CHECK(h_q % 64 == 0 && h_q >= 64); + + CHECK_SHAPE(q, s_q, h_q, d_qk); + CHECK_SHAPE(kv, s_kv, h_kv, d_qk); + CHECK_SHAPE(indices, s_q, h_kv, topk); + CHECK_SHAPE(indices_all_valid_per_q, s_q, 1); + + TORCH_CHECK(q.stride(-1) == 1); + TORCH_CHECK(kv.stride(-1) == 1); + TORCH_CHECK(indices.stride(-1) == 1); + + at::cuda::CUDAGuard device_guard{(char)q.get_device()}; + auto opts = q.options(); + at::Tensor out = torch::empty({s_q, h_q, d_v}, opts); + CHECK_CONTIGUOUS(out); + + at::Tensor buf_attn_score, max_logits, lse, p_sum; + max_logits = torch::empty({s_q, h_q}, opts.dtype(torch::kFloat)); + lse = torch::empty({s_q, h_q}, opts.dtype(torch::kFloat)); + CHECK_CONTIGUOUS(max_logits); + CHECK_CONTIGUOUS(lse); + + SparsePrefillParams params = { + s_q, s_kv, h_q, h_kv, d_qk, d_v, topk, + sm_scale, sm_scale * 1.44269504f, + arch, + (mctlass::bfloat16_t*)q.data_ptr(), + (mctlass::bfloat16_t*)kv.data_ptr(), + (int*)indices.data_ptr(), + (bool*)indices_all_valid_per_q.data_ptr(), + + int64_stride_to_int(q.stride(0)), int64_stride_to_int(q.stride(1)), + int64_stride_to_int(kv.stride(0)), int64_stride_to_int(kv.stride(1)), + int64_stride_to_int(indices.stride(0)), int64_stride_to_int(indices.stride(1)), + int64_stride_to_int(out.stride(0)),int64_stride_to_int(out.stride(1)), + + (mctlass::bfloat16_t*)out.data_ptr(), + (float*)max_logits.data_ptr(), + (float*)lse.data_ptr(), + + at::cuda::getCurrentCUDAStream().stream() + }; + + run_mla_fwd(params); + + return {out, max_logits, lse}; +} + +std::vector +get_mla_decoding_metadata( + at::Tensor &seqlens_k, + const int num_q_tokens_per_head_k, + const int h_k, + const std::optional h_q, + const bool is_fp8_kvcache, + const std::optional topk +) { + auto dprops = flash::mcGetCurrentDeviceProperties(); + int arch = dprops.major * 100 + dprops.minor; + // This should match the logic in the MLA kernel. + const int seqlen_q = num_q_tokens_per_head_k * h_k; + bool is_sparse_attn = topk.has_value(); + const auto [block_size_m, block_size_n] = get_tile_size(arch, seqlen_q, is_sparse_attn); + + CHECK_DEVICE(seqlens_k); + TORCH_CHECK(seqlens_k.is_contiguous()); + TORCH_CHECK(seqlens_k.dtype() == torch::kInt32); + if (is_sparse_attn) + TORCH_CHECK(h_q.has_value(), "num_heads_q must be provided when topk is provided"); + + CHECK_DEVICE(seqlens_k); + TORCH_CHECK(seqlens_k.is_contiguous()); + TORCH_CHECK(seqlens_k.dtype() == torch::kInt32); + + int batch_size = seqlens_k.size(0); + int *seqlens_k_ptr = seqlens_k.data_ptr(); + auto options = seqlens_k.options(); + + int sm_count = dprops.multiProcessorCount; + const char* val = std::getenv("FMLA_SM"); + if(val != nullptr){ + sm_count = std::stoi(val); + } + + DecodingAttnImplMeta attn_impl_meta = get_attn_impl_meta(arch, sm_count, num_q_tokens_per_head_k, h_k, block_size_m, block_size_n, h_q, is_fp8_kvcache, is_sparse_attn); + if(std::getenv("FMLA_LOG")){ + printf("block_size_m %d, num_q_tokens_per_head_k %d, h_k %d, seqlen_q %d, sm_count %d, sm_parts %d \n", + block_size_m, num_q_tokens_per_head_k, h_k, seqlen_q, sm_count, attn_impl_meta.num_sm_parts); + } + + auto tile_scheduler_metadata = torch::empty({attn_impl_meta.num_sm_parts, TileSchedulerMetaDataSize}, options); + auto num_splits = torch::empty({batch_size + 1}, options); + int *tile_scheduler_metadata_ptr = tile_scheduler_metadata.data_ptr(); + int *num_splits_ptr = num_splits.data_ptr(); + + at::cuda::CUDAGuard device_guard{(char)seqlens_k.get_device()}; + auto stream = at::cuda::getCurrentCUDAStream().stream(); + GetDecodingMetadataParams params = {}; + params.seqlens_k_ptr = seqlens_k_ptr; + params.tile_scheduler_metadata_ptr = tile_scheduler_metadata_ptr; + params.num_splits_ptr = num_splits_ptr; + params.batch_size = batch_size; + params.block_size_n = attn_impl_meta.k_block_size; + params.fixed_overhead_num_blocks = attn_impl_meta.fixed_overhead_num_blocks; + params.num_sm_parts = attn_impl_meta.num_sm_parts; + params.topk = is_sparse_attn ? topk.value() : -1; + run_get_mla_metadata_kernel(params, stream); + + return {tile_scheduler_metadata, num_splits}; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.doc() = "FlashMLA"; + m.def("get_mla_metadata", &get_mla_decoding_metadata); + m.def("fwd_kvcache_mla", &fwd_kvcache_mla); + m.def("sparse_prefill_fwd", &sparse_prefill_fwd); } diff --git a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_bf16_True_True_sm80.cu b/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_bf16_True_True_sm80.cu deleted file mode 100644 index c5f6887..0000000 --- a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_bf16_True_True_sm80.cu +++ /dev/null @@ -1,18 +0,0 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - -#include "flash_mla.h" -#include "flash_run_fwd_template_impl.h" -#include - -template void run_flash_splitkv_fwd_template< - 576, - 16, - 16, - 4, - true, - true, - cutlass::bfloat16_t, - false, - 512, - 2 - >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); diff --git a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_bf16_True_True_sm80.cu b/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_bf16_True_True_sm80.cu deleted file mode 100644 index 9ea11cf..0000000 --- a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_bf16_True_True_sm80.cu +++ /dev/null @@ -1,18 +0,0 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - -#include "flash_mla.h" -#include "flash_run_fwd_template_impl.h" -#include - -template void run_flash_splitkv_fwd_template< - 576, - 32, - 16, - 4, - true, - true, - cutlass::bfloat16_t, - false, - 512, - 2 - >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); diff --git a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_fp16_True_True_sm80.cu b/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_fp16_True_True_sm80.cu deleted file mode 100644 index fa53b01..0000000 --- a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_fp16_True_True_sm80.cu +++ /dev/null @@ -1,18 +0,0 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - -#include "flash_mla.h" -#include "flash_run_fwd_template_impl.h" -#include - -template void run_flash_splitkv_fwd_template< - 576, - 32, - 16, - 4, - true, - true, - cutlass::half_t, - false, - 512, - 2 - >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); diff --git a/csrc/flash_api/flash_mla.h b/csrc/flash_api/flash_mla.h index 553e446..b701cb7 100644 --- a/csrc/flash_api/flash_mla.h +++ b/csrc/flash_api/flash_mla.h @@ -7,7 +7,8 @@ #include #include - +#include "mctlass/bfloat16.h" +#include "static_switch.h" constexpr int maxValidBlockSizeM = 128; namespace mcFlashAttn { @@ -30,6 +31,10 @@ struct Qkv_params { index_t q_head_stride; index_t k_head_stride; index_t v_head_stride; + index_t indices_batch_stride; + index_t indices_row_stride; + index_t indices_all_valid_per_q_batch_stride; + index_t indices_all_valid_per_q_row_stride; // The number of heads. int h, h_k; @@ -62,6 +67,8 @@ struct Flash_fwd_mla_params : public Qkv_params { // The dimensions. int b, seqlen_q, seqlen_k, seqlen_knew, d, d_v, seqlen_q_rounded, seqlen_k_rounded, d_rounded, rotary_dim, total_q; int ngroups; + bool is_sparse_attn = false; + int topk; // The scaling factors for the kernel. float scale_softmax; @@ -71,32 +78,15 @@ struct Flash_fwd_mla_params : public Qkv_params { int * __restrict__ cu_seqlens_q; int * __restrict__ cu_seqlens_k; int * __restrict__ leftpad_k; + int *__restrict__ indices_ptr; // [batch, s_q, topk] // If provided, the actual length of each k sequence. int * __restrict__ seqused_k; - int *__restrict__ blockmask; - // The K_new and V_new matrices. void * __restrict__ knew_ptr; void * __restrict__ vnew_ptr; - // The stride between rows of the Q, K and V matrices. - index_t knew_batch_stride; - index_t vnew_batch_stride; - index_t knew_row_stride; - index_t vnew_row_stride; - index_t knew_head_stride; - index_t vnew_head_stride; - - // kv cache dequant - index_t kscale_batch_stride; - index_t vscale_batch_stride; - index_t kscale_row_stride; - index_t vscale_row_stride; - index_t kscale_head_stride; - index_t vscale_head_stride; - // The cos and sin matrices for rotary embedding. void * __restrict__ rotary_cos_ptr; void * __restrict__ rotary_sin_ptr; @@ -115,65 +105,45 @@ struct Flash_fwd_mla_params : public Qkv_params { void *__restrict__ k_scale_ptr; void *__restrict__ v_scale_ptr; - // The dropout probability (probability of keeping an activation). - float p_dropout; - // uint32_t p_dropout_in_uint; - // uint16_t p_dropout_in_uint16_t; - uint8_t p_dropout_in_uint8_t; - // Scale factor of 1 / (1 - p_dropout). float rp_dropout; - float scale_softmax_rp_dropout; - - // Local window size - int window_size_left, window_size_right; - - // ratio of softcapping attention - // S = exp2(log2(e) * softcap * tanh(S * softmax_scale / softcap)) - // only value > 0.0 will take effect - float softcap; - - // Random state. - // at::PhiloxCudaState philox_args; // the RNG seed and offset . uint64_t rng_state_seed = 0; uint64_t rng_state_offset = 0; bool is_bf16; + bool is_fp8 = false; bool is_causal; + bool* indices_all_valid_per_q_ptr; // If is_seqlens_k_cumulative, then seqlen_k is cu_seqlens_k[bidb + 1] - cu_seqlens_k[bidb]. // Otherwise it's cu_seqlens_k[bidb], i.e., we use cu_seqlens_k to store the sequence lengths of K. bool is_seqlens_k_cumulative; - bool is_rotary_interleaved; - int num_splits; // For split-KV version - void * __restrict__ alibi_slopes_ptr; - index_t alibi_slopes_batch_stride; - - // attn_mask support for bert model Jira[C500-21935] - bool has_attn_mask; - void * __restrict__ attn_mask_ptr = nullptr; - index_t attn_mask_batch_stride = 0; - index_t attn_mask_nheads_stride = 0; - index_t attn_mask_row_stride = 0; - index_t attn_mask_col_stride = 1; - - index_t attn_mask_batch_shape = 1; - index_t attn_mask_nheads_shape = 1; - index_t attn_mask_row_shape = 1; - index_t attn_mask_col_shape = 1; - bool unpadded_lse; // For varlen paths: LSE is in [nheads, total_seqlen_q] format instead of [b, nheads, seqlen_q]. bool seqlenq_ngroups_swapped; // q has been transposed from (b, 1, (nheads_kv ngroups), d) to (b, ngroups, nheads_kv, d). int d_value; int d_value_rounded; + int arch; - bool is_support_splitkv = false; + int *__restrict__ tile_scheduler_metadata_ptr; + int num_sm_parts; + int *__restrict__ num_splits_ptr; + + // fp8 params + float* __restrict__ descale_q_ptr = nullptr; + float* __restrict__ descale_k_ptr = nullptr; + + // CP (Context Parallelism) parameters + int cp_world_size; + int cp_rank; + int *__restrict__ cp_tot_seqused_k; + + cudaStream_t stream; }; @@ -187,12 +157,46 @@ struct Flash_launch_params { Flash_launch_params(): is_balance(false),rowblock_parallel(0),block_type(0),performance_mode(false){} }; - } +struct SparsePrefillParams { + int s_q, s_kv, h_q, h_kv, d_qk, d_v, topk; + float sm_scale, sm_scale_div_log2; + int arch; + + // Input tensors + mctlass::bfloat16_t* __restrict__ q_ptr; // [s_q, h_q, d_qk] + mctlass::bfloat16_t* __restrict__ kv_ptr; // [s_kv, h_kv, d_qk] + int* __restrict__ indices_ptr; // [s_q, h_kv, topk] + bool* indices_all_valid_per_q_ptr; // [1] + + // int stride_q_s_q; int stride_q_h_q; + // int stride_kv_s_kv; int stride_kv_h_kv; + int q_row_stride;int q_head_stride; + int k_row_stride;int k_head_stride; + int stride_indices_s_q; int stride_indices_h_kv; + int o_row_stride;int o_head_stride; + // Output tensors + mctlass::bfloat16_t* __restrict__ out_ptr; // [s_q, h_q, d_v] + float* __restrict__ max_logits; // [s_q, h_q] + float* __restrict__ lse_ptr; // [s_q, h_q] + + cudaStream_t stream; +}; + static constexpr int TileSchedulerMetaDataSize = 8; // [begin_idx, begin_seqlen, end_idx, end_seqlen, begin_n_split_idx, _, _, _] +struct GetDecodingMetadataParams { + int *__restrict__ seqlens_k_ptr; + int *__restrict__ tile_scheduler_metadata_ptr; + int *__restrict__ num_splits_ptr; + int batch_size; + int block_size_n; + int fixed_overhead_num_blocks; + int num_sm_parts; + int topk; +}; //////////////////////////////////////////////////////////////////////////////////////////////////// struct Mla_metadata_params { @@ -205,4 +209,4 @@ struct Mla_metadata_params { int num_sm_parts; }; -void get_mla_metadata_func(Mla_metadata_params ¶ms, cudaStream_t stream); +void run_get_mla_metadata_kernel(GetDecodingMetadataParams ¶ms, cudaStream_t stream); diff --git a/csrc/flash_dispatch/flash_fwd_dispatch_template.h b/csrc/flash_dispatch/flash_fwd_dispatch_template.h index 931e1fb..13d65f0 100644 --- a/csrc/flash_dispatch/flash_fwd_dispatch_template.h +++ b/csrc/flash_dispatch/flash_fwd_dispatch_template.h @@ -21,41 +21,131 @@ template< typename elem_type, bool Is_splits = false, int kHeadDimV = kHeadDim, - int Num_Stages = 1 + int Num_Stages = 1, + Arch arch = Arch::xcore1000 > -void run_flash_splitkv_fwd_template(Flash_fwd_mla_params ¶ms, cudaStream_t stream); +void run_flash_splitkv_fwd_mla_template(mcFlashAttn::Flash_fwd_mla_params ¶ms, cudaStream_t stream); + +template< + int kHeadDim, + int kBlockM, + int kBlockN, + int kNWarps, + bool Is_Q_in_regs, + bool Share_Q_K_smem, + typename elem_type, + bool Is_splits = false, + int kHeadDimV = kHeadDim, + int Num_Stages = 1, + Arch arch = Arch::xcore1000 +> +void run_flash_splitkv_fwd_sparse_mla_template(mcFlashAttn::Flash_fwd_mla_params ¶ms, cudaStream_t stream); + +template< + int kHeadDim, + int kBlockM, + int kBlockN, + int kNWarps, + bool Is_Q_in_regs, + bool Share_Q_K_smem, + typename elem_type, + bool Is_splits = false, + int kHeadDimV = kHeadDim, + int Num_Stages = 1, + Arch arch = Arch::xcore1000 +> +void run_flash_mla_sparse_prefill_template(SparsePrefillParams ¶ms, cudaStream_t stream); namespace mcFlashAttn { - template - void run_mha_fwd_splitkv_dispatch(Flash_fwd_mla_params ¶ms, const cudaStream_t stream); + template + void run_mla_fwd_splitkv_dispatch(Flash_fwd_mla_params ¶ms, const cudaStream_t stream); + + template + void run_flash_mla_sparse_prefill_dispatch(SparsePrefillParams ¶ms, const cudaStream_t stream); template<> - inline void run_mha_fwd_splitkv_dispatch<576>(Flash_fwd_mla_params ¶ms, const cudaStream_t stream) { + inline void run_flash_mla_sparse_prefill_dispatch<576, Arch::xcore1000>(SparsePrefillParams ¶ms, const cudaStream_t stream) { constexpr static int HeaddimQK = 576; constexpr static int HeaddimVO = 512; constexpr static int Num_Stages = 2; + constexpr Arch arch = Arch::xcore1000; + assert(params.is_bf16 && "sparse prefill only support bf16"); + constexpr static int kBlockM = 64; + constexpr static int kBlockN = 16; + constexpr static int kNWarps = 8; + run_flash_mla_sparse_prefill_template(params, stream); + } + + template<> + inline void run_flash_mla_sparse_prefill_dispatch<576, Arch::xcore1500>(SparsePrefillParams ¶ms, const cudaStream_t stream) { + constexpr static int HeaddimQK = 576; + constexpr static int HeaddimVO = 512; + + constexpr static int Num_Stages = 2; + constexpr Arch arch = Arch::xcore1500; + assert(params.is_bf16 && "sparse prefill only support bf16"); + constexpr static int kBlockM = 64; + constexpr static int kBlockN = 16; + constexpr static int kNWarps = 8; + run_flash_mla_sparse_prefill_template(params, stream); + } + template<> + inline void run_mla_fwd_splitkv_dispatch<576, Arch::xcore1000>(Flash_fwd_mla_params ¶ms, const cudaStream_t stream) { + constexpr static int HeaddimQK = 576; + constexpr static int HeaddimVO = 512; + + constexpr static int Num_Stages = 2; + constexpr Arch arch = Arch::xcore1000; FP16_SWITCH(!params.is_bf16, [&] { - BOOL_SWITCH(params.num_splits > 1, Is_splits, [&] { + if (!params.is_sparse_attn){ if (params.seqlen_q >= 64) { constexpr static int kBlockM = 64; constexpr static int kBlockN = 16; constexpr static int kNWarps = 8; - run_flash_splitkv_fwd_template(params, stream); - } else if (params.seqlen_q >= 32) { + run_flash_splitkv_fwd_mla_template(params, stream); + } + else if (params.seqlen_q >= 32) { constexpr static int kBlockM = 32; constexpr static int kBlockN = 16; constexpr static int kNWarps = 4; - run_flash_splitkv_fwd_template(params, stream); + run_flash_splitkv_fwd_mla_template(params, stream); } else { constexpr static int kBlockM = 16; constexpr static int kBlockN = 16; constexpr static int kNWarps = 4; - run_flash_splitkv_fwd_template(params, stream); + run_flash_splitkv_fwd_mla_template(params, stream); } - }); + }else{ + constexpr static int kBlockM = 64; + constexpr static int kBlockN = 16; + constexpr static int kNWarps = 8; + run_flash_splitkv_fwd_sparse_mla_template(params, stream); + } }); } + template<> + inline void run_mla_fwd_splitkv_dispatch<576, Arch::xcore1500>(Flash_fwd_mla_params ¶ms, const cudaStream_t stream) { + constexpr static int HeaddimQK = 576; + constexpr static int HeaddimVO = 512; + + constexpr static int Num_Stages = 2; + constexpr Arch arch = Arch::xcore1500; + FP16_SWITCH(!params.is_bf16, [&] { + if (!params.is_sparse_attn) { + // NOTE: only support blockM=64 in xcore1500 now + constexpr static int kBlockM = 64; + constexpr static int kBlockN = 32; + constexpr static int kNWarps = 8; + run_flash_splitkv_fwd_mla_template(params, stream); + } else { + constexpr static int kBlockM = 64; + constexpr static int kBlockN = 16; + constexpr static int kNWarps = 8; + run_flash_splitkv_fwd_sparse_mla_template(params, stream); + } + }); + } } // namespace mcFlashAttn end diff --git a/csrc/flash_dispatch/flash_fwd_launch_template.h b/csrc/flash_dispatch/flash_fwd_launch_template.h index 887576b..a0ed9f7 100644 --- a/csrc/flash_dispatch/flash_fwd_launch_template.h +++ b/csrc/flash_dispatch/flash_fwd_launch_template.h @@ -6,59 +6,143 @@ #include "flash_mla.h" #include "static_switch.h" -#include "flash_fwd_split_kernel.h" -#include "feature/attn_mask.h" +#include "flash_dense_mla_decode_kernel.h" +#include "flash_sparse_mla_decode_kernel.h" +#include "flash_fwd_splitkv_mla_combine_kernel.h" +#include "xcore1000/sparse_prefill_kernel_64x16_8waves_xcore1000.h" +#include "print_parameter.h" using namespace mcFlashAttn; -template -__global__ void flash_fwd_splitkv_kernel(const Flash_fwd_mla_params params, const int num_m_block) { - flash::compute_attn_splitkv(params, num_m_block); +template +__global__ void flash_fwd_splitkv_mla_kernel(const Flash_fwd_mla_params params, const int num_m_block) { + flash::compute_attn_1rowblock_splitkv_mla(params, num_m_block); } - +template +__global__ void flash_fwd_splitkv_sparse_mla_kernel(const Flash_fwd_mla_params params, const int num_m_block) { + flash::compute_attn_1rowblock_splitkv_sparse_mla(params, num_m_block); +} +template +__global__ void sparse_attn_global_fwd_kernel(const SparsePrefillParams params) { + flash::sparse_attn_fwd_kernel(params); +} +//flash-meta combine kernel template -__global__ void flash_fwd_splitkv_combine_kernel(const Flash_fwd_mla_params params) { +__global__ void flash_fwd_splitkv_mla_combine_kernel(const Flash_fwd_mla_params params) { static_assert(Log_max_splits >= 1); - flash::combine_attn_seqk_parallel(params); + flash::combine_attn_seqk_parallel_splitkv_mla(params); } -template -void run_flash_splitkv_fwd(Flash_fwd_mla_params ¶ms, cudaStream_t stream) { - - constexpr size_t smem_size = Kernel_traits::kSmemSize; - const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM; - dim3 grid(num_m_block, params.num_splits > 1 ? params.num_splits : params.b, params.num_splits > 1 ? params.b * params.h : params.h); +template +void run_flash_splitkv_fwd_mla(Flash_fwd_mla_params ¶ms, cudaStream_t stream) { + constexpr int max_smem_size = arch == Arch::xcore1000 ? 64 * 1024 : 128 * 1024; + constexpr int smem_size = std::min(Kernel_traits::kSmemSize, max_smem_size); + // const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM; + const int num_m_block = cute::ceil_div(params.seqlen_q, Kernel_traits::kBlockM); + // dim3 grid(num_m_block, params.num_splits > 1 ? params.num_splits : params.b, params.num_splits > 1 ? params.b * params.h : params.h); + dim3 grid(num_m_block, params.h, params.num_sm_parts); static_assert(Kernel_traits::kHeadDim == 576 && Kernel_traits::kHeadDimV == 512); - const bool is_even_MN = params.cu_seqlens_q == nullptr && params.cu_seqlens_k == nullptr && params.seqlen_k % Kernel_traits::kBlockN == 0 && params.seqlen_q % Kernel_traits::kBlockM == 0; const bool is_even_K = params.d == Kernel_traits::kHeadDim && params.d_v == Kernel_traits::kHeadDimV; - EVENK_SWITCH(is_even_K, IsEvenKConst, [&] { - LOCAL_SWITCH_AND_CONST_PRECOND((!Is_causal), (params.window_size_left >= 0 || params.window_size_right >= 0) && !Is_causal, Is_local, [&] { - BOOL_SWITCH(params.num_splits > 1, Split, [&] { - auto kernel = &flash_fwd_splitkv_kernel; - if (smem_size >= 32 * 1024) { - CUDA_CHECK(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - } - kernel<<>>(params, num_m_block); - CUDA_KERNEL_LAUNCH_CHECK(); - }); + BOOL_SWITCH(params.cp_world_size > 1 && params.cp_tot_seqused_k != nullptr, IsEnableDcp, [&] { + EVENK_SWITCH(is_even_K, IsEvenKConst, [&] { + if (std::getenv("MHA_PRINT_PARA")) { + shape_print(params, Is_causal, "mla_dense_decode"); + } + if (std::getenv("MHA_DEBUG_PARA")){ + debug_print(params, "mla_dense_decode"); + } + auto kernel = &flash_fwd_splitkv_mla_kernel; + if (smem_size >= 32 * 1024) { + CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + kernel<<>>(params, num_m_block); + CUDA_KERNEL_LAUNCH_CHECK(); }); }); - if (params.num_splits > 1) { - // We want kBlockM to be as small as possible for more parallelism. - // With 128 threads we can load 512 elements at a time, so if headdim is divisible by 128, kBlockM = 4. - // If headdim is divisible by 64, then we set kBlockM = 8, etc. - //constexpr static int kBlockM = Kernel_traits::kHeadDim % 128 == 0 ? 4 : (Kernel_traits::kHeadDim % 64 == 0 ? 8 : 16); - constexpr static int kBlockM = Kernel_traits::kHeadDim % 128 == 0 ? 8 : (Kernel_traits::kHeadDim % 64 == 0 ? 16 : 32); + // We want kBlockM to be as small as possible for more parallelism. + // In MLA case head_dim_vo = 512, we will switch different kBlockM for different case to get better performance + COMBINE_BLOCKM_SWITCH(params.b, params.h, params.seqlen_q, kBlockM, [&] { dim3 grid_combine((params.b * params.h * params.seqlen_q + kBlockM - 1) / kBlockM); - const int kNThreads = 256; /*Kernel_traits::kNThreads;*/ EVENK_SWITCH(is_even_K, IsEvenKConst, [&] { - NUMSPLITS_SWITCH(params.num_splits, kLogMaxSplits, [&] { - flash_fwd_splitkv_combine_kernel<<>>(params); + NUMSPLITS_SWITCH(params.num_sm_parts, kLogMaxSplits, [&] { + const int kNThreads = kBlockM == 1 ? 64 : (kBlockM == 2 ? 128 : 256); + flash_fwd_splitkv_mla_combine_kernel<<>>(params); CUDA_KERNEL_LAUNCH_CHECK(); }); }); - } + }); } + +template +void run_flash_splitkv_fwd_sparse_mla(Flash_fwd_mla_params ¶ms, cudaStream_t stream) { + constexpr int max_smem_size = arch == Arch::xcore1000 ? 64 * 1024 : 128 * 1024; + constexpr int smem_size = std::min(Kernel_traits::kSmemSize, max_smem_size); + // const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM; + const int num_m_block = cute::ceil_div(params.seqlen_q, Kernel_traits::kBlockM); + // dim3 grid(num_m_block, params.num_splits > 1 ? params.num_splits : params.b, params.num_splits > 1 ? params.b * params.h : params.h); + dim3 grid(num_m_block, params.h, params.num_sm_parts); + static_assert(Kernel_traits::kHeadDim == 576 && Kernel_traits::kHeadDimV == 512); + + // const bool is_even_MN = params.cu_seqlens_q == nullptr && params.cu_seqlens_k == nullptr && params.seqlen_k % Kernel_traits::kBlockN == 0 && params.seqlen_q % Kernel_traits::kBlockM == 0; + const bool is_even_K = true; // is_even_k is always true in mla case; + const bool is_even_topK = params.topk % Kernel_traits::kBlockN == 0; + BOOL_SWITCH(is_even_topK, IsEvenTopKConst, [&] { + if (std::getenv("MHA_PRINT_PARA")) { + shape_print(params, Is_causal, "mla_sparse_decode"); + } + if (std::getenv("MHA_DEBUG_PARA")){ + debug_print(params, "mla_sparse_decode"); + } + auto kernel = &flash_fwd_splitkv_sparse_mla_kernel; + if (smem_size >= 32 * 1024) { + CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + kernel<<>>(params, num_m_block); + CUDA_KERNEL_LAUNCH_CHECK(); + }); + // We want kBlockM to be as small as possible for more parallelism. + // In MLA case head_dim_vo = 512, we will switch different kBlockM for different case to get better performance + COMBINE_BLOCKM_SWITCH(params.b, params.h, params.seqlen_q, kBlockM, [&] { + dim3 grid_combine((params.b * params.h * params.seqlen_q + kBlockM - 1) / kBlockM); + EVENK_SWITCH(is_even_K, IsEvenKConst, [&] { + NUMSPLITS_SWITCH(params.num_sm_parts, kLogMaxSplits, [&] { + const int kNThreads = kBlockM == 1 ? 64 : (kBlockM == 2 ? 128 : 256); + flash_fwd_splitkv_mla_combine_kernel<<>>(params); + CUDA_KERNEL_LAUNCH_CHECK(); + }); + }); + }); +} + +template +void run_sparse_prefill(SparsePrefillParams ¶ms, cudaStream_t stream) { + constexpr int max_smem_size = arch == Arch::xcore1000 ? 64 * 1024 : 128 * 1024; + constexpr int smem_size = std::min(Kernel_traits::kSmemSize, max_smem_size); + + dim3 grid((params.h_q/Kernel_traits::kBlockM)*params.s_q, 1, 1); + static_assert(Kernel_traits::kHeadDim == 576 && Kernel_traits::kHeadDimV == 512); + + const bool is_even_topK = params.topk % Kernel_traits::kBlockN == 0; + // WARNING: Be aware of the correctness of this condition + BOOL_SWITCH(is_even_topK, IsEvenTopKConst, [&] { + if (std::getenv("MHA_PRINT_PARA")) { + shape_print(params, Is_causal, "mla_sparse_prefill"); + } + if (std::getenv("MHA_DEBUG_PARA")){ + debug_print(params, "mla_sparse_prefill"); + } + auto kernel = &sparse_attn_global_fwd_kernel; + if (smem_size >= 32 * 1024) { + CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + kernel<<>>(params); + CUDA_KERNEL_LAUNCH_CHECK(); + }); + + +} \ No newline at end of file diff --git a/csrc/flash_dispatch/flash_fwd_run_template.h b/csrc/flash_dispatch/flash_fwd_run_template.h new file mode 100644 index 0000000..053e1e1 --- /dev/null +++ b/csrc/flash_dispatch/flash_fwd_run_template.h @@ -0,0 +1,69 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#pragma once + +#include + +#include "flash_fwd_launch_template.h" +#include "flash_mla.h" +#include "static_switch.h" + +template< + int kHeadDim, + int kBlockM, + int kBlockN, + int kNWarps, + bool Is_Q_in_regs, + bool Share_Q_K_smem, + typename elem_type, + bool Is_splits = false, + int kHeadDimV = kHeadDim, + int Num_Stages = 1, + Arch arch = Arch::xcore1000 +> +void run_flash_splitkv_fwd_mla_template(mcFlashAttn::Flash_fwd_mla_params ¶ms, cudaStream_t stream){ + using Kernel_traits = Flash_fwd_kernel_traits; + BOOL_SWITCH(params.is_causal, Is_causal, [&] { + run_flash_splitkv_fwd_mla(params, stream); + }); +} + +template< + int kHeadDim, + int kBlockM, + int kBlockN, + int kNWarps, + bool Is_Q_in_regs, + bool Share_Q_K_smem, + typename elem_type, + bool Is_splits = false, + int kHeadDimV = kHeadDim, + int Num_Stages = 1, + Arch arch = Arch::xcore1000 +> +void run_flash_splitkv_fwd_sparse_mla_template(mcFlashAttn::Flash_fwd_mla_params ¶ms, cudaStream_t stream){ + using Kernel_traits = Flash_fwd_kernel_traits; + BOOL_SWITCH(params.is_causal, Is_causal, [&] { + run_flash_splitkv_fwd_sparse_mla(params, stream); + }); +} + +template< + int kHeadDim, + int kBlockM, + int kBlockN, + int kNWarps, + bool Is_Q_in_regs, + bool Share_Q_K_smem, + typename elem_type, + bool Is_splits = false, + int kHeadDimV = kHeadDim, + int Num_Stages = 1, + Arch arch = Arch::xcore1000 +> +void run_flash_mla_sparse_prefill_template(SparsePrefillParams ¶ms, cudaStream_t stream){ + using Kernel_traits = Flash_fwd_kernel_traits; + + run_sparse_prefill(params, stream); + +} \ No newline at end of file diff --git a/csrc/flash_dispatch/flash_run_fwd_template_impl.h b/csrc/flash_dispatch/flash_run_fwd_template_impl.h deleted file mode 100644 index cf316aa..0000000 --- a/csrc/flash_dispatch/flash_run_fwd_template_impl.h +++ /dev/null @@ -1,28 +0,0 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - -#pragma once - -#include - -#include "flash_fwd_launch_template.h" -#include "flash_mla.h" -#include "static_switch.h" - -template< - int kHeadDim, - int kBlockM, - int kBlockN, - int kNWarps, - bool Is_Q_in_regs, - bool Share_Q_K_smem, - typename elem_type, - bool Is_splits = false, - int kHeadDimV = kHeadDim, - int Num_Stages = 1 -> -void run_flash_splitkv_fwd_template(mcFlashAttn::Flash_fwd_mla_params ¶ms, cudaStream_t stream){ - using Kernel_traits = Flash_fwd_kernel_traits; - BOOL_SWITCH(params.is_causal, Is_causal, [&] { - run_flash_splitkv_fwd(params, stream); - }); -} \ No newline at end of file diff --git a/csrc/flash_kernel/feature/alibi.h b/csrc/flash_kernel/feature/alibi.h deleted file mode 100644 index d43a981..0000000 --- a/csrc/flash_kernel/feature/alibi.h +++ /dev/null @@ -1,76 +0,0 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - -#pragma once - -#include - -#include - -#include -#include - -#include "utils.h" - -namespace flash { - -using namespace cute; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -struct Alibi { - - const float alibi_slope; - const int max_seqlen_k, max_seqlen_q; - - __forceinline__ __device__ Alibi(const float alibi_slope, const int max_seqlen_k, const int max_seqlen_q) - : alibi_slope(alibi_slope) - , max_seqlen_k(max_seqlen_k) - , max_seqlen_q(max_seqlen_q) { - }; - - - template - __forceinline__ __device__ void apply_alibi(Tensor &tensor, - const int col_idx_offset_, - const int row_idx_offset, - const int warp_row_stride, - const int warp_col_stride = 16) { - // tensor has shape (nrow=(1, MMA_M), ncol=(4, MMA_N)) - static_assert(Layout::rank == 2, "Only support 2D Tensor"); - static_assert(decltype(size<0, 0>(tensor))::value == 1); - static_assert(decltype(size<1, 0>(tensor))::value == 4); - const int col_idx_offset = col_idx_offset_ + ((__lane_id() >> 4) << 2); - if constexpr (Is_causal) { // Simpler, we add the same bias vector to all rows - #pragma unroll - for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { - const int col_idx_base = col_idx_offset + nj * warp_col_stride; - #pragma unroll - for (int j = 0; j < size<1, 0>(tensor); ++j) { - const int col_idx = col_idx_base + j; - #pragma unroll - for (int mi = 0; mi < size<0>(tensor); ++mi) { - tensor(mi, make_coord(j, nj)) += alibi_slope * col_idx; - } - } - } - } else { // Bias depends on both row_idx and col_idx - #pragma unroll - for (int mi = 0; mi < size<0, 1>(tensor); ++mi) { - const int row_idx = row_idx_offset + mi * warp_row_stride; - #pragma unroll - for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { - const int col_idx_base = col_idx_offset + nj * warp_col_stride; - #pragma unroll - for (int j = 0; j < size<1, 0>(tensor); ++j) { - const int col_idx = col_idx_base + j; - tensor(make_coord(0, mi), make_coord(j, nj)) -= alibi_slope * abs(row_idx + max_seqlen_k - max_seqlen_q - col_idx); - } - } - } - } - } - -}; - -} // namespace flash diff --git a/csrc/flash_kernel/feature/attn_mask.h b/csrc/flash_kernel/feature/attn_mask.h deleted file mode 100644 index 5056eb5..0000000 --- a/csrc/flash_kernel/feature/attn_mask.h +++ /dev/null @@ -1,166 +0,0 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - -#pragma once - -#include - -#include - -#include -#include - -#include "utils.h" -#include "block_info.h" - -namespace flash { - -using namespace cute; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -inline __host__ __device__ bool use_attn_mask_merge_ldg(Params ¶ms) { - /*merged impl require - 1. bias_col_shape % 4 == 0 - 2. bias_col_stride == 1 - */ - if((params.attn_mask_col_shape % 4) == 0 - && params.attn_mask_col_stride == 1) { - return true; - } else { - return false; - } -} - -template -inline __device__ void apply_attn_mask(Tensor &tensor, - const int col_idx_offset_, - const int max_seqlen_k, - const int row_idx_offset_, - const int max_seqlen_q, - const int warp_row_stride, - const int warp_col_stride, - const float softmax_scale, - T *bias, - const int bias_row_stride, - const int bias_col_stride) { - // tensor has shape (nrow=(1, MMA_M), ncol=(4, MMA_N) - CUTE_STATIC_ASSERT_V((size<1, 0>(tensor)) == Int<4>{}); - static_assert(Layout::rank == 2, "Only support 2D Tensor"); - typedef __NATIVE_VECTOR__(2, int) VecType; - const int lane_id = threadIdx.x % 64; - const int row_idx_offset = row_idx_offset_; - const int col_idx_offset = col_idx_offset_ + (lane_id / 16) * 4; - if constexpr (mergeLdg) { - #pragma unroll - for (int mi = 0; mi < size<0, 1>(tensor); ++mi) { - const int row_idx = row_idx_offset + mi * warp_row_stride; - #pragma unroll - for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { - const int col_idx_base = col_idx_offset + nj * warp_col_stride; - T bias_16[4]; - if constexpr (Is_even_MN) { - uint64_t *bias_64 = reinterpret_cast(bias_16); - bias_64[0] = *((uint64_t *)(bias + row_idx * bias_row_stride + col_idx_base)); - } else { - bool mask = row_idx < max_seqlen_q && col_idx_base < max_seqlen_k; - VecType *dst_ptr = reinterpret_cast(bias_16); - VecType *src_ptr = reinterpret_cast(bias + row_idx * bias_row_stride + col_idx_base); - *dst_ptr = __builtin_mxc_ldg_b64_predicator(src_ptr, 0, true, true, false, false, - mask, 1, MACA_ICMP_EQ); - } - #pragma unroll - for (int j = 0; j < size<1, 0>(tensor); ++j) { - if (row_idx < max_seqlen_q && col_idx_base < max_seqlen_k) { - tensor(make_coord(0, mi), make_coord(j, nj)) += bias_16[j] / softmax_scale; - } - } - } - } - } else { - #pragma unroll - for (int mi = 0; mi < size<0, 1>(tensor); ++mi) { - const int row_idx = row_idx_offset + mi * warp_row_stride; - bool row_mask = Is_even_MN || row_idx < max_seqlen_q; - #pragma unroll - for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { - const int col_idx_base = col_idx_offset + nj * warp_col_stride; - #pragma unroll - for (int j = 0; j < size<1, 0>(tensor); ++j) { - /*naive impl will have ldg_u16*/ - const int col_idx = col_idx_base + j; - if (row_mask && col_idx < max_seqlen_k) { - tensor(make_coord(0, mi), make_coord(j, nj)) += *(bias + row_idx * bias_row_stride + col_idx * bias_col_stride) / softmax_scale; - } - } - } - } - } -} - -template -inline __device__ void apply_attn_mask(Tensor0& acc_s, Tensor1& tSrMask, const float softmax_scale) { - // acc_s shape is (4, MMA_M, MMA_N) - CUTE_STATIC_ASSERT_V(size<0>(acc_s) == _4{}); - CUTE_STATIC_ASSERT_V(size<0>(acc_s) == size<0>(tSrMask)); - CUTE_STATIC_ASSERT_V(size<1>(acc_s) == size<1>(tSrMask)); - CUTE_STATIC_ASSERT_V(size<2>(acc_s) == size<2>(tSrMask)); - using T = typename Tensor1::value_type; - CONVERT_TENSOR_TYPE(T, float, tSrMask, rMask) - #pragma unroll - for (int m = 0; m < size<1>(acc_s); m++) { - #pragma unroll - for (int n = 0; n < size<2>(acc_s); n++) { - #pragma unroll - for (int i = 0; i < size<0>(acc_s); i++) { - acc_s(i, m, n) += rMask(i, m, n) / softmax_scale; - } - } - } -} - -template -inline __device__ void load_attn_mask(Tensor0& tSgMask, Tensor1& tSrMask, Tensor2& tScMask, const int max_N, const int max_M) { - // load attn_mask bias from global -> register - // tSgMask shape is (4, MMA_M, MMA_N) - CUTE_STATIC_ASSERT_V(size<0>(tSgMask) == _4{}); - CUTE_STATIC_ASSERT_V(size<0>(tSgMask) == size<0>(tSrMask)); - CUTE_STATIC_ASSERT_V(size<1>(tSgMask) == size<1>(tSrMask)); - CUTE_STATIC_ASSERT_V(size<2>(tSgMask) == size<2>(tSrMask)); - typedef __NATIVE_VECTOR__(2, int) VecType; - if constexpr (mergeLdg) { - #pragma unroll - for (int m = 0; m < size<1>(tSgMask); m++) { - bool row_mask = Is_even_MN || get<0>(tScMask(0, m, 0)) < max_M; - #pragma unroll - for (int n = 0; n < size<2>(tSgMask); n++) { - bool col_mask = Is_even_MN || get<1>(tScMask(0, 0, n)) < max_N; - auto src_ptr = (VecType *)(tSgMask(_, m, n).data().get()); // gmem - auto dst_ptr = (VecType *)(tSrMask(_, m, n).data()); // rf - if constexpr (Is_even_MN) { - *dst_ptr = __builtin_mxc_ldg_b64(src_ptr, 0, -1, true, true, false, false); - } else{ - *dst_ptr = __builtin_mxc_ldg_b64_predicator(src_ptr, 0, true, true, false, false, - row_mask && col_mask, 1, MACA_ICMP_EQ); - } - } - } - } else { - #pragma unroll - for (int m = 0; m < size<1>(tSgMask); m++) { - bool row_mask = Is_even_MN || get<0>(tScMask(0, m, 0)) < max_M; - #pragma unroll - for (int n = 0; n < size<2>(tSgMask); n++) { - int col_base_idx = get<1>(tScMask(0, 0, n)); - #pragma unroll - for (int i = 0; i < size<0>(tSgMask); i++) { - if (row_mask && col_base_idx + i < max_N) { - tSrMask(i, m, n) = tSgMask(i, m, n); - } - } - } - } - } -} - -} // namespace flash \ No newline at end of file diff --git a/csrc/flash_kernel/feature/dropout.h b/csrc/flash_kernel/feature/dropout.h deleted file mode 100644 index cb2748c..0000000 --- a/csrc/flash_kernel/feature/dropout.h +++ /dev/null @@ -1,206 +0,0 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - -/****************************************************************************** - * Copyright (c) 2024, Tri Dao. - ******************************************************************************/ - -#pragma once - -#include "philox.cuh" -#include "utils.h" - -namespace flash { - -struct Dropout { - - const unsigned long long seed, offset; - const uint8_t p_dropout_in_uint8_t; - - __forceinline__ __device__ Dropout(const unsigned long long seed, const unsigned long long offset, - const uint8_t p_dropout_in_uint8_t, - const int bid, const int hid, const int tid, const int nheads) - : seed(seed) - , offset(offset + (bid * nheads + hid) * 64 + tid % 64) - , p_dropout_in_uint8_t(p_dropout_in_uint8_t) { - } - - template - __forceinline__ __device__ void apply_dropout(Tensor &tensor_, - int block_row_start, int block_col_start, int block_row_stride) { - // convert shape from (4, MMA_M, MMA_N) to (8, MMA_M, MMA_N / 2) - Tensor tensor = make_tensor(tensor_.data(), flash::convert_layout_acc_dropout(tensor_.layout())); - using T = typename Engine::value_type; - auto encode_dropout = [](bool keep, T val) { - return keep ? val : (encode_dropout_in_sign_bit ? -val : T(0)); - }; - static_assert(decltype(size<2>(tensor))::value % 2 == 0); - const uint16_t p_dropout_8bit_in_uint16_t = uint16_t(p_dropout_in_uint8_t); - const uint32_t p_dropout_8bit_in_uint32_t = (uint32_t(p_dropout_8bit_in_uint16_t) << 16) | uint32_t(p_dropout_8bit_in_uint16_t); - // if (cute::thread0()) { printf("threshold2 = 0x%x\n", p_dropout_8bit_in_uint32_t); } - #pragma unroll - for (int m = 0; m < size<1>(tensor); ++m, block_row_start += block_row_stride) { - uint2 rowcol = make_uint2(block_row_start, block_col_start); - #pragma unroll - for (int n = 0; n < size<2>(tensor) / 2; ++n, ++rowcol.y) { - // if (cute::thread(32, 0)) { printf("m = %d, n = %d, row = %d, col = %d\n", m, n, int(rowcol.x), int(rowcol.y));} - uint4 random_uint4 = flash::philox(seed, reinterpret_cast(rowcol), offset); - // if (cute::thread0()) { printf("philox = %u, %d, %d, %d\n", random_uint4.x, random_uint4.y, random_uint4.z, random_uint4.w);} - uint8_t (&rnd_8)[16] = reinterpret_cast(random_uint4); - // Special implementation for 16-bit types: we duplicate the threshold to the - // low and high 16 bits of a 32-bit value, then use the f16x2 comparison instruction - // to get a mask. The low 16 bits of the mask will be either 0xffff or 0x0000, - // and the high 16 bits will be either 0xffff or 0x0000, depending on whether - // the random value is less than the threshold. - // We then do a bit-wise AND between the mask and the original value (in 32-bit). - // We're exploiting the fact that floating point comparison is equivalent to integer - // comparison, since we're comparing unsigned integers whose top 8-bits are zero. - if (!encode_dropout_in_sign_bit - && (std::is_same::value || std::is_same::value)) { - uint16_t rnd_16[16]; - #pragma unroll - for (int i = 0; i < 16; i++) { rnd_16[i] = uint16_t(rnd_8[i]); } - uint32_t (&rnd_32)[8] = reinterpret_cast(rnd_16); - #pragma unroll - for (int j = 0; j < 2; j++) { - Tensor tensor_uint32 = recast(tensor(_, m, n * 2 + j)); - // if (cute::thread0()) { printf("random = 0x%x, 0x%x, 0x%x, 0x%x\n", rnd_32[j * 4 + 0], rnd_32[j * 4 + 1], rnd_32[j * 4 + 2], rnd_32[j * 4 + 3]); } - // if (cute::thread0()) { printf("tensor_uint32 = 0x%x, 0x%x, 0x%x, 0x%x\n", tensor_uint32(0), tensor_uint32(1), tensor_uint32(2), tensor_uint32(3)); } - #pragma unroll - for (int i = 0; i < 4; i++) { - uint32_t mask; - //asm volatile("set.le.u32.f16x2 %0, %1, %2;\n" : "=r"(mask) : "r"(rnd_32[j * 4 + i]), "r"(p_dropout_8bit_in_uint32_t)); - tensor_uint32(i) &= mask; - } - // if (cute::thread0()) { printf("tensor_uint32 = 0x%x, 0x%x, 0x%x, 0x%x\n", tensor_uint32(0), tensor_uint32(1), tensor_uint32(2), tensor_uint32(3)); } - } - } else { - #pragma unroll - for (int j = 0; j < 2; j++) { - #pragma unroll - for (int i = 0; i < 8; i++) { - tensor(i, m, n * 2 + j) = encode_dropout(rnd_8[j * 8 + i] <= p_dropout_in_uint8_t, tensor(i, m, n * 2 + j)); - } - Tensor tensor_uint32 = recast(tensor(_, m, n * 2 + j)); - // if (cute::thread0()) { printf("tensor_uint32 = 0x%x, 0x%x, 0x%x, 0x%x\n", tensor_uint32(0), tensor_uint32(1), tensor_uint32(2), tensor_uint32(3)); } - } - } - // // if ((threadIdx.x == 0) && (blockIdx.x == 0) && (blockIdx.y == 0)) { - // // printf("n = %d, ph Philox: %u, %u, %u, %u\n", n, rnd_8.x, rnd_8.y, rnd_8.z, rnd_8.w); - // // } - } - } - } - - template - __forceinline__ __device__ void mc_apply_dropout(Tensor &tensor, - int block_row_start, int block_col_start, - int block_row_stride, - int kBlockN, int n_block) { - using T = typename Engine::value_type; - auto encode_dropout = [](bool keep, T val) { - return keep ? val : (encode_dropout_in_sign_bit ? -val : T(0)); - }; - static_assert(decltype(size<0>(tensor))::value == 4); - int warp_id = threadIdx.x / 64; - - #pragma unroll - for (int m = 0; m < size<1>(tensor); ++m, block_row_start += block_row_stride) { - // use block_col_offset to control rnd_8 apply - int block_col_offset = ((kBlockN * n_block) % 64) / 16; - int warp_col_offset = warp_id / block_row_stride; - int block_col_start_tmp = block_col_start; - - #pragma unroll - for (int n = 0; n < size<2>(tensor); ++n, warp_col_offset += AtomLayoutNS) { - // when blockN=128, one block contain two 64 in col, so we need to update block_col_start_tmp - block_col_start_tmp += warp_col_offset / 4; - warp_col_offset %= 4; - // rnd_8 contains 16 nums for kBlockN 64 in 1 block, to process 64 in col - // if use kBlockN 32 in 2 block to process 64 in col, rnd_8 should apply in two differen block - uint2 rowcol = make_uint2(block_row_start, block_col_start_tmp); - uint4 random_uint4 = flash::philox(seed, reinterpret_cast(rowcol), offset); - uint8_t (&rnd_8)[16] = reinterpret_cast(random_uint4); - #pragma unroll - for (int i = 0; i < 4; ++i) { - int rng_idx = (block_col_offset + warp_col_offset) * 4 + i; - // This implementation is a native perf version,here assembly instruction will have ldp_u8. - tensor(i, m, n) = encode_dropout(rnd_8[rng_idx] <= p_dropout_in_uint8_t, tensor(i, m, n)); - } - } - } - } - - // it works for: - // blockN=64, waves layout: 4x1, 2x2, 4x2, 1x4, 2x4 - // blockN=128, waves layout: 1x4, 2x4 - template - __forceinline__ __device__ void mc_apply_dropout(Tensor &tensor, int block_row_start, int block_col_start) { - using T = typename Engine::value_type; - auto encode_dropout = [](bool keep, T val) { - return keep ? val : (encode_dropout_in_sign_bit ? -val : T(0)); - }; - static_assert(decltype(size<0>(tensor))::value == 4); - const int wave_col = threadIdx.x / 64 / AtomLayoutMS; - - if constexpr (AtomLayoutNS == 1) { - #pragma unroll - for (int m = 0; m < size<1>(tensor); ++m, block_row_start += AtomLayoutMS) { - uint2 rowcol = make_uint2(block_row_start, block_col_start); - uint4 random_uint4 = flash::philox(seed, reinterpret_cast(rowcol), offset); - uint8_t (&rnd_8)[16] = reinterpret_cast(random_uint4); - #pragma unroll - for (int n = 0; n < size<2>(tensor); ++n) { - #pragma unroll - for (int i = 0; i < 4; ++i) { - // w0|w0|w0|w0 - tensor(i, m, n) = encode_dropout(rnd_8[n * 4 + i] <= p_dropout_in_uint8_t, tensor(i, m, n)); - } - } - } - } else if constexpr (AtomLayoutNS == 2) { - #pragma unroll - for (int m = 0; m < size<1>(tensor); ++m, block_row_start += AtomLayoutMS) { - uint2 rowcol = make_uint2(block_row_start, block_col_start); - uint4 random_uint4 = flash::philox(seed, reinterpret_cast(rowcol), offset); - uint8_t (&rnd_8)[16] = reinterpret_cast(random_uint4); - #pragma unroll - for (int n = 0; n < size<2>(tensor); ++n) { - #pragma unroll - for (int i = 0; i < 4; ++i) { - // e.g., w0|w2|w0|w2 - if (wave_col == 0) { - tensor(i, m, n) = encode_dropout(rnd_8[n * 8 + i] <= p_dropout_in_uint8_t, tensor(i, m, n)); - } else { - tensor(i, m, n) = encode_dropout(rnd_8[n * 8 + i + 4] <= p_dropout_in_uint8_t, tensor(i, m, n)); - } - } - } - } - } else if constexpr (AtomLayoutNS == 4) { - #pragma unroll - for (int m = 0; m < size<1>(tensor); ++m, block_row_start += AtomLayoutMS) { - #pragma unroll - for (int n = 0; n < size<2>(tensor); ++n, block_col_start += 1) { - uint2 rowcol = make_uint2(block_row_start, block_col_start); - uint4 random_uint4 = flash::philox(seed, reinterpret_cast(rowcol), offset); - uint8_t (&rnd_8)[16] = reinterpret_cast(random_uint4); - #pragma unroll - for (int i = 0; i < 4; ++i) { - // e.g., w0|w2|w4|w6 - if (wave_col == 0) { - tensor(i, m, n) = encode_dropout(rnd_8[i] <= p_dropout_in_uint8_t, tensor(i, m, n)); - } else if (wave_col == 1) { - tensor(i, m, n) = encode_dropout(rnd_8[i + 4] <= p_dropout_in_uint8_t, tensor(i, m, n)); - } else if (wave_col == 2) { - tensor(i, m, n) = encode_dropout(rnd_8[i + 8] <= p_dropout_in_uint8_t, tensor(i, m, n)); - } else { - tensor(i, m, n) = encode_dropout(rnd_8[i + 12] <= p_dropout_in_uint8_t, tensor(i, m, n)); - } - } - } - } - } - } -}; - -} // namespace flash diff --git a/csrc/flash_kernel/feature/mask.h b/csrc/flash_kernel/feature/mask.h index 36441be..e77c982 100644 --- a/csrc/flash_kernel/feature/mask.h +++ b/csrc/flash_kernel/feature/mask.h @@ -7,6 +7,7 @@ #pragma once #include +#include "utils.h" namespace flash { @@ -37,113 +38,38 @@ __forceinline__ __device__ void apply_mask(Tensor &tensor, const } } -// [warp_col_stride] tiled mma 4x1: 16, tiled mma 2x2: 32 -template -__forceinline__ __device__ void apply_mask_local(Tensor &tensor, const int col_idx_offset_, - const int max_seqlen_k, const int row_idx_offset, - const int max_seqlen_q, const int warp_row_stride, - const int window_size_left, const int window_size_right, - const int warp_col_stride = 16) { - // tensor has shape (nrow=(1, MMA_M), ncol=(4, MMA_N)) - static_assert(Layout::rank == 2, "Only support 2D Tensor"); - static_assert(decltype(size<0, 0>(tensor))::value == 1); - static_assert(decltype(size<1, 0>(tensor))::value == 4); - const int col_idx_offset = col_idx_offset_ + ((__lane_id() >> 4) << 2); - #pragma unroll - for (int mi = 0; mi < size<0, 1>(tensor); ++mi) { - const int row_idx = row_idx_offset + mi * warp_row_stride; - const int col_idx_limit_left = std::max(0, row_idx + max_seqlen_k - max_seqlen_q - window_size_left); - const int col_idx_limit_right = std::min(max_seqlen_k, row_idx + 1 + max_seqlen_k - max_seqlen_q + window_size_right); - #pragma unroll - for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { - const int col_idx_base = col_idx_offset + nj * warp_col_stride; - #pragma unroll - for (int j = 0; j < size<1, 0>(tensor); ++j) { - const int col_idx = col_idx_base + j; - if (col_idx >= col_idx_limit_right || (HasWSLeft && col_idx < col_idx_limit_left)) { - tensor(make_coord(0, mi), make_coord(j, nj)) = -INFINITY; - } - } - } - // if (cute::thread0()) { - // printf("mi = %d, i = %d, row_idx = %d, max_seqlen_k = %d\n", mi, i, row_idx, max_seqlen_k); - // print(tensor(make_coord(i, mi), _)); - // // print(tensor(_, j + nj * size<1, 0>(tensor))); - // } - } -} -template -__forceinline__ __device__ void apply_mask_causal(Tensor &tensor, const int col_idx_offset_, - const int max_seqlen_k, const int row_idx_offset, - const int max_seqlen_q, const int warp_row_stride, - const int warp_col_stride = 16) { - // Causal masking is equivalent to local masking with window_size_left = infinity and window_size_right = 0 - apply_mask_local(tensor, col_idx_offset_, max_seqlen_k, row_idx_offset, - max_seqlen_q, warp_row_stride, -1, 0, warp_col_stride); -} - -template -__forceinline__ __device__ void apply_mask_causal_w_idx( - Tensor &tensor, Tensor const &idx_rowcol, - const int col_idx_offset_, const int max_seqlen_k, const int row_idx_offset) -{ - // tensor has shape (ncol=(2, MMA_M), nrow=(2, MMA_N)) - static_assert(Layout0::rank == 2, "Only support 2D Tensor"); - static_assert(Layout1::rank == 2, "Only support 2D Tensor"); - CUTE_STATIC_ASSERT_V(size<0>(tensor) == size<0>(idx_rowcol)); - CUTE_STATIC_ASSERT_V(size<1>(tensor) == size<1>(idx_rowcol)); - #pragma unroll - for (int mi = 0; mi < size<0>(tensor); ++mi) { - const int col_idx_limit = std::min(max_seqlen_k, 1 + row_idx_offset + get<0>(idx_rowcol(mi, 0))); - #pragma unroll - for (int ni = 0; ni < size<1, 1>(tensor); ++ni) { - if (col_idx_offset_ + get<1>(idx_rowcol(0, ni)) >= col_idx_limit) { - tensor(mi, ni) = -INFINITY; - } - } - // if (cute::thread0()) { - // printf("ni = %d, j = %d, col_idx = %d, max_seqlen_k = %d\n", ni, j, col_idx, max_seqlen_k); - // print(tensor(_, make_coord(j, ni))); - // // print(tensor(_, j + ni * size<1, 0>(tensor))); - // } - } -} - -template +template struct Mask { const int max_seqlen_k, max_seqlen_q, ngroups; - const int window_size_left, window_size_right; - const float alibi_slope; + // CP (Context Parallelism) parameters + const int tot_seqlen_k, cp_world_size, cp_rank; - __forceinline__ __device__ Mask(const int max_seqlen_k, const int max_seqlen_q, const int ngroups, - const int window_size_left, const int window_size_right, - const float alibi_slope=0.f) + __forceinline__ __device__ Mask(const int max_seqlen_k, const int max_seqlen_q, const int ngroups, const int tot_seqlen_k = 0, const int cp_world_size = 1, const int cp_rank = 0) : max_seqlen_k(max_seqlen_k) , max_seqlen_q(max_seqlen_q) , ngroups(ngroups) - , window_size_left(window_size_left) - , window_size_right(window_size_right) - , alibi_slope(!Has_alibi ? 0.0 : alibi_slope) { + , tot_seqlen_k(tot_seqlen_k) + , cp_world_size(cp_world_size) + , cp_rank(cp_rank) { }; // Causal_mask: whether this particular iteration needs causal masking - template + template __forceinline__ __device__ void apply_mask(Tensor &tensor_, const int col_idx_offset_, const int row_idx_offset, const int warp_row_stride) { - static_assert(!(Causal_mask && Is_local), "Cannot be both causal and local"); static_assert(Layout::rank == 3, "Only support 3D Tensor"); - static_assert(decltype(size<0>(tensor_))::value == 4, "First dimension must be 4"); - static constexpr bool Need_masking = Has_alibi || Causal_mask || Is_local || !Is_even_MN; - // if (cute::thread0()) { printf("Has_alibi = %d, Causal_mask=%d, Is_local=%d, Is_even_MN = %d, Need_masking = %d\n", Has_alibi, Causal_mask, Is_local, Is_even_MN, Need_masking); } + static_assert(decltype(size<0>(tensor_))::value == Elem_per_thread, "The tensor_ first dimension not match the Elem_per_thread"); + static constexpr bool Need_masking = Causal_mask || !Is_even_MN; + if constexpr (Need_masking) { // Reshape tensor_ from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N)) Tensor tensor = make_tensor(tensor_.data(), flash::convert_layout_acc_rowcol(tensor_.layout())); // Do we need both row and column indices, or just column incides? - static constexpr bool Col_idx_only = !(Has_alibi && !Is_causal) && !Is_local && !Causal_mask; + static constexpr bool Col_idx_only = !Causal_mask; const int col_idx_offset = col_idx_offset_ + ((__lane_id() >> 4) << 2); if constexpr (Col_idx_only) { #pragma unroll @@ -155,9 +81,6 @@ struct Mask { #pragma unroll for (int mi = 0; mi < size<0>(tensor); ++mi) { // No causal, no local - if constexpr (Has_alibi) { - tensor(mi, make_coord(j, nj)) += alibi_slope * col_idx; - } if constexpr (!Is_even_MN) { if (col_idx >= max_seqlen_k) { tensor(mi, make_coord(j, nj)) = -INFINITY; } } @@ -171,38 +94,26 @@ struct Mask { #pragma unroll for (int i = 0; i < size<0, 0>(tensor); ++i) { const int row_idx = row_idx_base + i * 16; - const int col_idx_limit_left = std::max(0, row_idx + max_seqlen_k - max_seqlen_q - window_size_left); - const int col_idx_limit_right = std::min(max_seqlen_k, row_idx / ngroups + 1 + max_seqlen_k - max_seqlen_q / ngroups + window_size_right); + const int col_idx_limit_right = std::min(tot_seqlen_k, row_idx / ngroups + 1 + tot_seqlen_k - max_seqlen_q / ngroups); #pragma unroll for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { const int col_idx_base = col_idx_offset + nj * 16; #pragma unroll for (int j = 0; j < size<1, 0>(tensor); ++j) { const int col_idx = col_idx_base + j; - if constexpr (Has_alibi) { - if constexpr (Is_causal) { - tensor(make_coord(i, mi), make_coord(j, nj)) += alibi_slope * col_idx; - } else { - tensor(make_coord(i, mi), make_coord(j, nj)) -= alibi_slope * abs(row_idx + max_seqlen_k - max_seqlen_q - col_idx); - + if constexpr (Is_enable_dcp) { + // casusal mask with dcp + const int actual_col_idx = col_idx * cp_world_size + cp_rank + 1; + // actual_col_idx start from 1 to tot_seqlen_k + if (actual_col_idx > col_idx_limit_right || col_idx >= max_seqlen_k) { + tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY; } - } - if constexpr (Causal_mask) { + } else { + // casusal mask without dcp if (col_idx >= col_idx_limit_right) { tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY; } } - if constexpr (Is_local) { - if (col_idx >= col_idx_limit_right || col_idx < col_idx_limit_left) { - tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY; - } - } - if constexpr (!Causal_mask && !Is_local && !Is_even_MN) { - // Causal and Local already handles MN masking - if (col_idx >= max_seqlen_k) { - tensor(make_coord(i, mi), make_coord(j, nj)) = -INFINITY; - } - } } } } @@ -211,6 +122,32 @@ struct Mask { } }; + template + __forceinline__ __device__ void apply_sparse_attn_mask(Tensor &tensor_, + const int col_idx_offset_, + int32_t* indices_smem_ptr, + bool is_indices_all_valid) { + Tensor tensor = make_tensor(tensor_.data(), flash::convert_layout_acc_rowcol(tensor_.layout())); + const int col_idx_offset = col_idx_offset_ + ((__lane_id() >> 4) << 2); + #pragma unroll + for (int nj = 0; nj < size<1, 1>(tensor); ++nj) { + const int col_idx_base = col_idx_offset + nj * 16; + #pragma unroll + for (int j = 0; j < size<1, 0>(tensor); ++j) { + const int col_idx = col_idx_base + j; + // const bool invalid_flag = !is_indices_all_valid && (CHECK_BIT(is_valid_indices[(col_idx % kBlockN) >> 5], col_idx % kBlockN) == false); + const bool invalid_flag = !is_indices_all_valid && indices_smem_ptr[col_idx % kBlockTopK] < 0; + #pragma unroll + for (int mi = 0; mi < size<0>(tensor); ++mi) { + if constexpr (!Is_even_MN) { + if (col_idx >= max_seqlen_k) { tensor(mi, make_coord(j, nj)) = -INFINITY; } + } + if (invalid_flag) tensor(mi, make_coord(j, nj)) = -INFINITY; + } + } + } + }; + }; } // namespace flash diff --git a/csrc/flash_kernel/feature/rotary.h b/csrc/flash_kernel/feature/rotary.h deleted file mode 100644 index 3664eae..0000000 --- a/csrc/flash_kernel/feature/rotary.h +++ /dev/null @@ -1,529 +0,0 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - -/****************************************************************************** - * Copyright (c) 2024, Tri Dao. - ******************************************************************************/ - -#pragma once - -#include - -#include "utils.h" - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -namespace flash { - -using namespace cute; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -__forceinline__ __device__ void copy_rotary_interleaved(Tensor const &S, - Tensor &D, - Tensor const &Cos, - Tensor const &Sin, - Tensor const &identity_MN, - const int max_MN, const int min_MN, - const int dim, const int rotary_dim) { - CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); - CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); - CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Cos)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Cos)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Sin)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Sin)); // MMA_K - CUTE_STATIC_ASSERT_V(size<0>(Cos) == size<0>(Sin)); // MMA_K - static_assert(decltype(size<0>(S))::value == decltype(size<0>(Cos))::value * 2); - static_assert(decltype(size<0>(Cos))::value % 2 == 0); // Since we do fast conversion from fp16/bf16 to fp32 - Tensor rCos = make_fragment_like(Cos); - Tensor rSin = make_fragment_like(Sin); - Tensor rS = make_fragment_like(S); - typedef __NATIVE_VECTOR__(2, float) Float2; - #pragma unroll - for (int m = 0; m < size<1>(S); ++m) { - if (get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN) { - #pragma unroll - for (int k = 0; k < size<2>(S); ++k) { - if (Is_even_K || get<1>(identity_MN(0, 0, k)) < dim) { - cute::copy(S(_, m, k), rS(_, m, k)); - if (get<1>(identity_MN(0, 0, k)) < rotary_dim) { - cute::copy(Cos(_, m, k), rCos(_, m, k)); - cute::copy(Sin(_, m, k), rSin(_, m, k)); - // Tensor S_fp32 = convert_type(rS(_, m, k)); - // Tensor cos_fp32 = convert_type(rCos(_, m, k)); - // Tensor sin_fp32 = convert_type(rSin(_, m, k)); - using T = typename Engine0::value_type; - using T_rotary = typename Engine2::value_type; - CONVERT_TENSOR_TYPE(T, float, rS(_, m, k), S_fp32) - CONVERT_TENSOR_TYPE(T_rotary, float, rCos(_, m, k), cos_fp32) - CONVERT_TENSOR_TYPE(T_rotary, float, rSin(_, m, k), sin_fp32) - - #pragma unroll - for (int i = 0; i < size<0>(rS) / 2; ++i) { - Float2 x_vec = {S_fp32(2 * i), S_fp32(2 * i + 1)}; - Float2 real_vec = {cos_fp32(i), sin_fp32(i)}; - Float2 imag_vec = {sin_fp32(i), cos_fp32(i)}; - Float2 beta_vec = {0.0f, 0.0f}; - real_vec = __builtin_mxc_pk_fma_f32(x_vec, real_vec, beta_vec); - imag_vec = __builtin_mxc_pk_fma_f32(x_vec, imag_vec, beta_vec); - S_fp32(2 * i) = real_vec[0] - real_vec[1]; - S_fp32(2 * i + 1) = imag_vec[0] + imag_vec[1]; - //float real = S_fp32(2 * i) * cos_fp32(i) - S_fp32(2 * i + 1) * sin_fp32(i); - //float imag = S_fp32(2 * i) * sin_fp32(i) + S_fp32(2 * i + 1) * cos_fp32(i); - //S_fp32(2 * i) = real; - //S_fp32(2 * i + 1) = imag; - } - // Idk but I need to copy for the convert_type to work - Tensor S_fp32_copy = make_fragment_like(S_fp32); - cute::copy(S_fp32, S_fp32_copy); - //Tensor S_og_type = convert_type(S_fp32_copy); - CONVERT_TENSOR_TYPE(float, T, S_fp32_copy, S_og_type) - cute::copy(S_og_type, rS(_, m, k)); - } - cute::copy(rS(_, m, k), D(_, m, k)); - } else if (Clear_OOB_K) { - cute::clear(D(_, m, k)); - } - } - } - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -__forceinline__ __device__ void copy_rotary_contiguous(Tensor const &S, - Tensor &D, - Tensor const &Cos, - Tensor const &Sin, - Tensor const &identity_MN, - const int max_MN, const int min_MN, - const int dim, const int rotary_dim) { - CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); - CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); - CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Cos)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Cos)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Sin)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Sin)); // MMA_K - CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(Cos)); // MMA - CUTE_STATIC_ASSERT_V(size<0>(Cos) == size<0>(Sin)); - static_assert(decltype(size<0>(Cos))::value % 2 == 0); // Since we do fast conversion from fp16/bf16 to fp32 - Tensor rCos = make_fragment_like(Cos); - Tensor rSin = make_fragment_like(Sin); - Tensor rS = make_fragment_like(S); - Tensor rS_other = make_fragment_like(rS(_, 0, 0)); - typedef __NATIVE_VECTOR__(2, float) Float2; - Float2 beta_vec = {0.0f, 0.0f}; - - const int rotary_dim_half = rotary_dim >> 1; - - #pragma unroll - for (int m = 0; m < size<1>(S); ++m) { - if (get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN) { - #pragma unroll - for (int k = 0; k < size<2>(S); ++k) { - if (Is_even_K || get<1>(identity_MN(0, 0, k)) < dim) { - cute::copy(S(_, m, k), rS(_, m, k)); - if (get<1>(identity_MN(0, 0, k)) < rotary_dim) { - const bool is_left = get<1>(identity_MN(0, 0, k)) < rotary_dim_half; - Tensor gS_other = make_tensor(S(_, m, k).data() + (is_left ? rotary_dim_half : -rotary_dim_half), S(_, m, k).layout()); - cute::copy(gS_other, rS_other); - // if (cute::thread0()) { print_tensor(rS(_, m, k)); print_tensor(rS_other); } - Tensor gCos = make_tensor(Cos(_, m, k).data() + (is_left ? 0 : -rotary_dim_half), Cos(_, m, k).layout()); - Tensor gSin = make_tensor(Sin(_, m, k).data() + (is_left ? 0 : -rotary_dim_half), Sin(_, m, k).layout()); - cute::copy(gCos, rCos(_, m, k)); - cute::copy(gSin, rSin(_, m, k)); - // if (cute::thread0()) { print_tensor(rCos(_, m, k)); print_tensor(rSin(_, m, k)); } - // Tensor S_fp32 = convert_type(rS(_, m, k)); - // Tensor S_other_fp32 = convert_type(rS_other); - // Tensor cos_fp32 = convert_type(rCos(_, m, k)); - // Tensor sin_fp32 = convert_type(rSin(_, m, k)); - - using T = typename Engine0::value_type; - using T_rotary = typename Engine2::value_type; - CONVERT_TENSOR_TYPE(T, float, rS(_,m,k), S_fp32) - CONVERT_TENSOR_TYPE(T, float, rS_other, S_other_fp32) - CONVERT_TENSOR_TYPE(T_rotary, float, rCos(_, m, k), cos_fp32) - CONVERT_TENSOR_TYPE(T_rotary, float, rSin(_, m, k), sin_fp32) - - #pragma unroll - for (int i = 0; i < size<0>(rS); ++i) { - //S_fp32(i) = S_fp32(i) * cos_fp32(i) + S_other_fp32(i) * (is_left ? -sin_fp32(i) : sin_fp32(i)); - Float2 x_vec = {S_fp32(i), S_other_fp32(i)}; - Float2 alpha_vec = {cos_fp32(i), is_left ? -sin_fp32(i) : sin_fp32(i)}; - Float2 y_vec = __builtin_mxc_pk_fma_f32(x_vec, alpha_vec, beta_vec); - S_fp32(i) = y_vec[0] + y_vec[1]; - } - // Idk but I need to copy for the convert_type to work - Tensor S_fp32_copy = make_fragment_like(S_fp32); - cute::copy(S_fp32, S_fp32_copy); - //Tensor S_og_type = convert_type(S_fp32_copy); - CONVERT_TENSOR_TYPE(float, T, S_fp32_copy, S_og_type) - - cute::copy(S_og_type, rS(_, m, k)); - // if (cute::thread0()) { print_tensor(rS(_, m, k)); } - } - cute::copy(rS(_, m, k), D(_, m, k)); - } else if (Clear_OOB_K) { - cute::clear(D(_, m, k)); - } - } - } - } - -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -__forceinline__ __device__ void copy_rotary_interleaved_to_reg(Tensor const &S, - uint32_t *D_ptr, - Tensor const &Cos, - Tensor const &Sin, - Tensor const &identity_MN, - const int max_MN, const int min_MN, - const int dim, const int rotary_dim) { - CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Cos)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Cos)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Sin)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Sin)); // MMA_K - CUTE_STATIC_ASSERT_V(size<0>(Cos) == size<0>(Sin)); // MMA_K - static_assert(decltype(size<0>(S))::value == decltype(size<0>(Cos))::value * 2); - static_assert(decltype(size<0>(Cos))::value % 2 == 0); // Since we do fast conversion from fp16/bf16 to fp32 - Tensor rCos = make_fragment_like(Cos); - Tensor rSin = make_fragment_like(Sin); - Tensor rS = make_fragment_like(S); - - typedef __NATIVE_VECTOR__(2, float) Float2; - typedef __NATIVE_VECTOR__(4, int) VecTypeB128; - typedef __NATIVE_VECTOR__(2, int) VecTypeB64; - - #pragma unroll - for (int m = 0; m < size<1>(S); ++m) { - bool row_mask = get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN; - #pragma unroll - for (int k = 0; k < size<2>(S); ++k) { - bool mask = row_mask && (Is_even_K || get<1>(identity_MN(0, 0, k)) < dim); - auto S_ptr = (VecTypeB128 *)(S(_, m, k).data().ptr_); - auto rS_ptr = (VecTypeB128 *)(rS(_, m, k).data()); - rS_ptr[0] = __builtin_mxc_ldg_b128_predicator(S_ptr, 0, true, true, false, false, - mask, 1, MACA_ICMP_EQ); - bool rotary_mask = mask && (get<1>(identity_MN(0, 0, k)) < rotary_dim); - - auto gCos_ptr = (VecTypeB64 *)(Cos(_, m, k).data().ptr_); - auto gSin_ptr = (VecTypeB64 *)(Sin(_, m, k).data().ptr_); - auto rCos_ptr = (VecTypeB64 *)(rCos(_, m, k).data()); - auto rSin_ptr = (VecTypeB64 *)(rSin(_, m, k).data()); - rCos_ptr[0] = __builtin_mxc_ldg_b64_predicator(gCos_ptr, 0, true, true, false, false, - rotary_mask, 1, MACA_ICMP_EQ); - rSin_ptr[0] = __builtin_mxc_ldg_b64_predicator(gSin_ptr, 0, true, true, false, false, - rotary_mask, 1, MACA_ICMP_EQ); - - if (rotary_mask) { - using T = typename Engine0::value_type; - using T_rotary = typename Engine1::value_type; - CONVERT_TENSOR_TYPE(T, float, rS(_, m, k), S_fp32) - CONVERT_TENSOR_TYPE(T_rotary, float, rCos(_, m, k), cos_fp32) - CONVERT_TENSOR_TYPE(T_rotary, float, rSin(_, m, k), sin_fp32) - - #pragma unroll - for (int i = 0; i < size<0>(rS) / 2; ++i) { - Float2 x_vec = {S_fp32(2 * i), S_fp32(2 * i + 1)}; - Float2 real_vec = {cos_fp32(i), sin_fp32(i)}; - Float2 imag_vec = {sin_fp32(i), cos_fp32(i)}; - Float2 beta_vec = {0.0f, 0.0f}; - real_vec = __builtin_mxc_pk_fma_f32(x_vec, real_vec, beta_vec); - imag_vec = __builtin_mxc_pk_fma_f32(x_vec, imag_vec, beta_vec); - S_fp32(2 * i) = real_vec[0] - real_vec[1]; - S_fp32(2 * i + 1) = imag_vec[0] + imag_vec[1]; - } - // Idk but I need to copy for the convert_type to work - Tensor S_fp32_copy = make_fragment_like(S_fp32); - cute::copy(S_fp32, S_fp32_copy); - //Tensor S_og_type = convert_type(S_fp32_copy); - CONVERT_TENSOR_TYPE(float, T, S_fp32_copy, S_og_type) - cute::copy(S_og_type, rS(_, m, k)); - } - const int idx = (m * size<2>(S) + k) << 2; - auto D = (VecTypeB128 *)(D_ptr + idx); - D[0] = rS_ptr[0]; - } - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -__forceinline__ __device__ void copy_rotary_contiguous_to_reg(Tensor const &S, - uint32_t *D_ptr, - Tensor const &Cos, - Tensor const &Sin, - Tensor const &identity_MN, - const int max_MN, const int min_MN, - const int dim, const int rotary_dim) { - CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Cos)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Cos)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Sin)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Sin)); // MMA_K - CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(Cos)); // MMA - CUTE_STATIC_ASSERT_V(size<0>(Cos) == size<0>(Sin)); - static_assert(decltype(size<0>(Cos))::value % 2 == 0); // Since we do fast conversion from fp16/bf16 to fp32 - Tensor rCos = make_fragment_like(Cos); - Tensor rSin = make_fragment_like(Sin); - Tensor rS = make_fragment_like(S); - Tensor rS_other = make_fragment_like(rS(_, 0, 0)); - typedef __NATIVE_VECTOR__(2, float) Float2; - Float2 beta_vec = {0.0f, 0.0f}; - - const int rotary_dim_half = rotary_dim >> 1; - - typedef __NATIVE_VECTOR__(4, int) VecType; - #pragma unroll - for (int m = 0; m < size<1>(S); ++m) { - bool row_mask = get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN; - #pragma unroll - for (int k = 0; k < size<2>(S); ++k) { - bool mask = row_mask && (Is_even_K || get<1>(identity_MN(0, 0, k)) < dim); - auto S_ptr = (VecType *)(S(_, m, k).data().ptr_); - auto rS_ptr = (VecType *)(rS(_, m, k).data()); - rS_ptr[0] = __builtin_mxc_ldg_b128_predicator(S_ptr, 0, true, true, false, false, - mask, 1, MACA_ICMP_EQ); - - bool rotary_mask = mask && (get<1>(identity_MN(0, 0, k)) < rotary_dim); - const bool is_left = get<1>(identity_MN(0, 0, k)) < rotary_dim_half; - auto gS_ptr = (VecType *)(S(_, m, k).data().ptr_ + (is_left ? rotary_dim_half : -rotary_dim_half)); - auto rS_other_ptr = (VecType *)(rS_other.data()); - rS_other_ptr[0] = __builtin_mxc_ldg_b128_predicator(gS_ptr, 0, true, true, false, false, - rotary_mask, 1, MACA_ICMP_EQ); - - auto gCos_ptr = (VecType *)(Cos(_, m, k).data().ptr_ + (is_left ? 0 : -rotary_dim_half)); - auto gSin_ptr = (VecType *)(Sin(_, m, k).data().ptr_ + (is_left ? 0 : -rotary_dim_half)); - auto rCos_ptr = (VecType *)(rCos(_, m, k).data()); - auto rSin_ptr = (VecType *)(rSin(_, m, k).data()); - rCos_ptr[0] = __builtin_mxc_ldg_b128_predicator(gCos_ptr, 0, true, true, false, false, - rotary_mask, 1, MACA_ICMP_EQ); - rSin_ptr[0] = __builtin_mxc_ldg_b128_predicator(gSin_ptr, 0, true, true, false, false, - rotary_mask, 1, MACA_ICMP_EQ); - if (rotary_mask) { - using T = typename Engine0::value_type; - using T_rotary = typename Engine1::value_type; - CONVERT_TENSOR_TYPE(T, float, rS(_, m, k), S_fp32) - CONVERT_TENSOR_TYPE(T, float, rS_other, S_other_fp32) - CONVERT_TENSOR_TYPE(T_rotary, float, rCos(_, m, k), cos_fp32) - CONVERT_TENSOR_TYPE(T_rotary, float, rSin(_, m, k), sin_fp32) - - #pragma unroll - for (int i = 0; i < size<0>(rS); ++i) { - // S_fp32(i) = S_fp32(i) * cos_fp32(i) + S_other_fp32(i) * (is_left ? -sin_fp32(i) : sin_fp32(i)); - Float2 x_vec = {S_fp32(i), S_other_fp32(i)}; - Float2 alpha_vec = {cos_fp32(i), is_left ? -sin_fp32(i) : sin_fp32(i)}; - Float2 y_vec = __builtin_mxc_pk_fma_f32(x_vec, alpha_vec, beta_vec); - S_fp32(i) = y_vec[0] + y_vec[1]; - } - // Idk but I need to copy for the convert_type to work - Tensor S_fp32_copy = make_fragment_like(S_fp32); - cute::copy(S_fp32, S_fp32_copy); - //Tensor S_og_type = convert_type(S_fp32_copy); - CONVERT_TENSOR_TYPE(float, T, S_fp32_copy, S_og_type) - cute::copy(S_og_type, rS(_, m, k)); - } - const int idx = (m * size<2>(S) + k) * 4; - auto D = (VecType *)(D_ptr + idx); - D[0] = rS_ptr[0]; - } - } - -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -__forceinline__ __device__ void copy_rotary_interleaved_to_global(Tensor const &S, - Tensor &D, - Tensor const &Cos, - Tensor const &Sin, - Tensor const &identity_MN, - const int max_MN, const int min_MN, - const int dim, const int rotary_dim) { - CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); - CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); - CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Cos)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Cos)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Sin)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Sin)); // MMA_K - CUTE_STATIC_ASSERT_V(size<0>(Cos) == size<0>(Sin)); // MMA_K - static_assert(decltype(size<0>(S))::value == decltype(size<0>(Cos))::value * 2); - static_assert(decltype(size<0>(Cos))::value % 2 == 0); // Since we do fast conversion from fp16/bf16 to fp32 - Tensor rCos = make_fragment_like(Cos); - Tensor rSin = make_fragment_like(Sin); - Tensor rS = make_fragment_like(S); - typedef __NATIVE_VECTOR__(2, float) Float2; - typedef __NATIVE_VECTOR__(4, int) VecTypeB128; - typedef __NATIVE_VECTOR__(2, int) VecTypeB64; - - #pragma unroll - for (int m = 0; m < size<1>(S); ++m) { - bool row_mask = get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN; - #pragma unroll - for (int k = 0; k < size<2>(S); ++k) { - bool mask = row_mask && (Is_even_K || get<1>(identity_MN(0, 0, k)) < dim); - auto S_ptr = (VecTypeB128 *)(S(_, m, k).data().ptr_); - auto rS_ptr = (VecTypeB128 *)(rS(_, m, k).data()); - rS_ptr[0] = __builtin_mxc_ldg_b128_predicator(S_ptr, 0, true, true, false, false, - mask, 1, MACA_ICMP_EQ); - bool rotary_mask = mask && (get<1>(identity_MN(0, 0, k)) < rotary_dim); - - auto gCos_ptr = (VecTypeB64 *)(Cos(_, m, k).data().ptr_); - auto gSin_ptr = (VecTypeB64 *)(Sin(_, m, k).data().ptr_); - auto rCos_ptr = (VecTypeB64 *)(rCos(_, m, k).data()); - auto rSin_ptr = (VecTypeB64 *)(rSin(_, m, k).data()); - rCos_ptr[0] = __builtin_mxc_ldg_b64_predicator(gCos_ptr, 0, true, true, false, false, - rotary_mask, 1, MACA_ICMP_EQ); - rSin_ptr[0] = __builtin_mxc_ldg_b64_predicator(gSin_ptr, 0, true, true, false, false, - rotary_mask, 1, MACA_ICMP_EQ); - - if (rotary_mask) { - using T = typename Engine0::value_type; - using T_rotary = typename Engine1::value_type; - CONVERT_TENSOR_TYPE(T, float, rS(_, m, k), S_fp32) - CONVERT_TENSOR_TYPE(T_rotary, float, rCos(_, m, k), cos_fp32) - CONVERT_TENSOR_TYPE(T_rotary, float, rSin(_, m, k), sin_fp32) - - #pragma unroll - for (int i = 0; i < size<0>(rS) / 2; ++i) { - Float2 x_vec = {S_fp32(2 * i), S_fp32(2 * i + 1)}; - Float2 real_vec = {cos_fp32(i), sin_fp32(i)}; - Float2 imag_vec = {sin_fp32(i), cos_fp32(i)}; - Float2 beta_vec = {0.0f, 0.0f}; - real_vec = __builtin_mxc_pk_fma_f32(x_vec, real_vec, beta_vec); - imag_vec = __builtin_mxc_pk_fma_f32(x_vec, imag_vec, beta_vec); - S_fp32(2 * i) = real_vec[0] - real_vec[1]; - S_fp32(2 * i + 1) = imag_vec[0] + imag_vec[1]; - } - // Idk but I need to copy for the convert_type to work - Tensor S_fp32_copy = make_fragment_like(S_fp32); - cute::copy(S_fp32, S_fp32_copy); - //Tensor S_og_type = convert_type(S_fp32_copy); - CONVERT_TENSOR_TYPE(float, T, S_fp32_copy, S_og_type) - cute::copy(S_og_type, rS(_, m, k)); - } - auto D_ptr = (VecTypeB128 *)(D(_, m, k).data().ptr_); - __builtin_mxc_stg_b128_predicator(D_ptr, 0, rS_ptr[0], true, false, true, mask, 1, MACA_ICMP_EQ); - } - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -__forceinline__ __device__ void copy_rotary_contiguous_to_global(Tensor const &S, - Tensor &D, - Tensor const &Cos, - Tensor const &Sin, - Tensor const &identity_MN, - const int max_MN, const int min_MN, - const int dim, const int rotary_dim) { - CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); - CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); - CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Cos)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Cos)); // MMA_K - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(Sin)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(Sin)); // MMA_K - CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(Cos)); // MMA - CUTE_STATIC_ASSERT_V(size<0>(Cos) == size<0>(Sin)); - static_assert(decltype(size<0>(Cos))::value % 2 == 0); // Since we do fast conversion from fp16/bf16 to fp32 - Tensor rCos = make_fragment_like(Cos); - Tensor rSin = make_fragment_like(Sin); - Tensor rS = make_fragment_like(S); - Tensor rS_other = make_fragment_like(rS(_, 0, 0)); - typedef __NATIVE_VECTOR__(2, float) Float2; - Float2 beta_vec = {0.0f, 0.0f}; - - const int rotary_dim_half = rotary_dim >> 1; - - typedef __NATIVE_VECTOR__(4, int) VecType; - #pragma unroll - for (int m = 0; m < size<1>(S); ++m) { - bool row_mask = get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN; - #pragma unroll - for (int k = 0; k < size<2>(S); ++k) { - bool mask = row_mask && (Is_even_K || get<1>(identity_MN(0, 0, k)) < dim); - auto S_ptr = (VecType *)(S(_, m, k).data().ptr_); - auto rS_ptr = (VecType *)(rS(_, m, k).data()); - rS_ptr[0] = __builtin_mxc_ldg_b128_predicator(S_ptr, 0, true, true, false, false, - mask, 1, MACA_ICMP_EQ); - - bool rotary_mask = mask && (get<1>(identity_MN(0, 0, k)) < rotary_dim); - const bool is_left = get<1>(identity_MN(0, 0, k)) < rotary_dim_half; - auto gS_ptr = (VecType *)(S(_, m, k).data().ptr_ + (is_left ? rotary_dim_half : -rotary_dim_half)); - auto rS_other_ptr = (VecType *)(rS_other.data()); - rS_other_ptr[0] = __builtin_mxc_ldg_b128_predicator(gS_ptr, 0, true, true, false, false, - rotary_mask, 1, MACA_ICMP_EQ); - - auto gCos_ptr = (VecType *)(Cos(_, m, k).data().ptr_ + (is_left ? 0 : -rotary_dim_half)); - auto gSin_ptr = (VecType *)(Sin(_, m, k).data().ptr_ + (is_left ? 0 : -rotary_dim_half)); - auto rCos_ptr = (VecType *)(rCos(_, m, k).data()); - auto rSin_ptr = (VecType *)(rSin(_, m, k).data()); - rCos_ptr[0] = __builtin_mxc_ldg_b128_predicator(gCos_ptr, 0, true, true, false, false, - rotary_mask, 1, MACA_ICMP_EQ); - rSin_ptr[0] = __builtin_mxc_ldg_b128_predicator(gSin_ptr, 0, true, true, false, false, - rotary_mask, 1, MACA_ICMP_EQ); - if (rotary_mask) { - using T = typename Engine0::value_type; - using T_rotary = typename Engine1::value_type; - CONVERT_TENSOR_TYPE(T, float, rS(_, m, k), S_fp32) - CONVERT_TENSOR_TYPE(T, float, rS_other, S_other_fp32) - CONVERT_TENSOR_TYPE(T_rotary, float, rCos(_, m, k), cos_fp32) - CONVERT_TENSOR_TYPE(T_rotary, float, rSin(_, m, k), sin_fp32) - - #pragma unroll - for (int i = 0; i < size<0>(rS); ++i) { - // S_fp32(i) = S_fp32(i) * cos_fp32(i) + S_other_fp32(i) * (is_left ? -sin_fp32(i) : sin_fp32(i)); - Float2 x_vec = {S_fp32(i), S_other_fp32(i)}; - Float2 alpha_vec = {cos_fp32(i), is_left ? -sin_fp32(i) : sin_fp32(i)}; - Float2 y_vec = __builtin_mxc_pk_fma_f32(x_vec, alpha_vec, beta_vec); - S_fp32(i) = y_vec[0] + y_vec[1]; - } - // Idk but I need to copy for the convert_type to work - Tensor S_fp32_copy = make_fragment_like(S_fp32); - cute::copy(S_fp32, S_fp32_copy); - //Tensor S_og_type = convert_type(S_fp32_copy); - CONVERT_TENSOR_TYPE(float, T, S_fp32_copy, S_og_type) - cute::copy(S_og_type, rS(_, m, k)); - } - auto D_ptr = (VecType *)(D(_, m, k).data().ptr_); - __builtin_mxc_stg_b128_predicator(D_ptr, 0, rS_ptr[0], true, false, true, mask, 1, MACA_ICMP_EQ); - } - } - -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace flash diff --git a/csrc/flash_kernel/feature/softmax.h b/csrc/flash_kernel/feature/softmax.h index ca42178..0ce6e68 100644 --- a/csrc/flash_kernel/feature/softmax.h +++ b/csrc/flash_kernel/feature/softmax.h @@ -12,7 +12,6 @@ #include -#include "philox.cuh" #include "utils.h" namespace flash { @@ -41,7 +40,7 @@ __device__ __forceinline__ void quad_allreduce_(Tensor &dst, T CUTE_STATIC_ASSERT_V(size(dst) == size(src)); #pragma unroll for (int i = 0; i < size(dst); i++){ - dst(i) = Allreduce<64>::run(src(i), op); + dst(i) = Partialreduce::run(src(i), op); } } @@ -220,6 +219,50 @@ struct Softmax { #pragma unroll for (int mi = 0; mi < size<0>(scores); mi++) { if constexpr(AddVec) { + Float2 x_vec = {row_sum(mi), 0.0f}; + Float2 scale_vec = {1.0f, 1.0f}; + #pragma unroll + for (int ni = 0; ni < size<1>(scores); ni += 2) { + Float2 beta_vec = {scores(mi, ni), scores(mi, ni + 1)}; + x_vec = __builtin_mxc_pk_fma_f32(x_vec, scale_vec, beta_vec); + } + row_sum(mi) = x_vec[0] + x_vec[1]; + } + else { + #pragma unroll + for (int ni = 0; ni < size<1>(scores); ni++) { + row_sum(mi) += scores(mi, ni); + } + } + } + } + }; + + template + __forceinline__ __device__ void softmax_rescale_o(Tensor0 &acc_s, Tensor1 &acc_o, Tensor2 &sRowMax, float softmax_scale_log2) { + // Reshape acc_s from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N)) + Tensor scores = make_tensor(acc_s.data(), flash::convert_layout_acc_rowcol(acc_s.layout())); + MaxOp max_op; + static_assert(decltype(size<0>(scores))::value == kNRows); + static_assert(decltype(size<1>(scores))::value % 2 == 0); + typedef __NATIVE_VECTOR__(2, float) Float2; + const int tidx = threadIdx.x; + const int wave_idx = tidx / 64; + const int lane_idx = tidx % 64; + const int wave_group_idx = wave_idx / 4; + const int row_offset = wave_idx % 4 * 16 + lane_idx % 16; + if constexpr (Is_first) { + flash::template thread_reduce_(scores, row_max, max_op); + flash::template quad_allreduce_(row_max, row_max, max_op); + if (lane_idx / 16 == 0) { + sRowMax(wave_group_idx, row_offset) = row_max(0); //sts row_max + } + flash::sync_threads(); + row_max(0) = max(row_max(0), sRowMax(wave_group_idx ^ 1, row_offset)); //lds row_max + flash::scale_apply_exp2(scores, row_max, softmax_scale_log2); + if constexpr(AddVec) { + #pragma unroll + for (int mi = 0; mi < size<0>(scores); mi++) { Float2 x_vec = { 0.0f, 0.0f}; Float2 scale_vec = {1.0f, 1.0f}; #pragma unroll @@ -227,7 +270,159 @@ struct Softmax { Float2 beta_vec = {scores(mi, ni), scores(mi, ni + 1)}; x_vec = __builtin_mxc_pk_fma_f32(x_vec, scale_vec, beta_vec); } - row_sum(mi) += x_vec[0] + x_vec[1]; + row_sum(mi) = x_vec[0] + x_vec[1]; + } + } else { + SumOp sum_op; + flash::thread_reduce_(scores, row_sum, sum_op); + } + } else { + Tensor scores_max_prev = make_fragment_like(row_max); + cute::copy(row_max, scores_max_prev); + flash::template thread_reduce_(scores, row_max, max_op); + flash::template quad_allreduce_(row_max, row_max, max_op); + if (lane_idx / 16 == 0) { + sRowMax(wave_group_idx, row_offset) = row_max(0); //sts row_max + } + flash::sync_threads(); + row_max(0) = max(row_max(0), sRowMax(wave_group_idx ^ 1, row_offset)); //lds row_max + // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K)) + Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout())); + static_assert(decltype(size<0>(acc_o_rowcol))::value == kNRows); + static_assert(decltype(size<1>(acc_o_rowcol))::value % 2 == 0); + #pragma unroll + for (int mi = 0; mi < size(row_max); ++mi) { + float scores_max_cur = !Check_inf + ? row_max(mi) + : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi)); + float scores_scale = __builtin_exp2f((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2); + row_sum(mi) *= scores_scale; + Float2 scale_vec = {scores_scale , scores_scale}; + Float2 beta_vec = {0.0f, 0.0f}; + #pragma unroll + for (int ni = 0; ni < size<1>(acc_o_rowcol); ni += 2) { + Float2 x_vec = {acc_o_rowcol(mi, ni), acc_o_rowcol(mi, ni + 1)}; + x_vec = __builtin_mxc_pk_fma_f32(x_vec, scale_vec, beta_vec); + acc_o_rowcol(mi, ni) = x_vec[0]; + acc_o_rowcol(mi, ni + 1) = x_vec[1]; + } + } + flash::scale_apply_exp2(scores, row_max, softmax_scale_log2); + #pragma unroll + for (int mi = 0; mi < size<0>(scores); mi++) { + if constexpr(AddVec) { + Float2 x_vec = {row_sum(mi), 0.0f}; + Float2 scale_vec = {1.0f, 1.0f}; + #pragma unroll + for (int ni = 0; ni < size<1>(scores); ni += 2) { + Float2 beta_vec = {scores(mi, ni), scores(mi, ni + 1)}; + x_vec = __builtin_mxc_pk_fma_f32(x_vec, scale_vec, beta_vec); + } + row_sum(mi) = x_vec[0] + x_vec[1]; + } + else { + #pragma unroll + for (int ni = 0; ni < size<1>(scores); ni++) { + row_sum(mi) += scores(mi, ni); + } + } + } + } + } + + template + __forceinline__ __device__ void get_row_max(Tensor0 &acc_s, Tensor1 &scores_max_prev, Tensor2 &sRowMax,float softmax_scale_log2) { + Tensor scores = make_tensor(acc_s.data(), flash::convert_layout_acc_rowcol(acc_s.layout())); + MaxOp max_op; + static_assert(decltype(size<0>(scores))::value == kNRows); + static_assert(decltype(size<1>(scores))::value % 2 == 0); + const int tidx = threadIdx.x; + const int wave_idx = tidx / 64; + const int lane_idx = tidx % 64; + const int wave_group_idx = wave_idx / 4; + const int row_offset = wave_idx % 4 * 16 + lane_idx % 16; + if constexpr (Is_first) { + flash::template thread_reduce_(scores, row_max, max_op); + flash::template quad_allreduce_(row_max, row_max, max_op); + if (lane_idx / 16 == 0) { + sRowMax(wave_group_idx, row_offset) = row_max(0); //sts row_max + } + flash::sync_threads(); + row_max(0) = max(row_max(0), sRowMax(wave_group_idx ^ 1, row_offset)); //lds row_max + } else { + cute::copy(row_max, scores_max_prev); + flash::template thread_reduce_(scores, row_max, max_op); + flash::template quad_allreduce_(row_max, row_max, max_op); + if (lane_idx / 16 == 0) { + sRowMax(wave_group_idx, row_offset) = row_max(0); //sts row_max + } + flash::sync_threads(); + row_max(0) = max(row_max(0), sRowMax(wave_group_idx ^ 1, row_offset)); //lds row_max + } + } + + template + __forceinline__ __device__ void softmax_rescale_o_without_row_max(Tensor0 &acc_s, Tensor1 &acc_o, Tensor2 &scores_max_prev, float softmax_scale_log2) { + // Reshape acc_s from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N)) + Tensor scores = make_tensor(acc_s.data(), flash::convert_layout_acc_rowcol(acc_s.layout())); + static_assert(decltype(size<0>(scores))::value == kNRows); + static_assert(decltype(size<1>(scores))::value % 2 == 0); + typedef __NATIVE_VECTOR__(2, float) Float2; + if constexpr (Is_first) { + flash::scale_apply_exp2(scores, row_max, softmax_scale_log2); + if constexpr (AddVec) { + #pragma unroll + for (int mi = 0; mi < size<0>(scores); mi++) { + Float2 x_vec = {0.0f, 0.0f}; + Float2 scale_vec = {1.0f, 1.0f}; + #pragma unroll + for (int ni = 0; ni < size<1>(scores); ni += 2) { + Float2 beta_vec = {scores(mi, ni), scores(mi, ni + 1)}; + x_vec = __builtin_mxc_pk_fma_f32(x_vec, scale_vec, beta_vec); + } + row_sum(mi) = x_vec[0] + x_vec[1]; + } + } + else { + SumOp sum_op; + flash::thread_reduce_(scores, row_sum, sum_op); + } + } else { + // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K)) + Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout())); + static_assert(decltype(size<0>(acc_o_rowcol))::value == kNRows); + static_assert(decltype(size<1>(acc_o_rowcol))::value % 2 == 0); + #pragma unroll + for (int mi = 0; mi < size(row_max); ++mi) { + float scores_max_cur = !Check_inf + ? row_max(mi) + : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi)); + float scores_scale = __builtin_exp2f((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2); + row_sum(mi) *= scores_scale; + // #pragma unroll + // for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) { acc_o_rowcol(mi, ni) *= scores_scale; } + Float2 scale_vec = {scores_scale , scores_scale}; + Float2 beta_vec = {0.0f, 0.0f}; + #pragma unroll + for (int ni = 0; ni < size<1>(acc_o_rowcol); ni += 2) { + Float2 x_vec = {acc_o_rowcol(mi, ni), acc_o_rowcol(mi, ni + 1)}; + x_vec = __builtin_mxc_pk_fma_f32(x_vec, scale_vec, beta_vec); + acc_o_rowcol(mi, ni) = x_vec[0]; + acc_o_rowcol(mi, ni + 1) = x_vec[1]; + } + } + flash::scale_apply_exp2(scores, row_max, softmax_scale_log2); + #pragma unroll + for (int mi = 0; mi < size<0>(scores); mi++) { + if constexpr(AddVec) { + Float2 x_vec = {row_sum(mi), 0.0f}; + Float2 scale_vec = {1.0f, 1.0f}; + #pragma unroll + for (int ni = 0; ni < size<1>(scores); ni += 2) { + Float2 beta_vec = {scores(mi, ni), scores(mi, ni + 1)}; + x_vec = __builtin_mxc_pk_fma_f32(x_vec, scale_vec, beta_vec); + } + row_sum(mi) = x_vec[0] + x_vec[1]; } else { #pragma unroll @@ -240,18 +435,17 @@ struct Softmax { }; template - __forceinline__ __device__ TensorT normalize_softmax_lse(Tensor0 &acc_o, float softmax_scale, float rp_dropout=1.0) { + __forceinline__ __device__ TensorT normalize_softmax_lse(Tensor0 &acc_o, float softmax_scale, float rp_dropout=1.0, float k_descale=1.0) { flash::quadreduce_sum(row_sum); TensorT lse = make_fragment_like(row_sum); Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout())); static_assert(decltype(size<0>(acc_o_rowcol))::value == kNRows); static_assert(decltype(size<1>(acc_o_rowcol))::value % 2 == 0); - typedef __NATIVE_VECTOR__(2, float) Float2; #pragma unroll for (int mi = 0; mi < size<0>(acc_o_rowcol); ++mi) { float sum = row_sum(mi); - float inv_sum = (sum == 0.f || sum != sum) ? 1.f : 1.f / sum; + float inv_sum = (sum == 0.f || sum != sum) ? 1.f : k_descale / sum; if (Return_lse) lse(mi) = (sum == 0.f || sum != sum) ? (Split ? -INFINITY : INFINITY) : row_max(mi) * softmax_scale + __logf(sum); float scale = !Is_dropout ? inv_sum : inv_sum * rp_dropout; @@ -271,6 +465,45 @@ struct Softmax { } return lse; }; + + template + __forceinline__ __device__ TensorT normalize_softmax_lse(Tensor0 &acc_o, Tensor1 &sRowSum, float softmax_scale, float rp_dropout=1.0) { + const int tidx = threadIdx.x; + const int wave_idx = tidx / 64; + const int lane_idx = tidx % 64; + const int wave_group_idx = wave_idx / 4; + const int row_offset = wave_idx % 4 * 16 + lane_idx % 16; + flash::quadreduce_sum(row_sum); + if (lane_idx / 16 == 0) { + sRowSum(wave_group_idx, row_offset) = row_sum(0); //sts row_max + } + flash::sync_threads(); + row_sum(0) += sRowSum(wave_group_idx ^ 1, row_offset); //lds row_max + TensorT lse = make_fragment_like(row_sum); + Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout())); + static_assert(decltype(size<0>(acc_o_rowcol))::value == kNRows); + static_assert(decltype(size<1>(acc_o_rowcol))::value % 2 == 0); + + typedef __NATIVE_VECTOR__(2, float) Float2; + #pragma unroll + for (int mi = 0; mi < size<0>(acc_o_rowcol); ++mi) { + float sum = row_sum(mi); + float inv_sum = (sum == 0.f || sum != sum) ? 1.f : 1.f / sum; + if (Return_lse) + lse(mi) = (sum == 0.f || sum != sum) ? (Split ? -INFINITY : INFINITY) : row_max(mi) * softmax_scale + __logf(sum); + float scale = !Is_dropout ? inv_sum : inv_sum * rp_dropout; + Float2 scale_vec = {scale, scale}; + Float2 beta_vec = {0.0f, 0.0f}; + #pragma unroll + for (int ni = 0; ni < size<1>(acc_o_rowcol); ni += 2) { + Float2 x_vec = {acc_o_rowcol(mi, ni), acc_o_rowcol(mi, ni + 1)}; + x_vec = __builtin_mxc_pk_fma_f32(x_vec, scale_vec, beta_vec); + acc_o_rowcol(mi, ni) = x_vec[0]; + acc_o_rowcol(mi, ni + 1) = x_vec[1]; + } + } + return lse; + }; }; } // namespace flash diff --git a/csrc/flash_kernel/flash_dense_mla_decode_kernel.h b/csrc/flash_kernel/flash_dense_mla_decode_kernel.h new file mode 100644 index 0000000..e37e538 --- /dev/null +++ b/csrc/flash_kernel/flash_dense_mla_decode_kernel.h @@ -0,0 +1,87 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#pragma once + +#include + +#include +#include +#include + +#include "block_info.h" +#include "kernel_traits.h" +#include "utils.h" +#include "softmax.h" +#include "mask.h" + + +#include "xcore1000/flash_fwd_mla_kernel_k64_16x16_4waves_stage1_xcore1000.h" +// #include "xcore1000/flash_fwd_mla_kernel_k64_16x16_4waves_xcore1000.h" +#include "xcore1000/flash_fwd_mla_kernel_k64_32x16_4waves_xcore1000.h" +#include "xcore1000/flash_fwd_mla_kernel_k64_64x16_8waves_xcore1000.h" +#include "xcore1500/flash_fwd_mla_kernel_k64_64x16_8waves_xcore1500.h" +#include "xcore1500/flash_fwd_mla_kernel_k64_64x32_8waves_xcore1500.h" + +#include "static_switch.h" + +namespace flash { + +using namespace cute; + +template +__forceinline__ __device__ void compute_attn_1rowblock_splitkv_mla(const Params ¶ms, const int m_block_max) { + constexpr int kBlockN = Kernel_traits::kBlockN; + const int m_block = blockIdx.x; + const int bidh = blockIdx.y; + const int partition_idx = blockIdx.z; + + extern __shared__ char shared_memory[]; + + int *tile_scheduler_metadata_ptr = params.tile_scheduler_metadata_ptr + partition_idx * TileSchedulerMetaDataSize; + // int4 tile_scheduler_metadata = __ldg(reinterpret_cast(tile_scheduler_metadata_ptr)); + int4 tile_scheduler_metadata = *reinterpret_cast(tile_scheduler_metadata_ptr); + int begin_idx = tile_scheduler_metadata.x; + int sched_begin_block_idx = tile_scheduler_metadata.y; + int end_idx = tile_scheduler_metadata.z; + int sched_end_block_idx = tile_scheduler_metadata.w; + if (begin_idx >= params.b || begin_idx < 0) return; + // int begin_n_split_idx = __ldg(tile_scheduler_metadata_ptr + 4); + int begin_n_split_idx = tile_scheduler_metadata_ptr[4]; + +#pragma unroll 1 + for (int batch_id = begin_idx; batch_id <= end_idx; ++batch_id) { + const int n_split_idx = batch_id == begin_idx ? begin_n_split_idx : 0; + const int seqlen_k = params.cu_seqlens_k[batch_id]; + const int n_block_min = batch_id == begin_idx ? sched_begin_block_idx : 0; + const int n_block_max = batch_id == end_idx ? sched_end_block_idx : cute::ceil_div(seqlen_k, kBlockN); + // [n_block_min, n_block_max) need be calculated in kernel + if (n_block_max <= n_block_min) continue; + const bool NoSplit = __ldg(params.num_splits_ptr + batch_id + 1) - __ldg(params.num_splits_ptr + batch_id) == 1; + if (batch_id > begin_idx) { + __syncthreads(); // Barrier between two tiles. + } + #if defined(__MACA_ARCH__) && (__MACA_ARCH__ == 1500) + if constexpr (Kernel_traits::kBlockM == 64 && Kernel_traits::kBlockN == 16 && Kernel_traits::kNWarps == 8) { + compute_attn_1rowblock_splitkv_mla_k64_64x16_8waves_xcore1500( + params, batch_id, bidh, m_block, n_split_idx, n_block_min, n_block_max, NoSplit); + }else if constexpr (Kernel_traits::kBlockM == 64 && Kernel_traits::kBlockN == 32 && Kernel_traits::kNWarps == 8) { + compute_attn_1rowblock_splitkv_mla_k64_64x32_8waves_xcore1500( + params, batch_id, bidh, m_block, n_split_idx, n_block_min, n_block_max, NoSplit); + } + #else defined(__MACA_ARCH__) && (__MACA_ARCH__ == 1000) + if constexpr (Kernel_traits::kBlockM == 32 && Kernel_traits::kBlockN == 16 && Kernel_traits::kNWarps == 4) { + compute_attn_1rowblock_splitkv_mla_k64_32x16_4waves_xcore1000( + params, batch_id, bidh, m_block, n_split_idx, n_block_min, n_block_max, NoSplit); + } else if constexpr (Kernel_traits::kBlockM == 64 && Kernel_traits::kBlockN == 16 && Kernel_traits::kNWarps == 8) { + compute_attn_1rowblock_splitkv_mla_k64_64x16_8waves_xcore1000( + params, batch_id, bidh, m_block, n_split_idx, n_block_min, n_block_max, NoSplit); + } else if constexpr (Kernel_traits::kBlockM == 16 && Kernel_traits::kBlockN == 16 && Kernel_traits::kNWarps == 4) { + compute_attn_1rowblock_splitkv_mla_k64_16x16_4waves_xcore1000( + params, batch_id, bidh, m_block, n_split_idx, n_block_min, n_block_max, NoSplit); + } + #endif + + } +} + +}// namespace flash diff --git a/csrc/flash_kernel/flash_fwd_split_kernel_k64_V1x8.h b/csrc/flash_kernel/flash_fwd_split_kernel_k64_V1x8.h deleted file mode 100644 index 78d16e6..0000000 --- a/csrc/flash_kernel/flash_fwd_split_kernel_k64_V1x8.h +++ /dev/null @@ -1,686 +0,0 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - -#pragma once - -#include - -#include -#include -#include - -#include "block_info.h" -#include "kernel_traits.h" -#include "utils.h" -#include "softmax.h" -#include "mask.h" -#include "rotary.h" -#include "attn_mask.h" - -namespace flash { - -using namespace cute; - -template -__forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_V1x8(const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int num_n_splits) { - - using Element = typename Kernel_traits::Element; - using ElementAccum = typename Kernel_traits::ElementAccum; - using index_t = typename Kernel_traits::index_t; - - // Shared memory. - extern __shared__ char smem_[]; - - // The thread index. - const int tidx = threadIdx.x; - - constexpr int kBlockM = Kernel_traits::kBlockM; - constexpr int kBlockN = Kernel_traits::kBlockN; - constexpr int kHeadDim = Kernel_traits::kHeadDim; - constexpr int kHeadDimV = Kernel_traits::kHeadDimV; - constexpr int kNWarps = Kernel_traits::kNWarps; - constexpr int kBlockKSmem = Kernel_traits::kBlockKSmem; - constexpr int kBlockKGmem = Kernel_traits::UseWarpsNx1 ? Kernel_traits::kBlockKSmem : 128; - constexpr int kAtomLayoutMS = Kernel_traits::kAtomLayoutMS; - constexpr int kAtomLayoutMO = Kernel_traits::kAtomLayoutMO; - - static_assert(kBlockKSmem == 64); - - using GmemTiledCopyO = std::conditional_t< - !Split, - typename Kernel_traits::GmemTiledCopyO, - typename Kernel_traits::GmemTiledCopyOaccum - >; - using ElementO = std::conditional_t; - - const BlockInfo binfo(params, bidb); - // if (threadIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { printf("Is_even_MN = %d, is_cumulativ = %d, seqlen_k_cache = %d, actual_seqlen_k = %d\n", Is_even_MN, params.is_seqlens_k_cumulative, binfo.seqlen_k_cache, binfo.actual_seqlen_k); } - // if (threadIdx.x == 0 && blockIdx.y == 1 && blockIdx.z == 0) { printf("params.knew_ptr = %p, seqlen_k_cache + seqlen_knew = %d\n", params.knew_ptr, binfo.seqlen_k_cache + (params.knew_ptr == nullptr ? 0 : params.seqlen_knew)); } - if (m_block * kBlockM >= binfo.actual_seqlen_q) return; - - const int n_blocks_per_split = ((binfo.actual_seqlen_k + kBlockN - 1) / kBlockN + num_n_splits - 1) / num_n_splits; - const int n_block_min = !Is_local - ? n_split_idx * n_blocks_per_split - : std::max(n_split_idx * n_blocks_per_split, (m_block * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q - params.window_size_left) / kBlockN); - int n_block_max = std::min(cute::ceil_div(binfo.actual_seqlen_k, kBlockN), (n_split_idx + 1) * n_blocks_per_split); - if (Is_causal || Is_local) { - n_block_max = std::min(n_block_max, - cute::ceil_div((m_block + 1) * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q / params.ngroups + params.window_size_right, kBlockN)); - } - - if (n_block_min >= n_block_max) { // This also covers the case where n_block_max <= 0 - // We exit early and write 0 to gOaccum and -inf to gLSEaccum. - // Otherwise we might read OOB elements from gK and gV, - // or get wrong results when we combine gOaccum from different blocks. - const index_t row_offset_o = binfo.q_offset(params.o_batch_stride, params.o_row_stride, bidb) - + m_block * kBlockM * params.o_row_stride + bidh * params.o_head_stride; - const index_t row_offset_oaccum = (((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q - + m_block * kBlockM) * params.d_v; - const index_t row_offset_lseaccum = ((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; - Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.oaccum_ptr : params.o_ptr) + (Split ? row_offset_oaccum : row_offset_o)), - Shape, Int>{}, - make_stride(Split ? kHeadDimV : params.o_row_stride, _1{})); - Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.softmax_lseaccum_ptr : params.softmax_lse_ptr) + row_offset_lseaccum), - Shape>{}, Stride<_1>{}); - - GmemTiledCopyO gmem_tiled_copy_Oaccum; - auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); - Tensor tOgOaccum = gmem_thr_copy_Oaccum.partition_D(gOaccum); - Tensor tOrOaccum = make_tensor(shape(tOgOaccum)); - clear(tOrOaccum); - // Construct identity layout for sO - Tensor cO = make_identity_tensor(make_shape(size<0>(gOaccum), size<1>(gOaccum))); // (BLK_M,BLK_K) -> (blk_m,blk_k) - // Repeat the partitioning with identity layouts - Tensor tOcO = gmem_thr_copy_Oaccum.partition_D(cO); - Tensor tOpO = make_tensor(make_shape(size<2>(tOgOaccum))); - if (!Is_even_K) { - #pragma unroll - for (int k = 0; k < size(tOpO); ++k) { tOpO(k) = get<1>(tOcO(0, 0, k)) < params.d_v; } - } - // Clear_OOB_K must be false since we don't want to write zeros to gmem - flash::copy( - gmem_tiled_copy_Oaccum, tOrOaccum, tOgOaccum, tOcO, tOpO, binfo.actual_seqlen_q - m_block * kBlockM - ); - #pragma unroll - for (int m = 0; m < size<1>(tOgOaccum); ++m) { - const int row = get<0>(tOcO(0, m, 0)); - if (row < binfo.actual_seqlen_q - m_block * kBlockM && get<1>(tOcO(0, m, 0)) == 0) { gLSEaccum(row) = Split ? -INFINITY : INFINITY; } - } - return; - } - - // We iterate over the blocks in reverse order. This is because the last block is the only one - // that needs masking when we read K and V from global memory. Moreover, iterating in reverse - // might save us 1 register (we just need n_block instead of both n_block and n_block_max). - const index_t row_offset_q = binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb) - + m_block * kBlockM * params.q_row_stride + bidh * params.q_head_stride; - // We move K and V to the last block. - const int bidb_cache = params.cache_batch_idx == nullptr ? bidb : params.cache_batch_idx[bidb]; - const int *block_table = params.block_table == nullptr ? nullptr : params.block_table + bidb * params.block_table_batch_stride; - const int block_table_idx = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN / params.page_block_size; - const int block_table_offset = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN - block_table_idx * params.page_block_size; - const index_t row_offset_k = block_table == nullptr - ? binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb_cache) - + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride - : (bidh / params.h_h_k_ratio) * params.k_head_stride; - const index_t row_offset_v = block_table == nullptr - ? binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb_cache) - + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride - : (bidh / params.h_h_k_ratio) * params.v_head_stride; - - Tensor gQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), - Shape, Int>{}, - make_stride(params.q_row_stride, _1{})); - Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + row_offset_k), - Shape, Int>{}, - make_stride(params.k_row_stride, _1{})); - // if (threadIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { printf("k_ptr = %p, row_offset_k = %d, gK_ptr = %p\n", params.k_ptr, row_offset_k, gK.data()); } - Tensor gV = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_ptr) + row_offset_v), - Shape, Int>{}, - make_stride(params.v_row_stride, _1{})); - - Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), - typename Kernel_traits::SmemLayoutQ{}); - //Tensor sK = make_tensor(sQ.data() + size(sQ), typename Kernel_traits::SmemLayoutKV{}); - Tensor sK = make_tensor(sQ.data() + (Kernel_traits::Share_Q_K_smem ? 0 : size(sQ)), - typename Kernel_traits::SmemLayoutK{}); - Tensor sV = make_tensor(sK.data() + size(sK), typename Kernel_traits::SmemLayoutVtNoSwizzle{}); - Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed{}); - Tensor sVtNoSwizzle = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{}); - - typename Kernel_traits::GmemTiledCopyB128 gmem_tiled_copy_Q; - auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx); - - Tensor tQgQ = gmem_thr_copy_Q.partition_S(gQ); - Tensor tQsQ = gmem_thr_copy_Q.partition_D(sQ); - - typename Kernel_traits::GmemTiledCopyB64 gmem_tiled_copy_KV; - auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx); - Tensor tKgK = gmem_thr_copy_KV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K) - Tensor tKsK = gmem_thr_copy_KV.partition_D(sK); - Tensor tVgV = gmem_thr_copy_KV.partition_S(gV); // (VCPY, VCPY_N, VCPY_K) - Tensor tVsV = gmem_thr_copy_KV.partition_D(sV); - Tensor tVrV = make_fragment_like(tVgV); - // wave0 and wave2 compute the same S, wave1 and wave3 compute the same S - int tidx_mma_s = tidx & 0x7F; - typename Kernel_traits::TiledMmaS tiled_mma_s; - auto thr_mma_s = tiled_mma_s.get_thread_slice(tidx_mma_s); - Tensor tSrQ = thr_mma_s.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K) - Tensor tSrK = thr_mma_s.partition_fragment_B(sK); // (MMA,MMA_N,MMA_K) - typename Kernel_traits::TiledMmaO tiled_mma_o; - auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); - Tensor tOrVt = thr_mma_o.partition_fragment_B(sVtNoSwizzle); // (MMA, MMA_K,MMA_N) - - Tensor acc_o = partition_fragment_C(tiled_mma_o, Shape, Int>{}); // MMA, MMA_M, MMA_K - - // - // Copy Atom retiling - // - - auto smem_tiled_copy_Q = make_tiled_copy_A(typename Kernel_traits::SmemCopyAtomB64{}, tiled_mma_s); - auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx_mma_s); - Tensor tSsQ = smem_thr_copy_Q.partition_S(sQ); - - auto smem_tiled_copy_K = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomB64{}, tiled_mma_s); - auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx_mma_s); - Tensor tSsK = smem_thr_copy_K.partition_S(sK); - - auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed{}, tiled_mma_o); - auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx); - Tensor tOsVt = smem_thr_copy_V.partition_S(sVtNoSwizzle); - - // PREDICATES - - // Construct identity layout for sQ and sK - Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) - Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) - - // Repeat the partitioning with identity layouts - Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) - Tensor tKVcKV = gmem_thr_copy_KV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) - - - // Prologue - - - // Read Q from gmem to smem, optionally apply rotary embedding. - Tensor tQrQ = make_fragment_like(tQgQ); - // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs - flash::copy_b128(tQgQ, tQrQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); - cute::copy(tQrQ, tQsQ); - - if constexpr (Kernel_traits::Is_Q_in_regs) { - flash::sync_threads(); - cute::copy(smem_tiled_copy_Q, tSsQ, tSrQ); - flash::sync_threads(); - } - - - int n_block = n_block_max - 1; - // We don't need to clear the sK smem tiles since we'll mask out the scores anyway. - Tensor tKrK = make_fragment_like(tKgK); - if constexpr (!Is_page_attn) { - flash::copy_b64(tKgK, tKrK, tKVcKV, params.d, binfo.actual_seqlen_k - n_block * kBlockN); - } else { - flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block, - block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, binfo.actual_seqlen_k - n_block * kBlockN); - } - - // flash::cp_async_wait<0>(); - // __syncthreads(); - // if (tidx == 0 && blockIdx.y == 0 && blockIdx.z == 0) { print(tKsK); } - // __syncthreads(); - - clear(acc_o); - - flash::Softmax(acc_o)> softmax; - - const float alibi_slope = !Has_alibi ? 0.0f : reinterpret_cast(params.alibi_slopes_ptr)[bidb * params.alibi_slopes_batch_stride + bidh] / params.scale_softmax; - flash::Mask mask(binfo.actual_seqlen_k, binfo.actual_seqlen_q, params.ngroups, params.window_size_left, params.window_size_right, alibi_slope); - - // For performance reason, we separate out two kinds of iterations: - // those that need masking on S, and those that don't. - // We need masking on S for the very last block when K and V has length not multiple of kBlockN. - // We also need masking on S if it's causal, for the last ceil_div(kBlockM, kBlockN) blocks. - // We will have at least 1 "masking" iteration. - - // If not even_N, then seqlen_k might end in the middle of a block. In that case we need to - // mask 2 blocks (e.g. when kBlockM == kBlockN), not just 1. - constexpr int n_masking_steps = (!Is_causal && !Is_local) - ? 1 - : ((Is_even_MN && Is_causal) ? cute::ceil_div(kBlockM, kBlockN) : cute::ceil_div(kBlockM, kBlockN) + 1); - #pragma unroll - for (int masking_step = 0; masking_step < n_masking_steps; ++masking_step, --n_block) { - Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) - cute::copy(tKrK, tKsK); - clear(acc_s); - - // Advance gV - if (masking_step > 0) { - if constexpr (!Is_page_attn) { - tVgV.data() = tVgV.data() + (-int(kBlockN * params.v_row_stride)); - flash::copy_b64(tVgV, tVrV, tKVcKV, params.d_v); - } else { - flash::copy_b64_page_one(gV, tVgV, tVrV, tKVcKV, params.d_v, n_block, - block_table, params.v_batch_stride, params.v_row_stride, params.page_block_size); - } - } else { - if constexpr (!Is_page_attn) { - // Clear the smem tiles to account for predicated off loads - flash::copy_b64( - tVgV, tVrV, tKVcKV, params.d_v, binfo.actual_seqlen_k - n_block * kBlockN - ); - } else { - flash::copy_b64_page_one(gV, tVgV, tVrV, tKVcKV, params.d_v, n_block, - block_table, params.v_batch_stride, params.v_row_stride, params.page_block_size, binfo.actual_seqlen_k - n_block * kBlockN); - } - } - flash::sync_threads(); - - flash::gemm_opt( - acc_s, tSrQ, tSrK, tSsQ, tSsK, tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, - smem_thr_copy_Q, smem_thr_copy_K - ); - // if (cute::thread0()) { print(acc_s); } - if constexpr (Is_softcap){ - flash::apply_softcap(acc_s, params.softcap); - } - - mask.template apply_mask( - acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 - ); - - cute::copy(tVrV, tVsV); - - if (n_block > n_block_min) { - // Advance gK - if constexpr (!Is_page_attn) { - tKgK.data() = tKgK.data() + (-int(kBlockN * params.k_row_stride)); - flash::copy_b64(tKgK, tKrK, tKVcKV, params.d); - } else { - flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, - block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size); - } - } - - // We have key_padding_mask so we'll need to Check_inf - masking_step == 0 - ? softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2) - : softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); - // if (cute::thread0()) { print(scores_max); print(scores_sum); print(scores); } - - // Convert acc_s from fp32 to fp16/bf16 - //Tensor rP = flash::convert_type(acc_s); - CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) - // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) - // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. - //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); - Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); - - flash::gemm_rs(acc_o, tOrP, tOrVt, tOsVt, tiled_mma_o, smem_tiled_copy_V, smem_thr_copy_V); - - // This check is at the end of the loop since we always have at least 1 iteration - if (n_masking_steps > 1 && n_block <= n_block_min) { - --n_block; - break; - } - } - - // These are the iterations where we don't need masking on S - for (; n_block >= n_block_min; --n_block) { - Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) - cute::copy(tKrK, tKsK); - clear(acc_s); - // Advance gV - if constexpr (!Is_page_attn) { - tVgV.data() = tVgV.data() + (-int(kBlockN * params.v_row_stride)); - flash::copy_b64(tVgV, tVrV, tKVcKV, params.d_v); - } else { - flash::copy_b64_page_one(gV, tVgV, tVrV, tKVcKV, params.d_v, n_block, - block_table, params.v_batch_stride, params.v_row_stride, params.page_block_size); - } - flash::sync_threads(); - - flash::gemm_opt( - acc_s, tSrQ, tSrK, tSsQ, tSsK, tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, - smem_thr_copy_Q, smem_thr_copy_K - ); - cute::copy(tVrV, tVsV); - - if (n_block > n_block_min) { - // Advance gK - if constexpr (!Is_page_attn) { - tKgK.data() = tKgK.data() + (-int(kBlockN * params.k_row_stride)); - flash::copy_b64(tKgK, tKrK, tKVcKV, params.d); - } else { - flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, - block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size); - } - } - - if constexpr (Is_softcap){ - flash::apply_softcap(acc_s, params.softcap); - } - - mask.template apply_mask( - acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 - ); - softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); - - //Tensor rP = flash::convert_type(acc_s); - CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) - // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) - // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. - //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); - Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); - - flash::gemm_rs(acc_o, tOrP, tOrVt, tOsVt, tiled_mma_o, smem_tiled_copy_V, smem_thr_copy_V); - } - - // Epilogue - - - Tensor lse = softmax.template normalize_softmax_lse(acc_o, params.scale_softmax); - // if (cute::thread0()) { print(lse); } - if constexpr (!Split) { - // use smem for O (mtreg->smem->mtreg->global) - Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) - // Partition sO to match the accumulator partitioning - using SmemTiledCopyO = typename Kernel_traits::SmemCopyAtomO; - auto smem_tiled_copy_Oaccum = make_tiled_copy_C(SmemTiledCopyO{}, tiled_mma_o); - auto smem_thr_copy_Oaccum = smem_tiled_copy_Oaccum.get_thread_slice(tidx); - //Tensor rO = flash::convert_type(acc_o); - CONVERT_TENSOR_TYPE(ElementAccum, ElementO, acc_o, rO) - Tensor taccOrOaccum = smem_thr_copy_Oaccum.retile_S(rO); // ((Atom,AtomNum), MMA_M, MMA_N) - Tensor taccOsOaccum = smem_thr_copy_Oaccum.partition_D(sOaccum); // ((Atom,AtomNum),PIPE_M,PIPE_N) - - // sOaccum is larger than sQ, so we need to syncthreads here - // TODO: allocate enough smem for sOaccum - if constexpr (Kernel_traits::Share_Q_K_smem) { flash::sync_threads(); } - - cute::copy(smem_tiled_copy_Oaccum, taccOrOaccum, taccOsOaccum); - - const index_t row_offset_o = binfo.q_offset(params.o_batch_stride, params.o_row_stride, bidb) - + m_block * kBlockM * params.o_row_stride + bidh * params.o_head_stride; - const index_t row_offset_oaccum = (((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q - + m_block * kBlockM) * params.d_v; - const index_t row_offset_lseaccum = ((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; - - Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.o_ptr) + (row_offset_o)), - Shape, Int>{}, - make_stride(params.o_row_stride, _1{})); - Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lse_ptr) + row_offset_lseaccum), - Shape>{}, Stride<_1>{}); - // if (tidx == 0) { printf("row_offset_o = %d, bidh = %d, gOaccum = %p\n", row_offset_o, bidh, gOaccum.data()); } - - GmemTiledCopyO gmem_tiled_copy_Oaccum; - auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); - Tensor tOsOaccum = gmem_thr_copy_Oaccum.partition_S(sOaccum); // ((Atom,AtomNum),ATOM_M,ATOM_N) - Tensor tOgOaccum = gmem_thr_copy_Oaccum.partition_D(gOaccum); - - flash::sync_threads(); - - Tensor tOrOaccum = make_tensor(shape(tOgOaccum)); - cute::copy(gmem_tiled_copy_Oaccum, tOsOaccum, tOrOaccum); - - Tensor caccO = make_identity_tensor(Shape, Int>{}); // (BLK_M,BLK_K) -> (blk_m,blk_k) - Tensor taccOcO = thr_mma_o.partition_C(caccO); // (MMA,MMA_M,MMA_K) - static_assert(decltype(size<0>(taccOcO))::value == 4); - // Convert to ((2, 2), MMA_M, MMA_K) then take only the row indices. - Tensor taccOcO_row = logical_divide(taccOcO, Shape<_4>{})(make_coord(0, _), _, 0); - CUTE_STATIC_ASSERT_V(size(lse) == size(taccOcO_row)); // MMA_M - if (get<1>(taccOcO_row(0)) == 0) { - #pragma unroll - for (int mi = 0; mi < size(lse); ++mi) { - const int row = get<0>(taccOcO_row(mi)); - if (row < binfo.actual_seqlen_q - m_block * kBlockM) { gLSEaccum(row) = lse(mi); } - } - } - - // Construct identity layout for sO - Tensor cO = make_identity_tensor(make_shape(size<0>(sOaccum), size<1>(sOaccum))); // (BLK_M,BLK_K) -> (blk_m,blk_k) - // Repeat the partitioning with identity layouts - Tensor tOcO = gmem_thr_copy_Oaccum.partition_D(cO); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) - // Clear_OOB_K must be false since we don't want to write zeros to gmem - flash::copy_reg_to_global( - tOrOaccum, tOgOaccum, tOcO, params.d_v, binfo.actual_seqlen_q - m_block * kBlockM - ); - } else { - // don't use smem for O (mtreg->global) - const index_t row_offset_oaccum = (((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q - + m_block * kBlockM) * params.d_v; - const index_t row_offset_lseaccum = ((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; - - Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.oaccum_ptr) + row_offset_oaccum), - Shape, Int>{}, - make_stride(kHeadDimV, _1{})); - Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lseaccum_ptr) + row_offset_lseaccum), - Shape>{}, Stride<_1>{}); - // if (tidx == 0) { printf("row_offset_o = %d, bidh = %d, gOaccum = %p\n", row_offset_o, bidh, gOaccum.data()); } - using GmemCopyAtomOaccum = typename Kernel_traits::SmemCopyAtomOaccum; - auto gmem_tiled_copy_Oaccum = make_tiled_copy_C(GmemCopyAtomOaccum{}, tiled_mma_o); - auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); - Tensor taccOrOaccum = gmem_thr_copy_Oaccum.retile_S(acc_o); // ((Atom,AtomNum), MMA_M, MMA_N) - Tensor taccOgOaccum = gmem_thr_copy_Oaccum.partition_D(gOaccum); - - - - Tensor caccO = make_identity_tensor(Shape, Int>{}); // (BLK_M,BLK_K) -> (blk_m,blk_k) - Tensor taccOcO = thr_mma_o.partition_C(caccO); // (MMA,MMA_M,MMA_K) - static_assert(decltype(size<0>(taccOcO))::value == 4); - // Convert to ((2, 2), MMA_M, MMA_K) then take only the row indices. - Tensor taccOcO_row = logical_divide(taccOcO, Shape<_4>{})(make_coord(0, _), _, 0); - CUTE_STATIC_ASSERT_V(size(lse) == size(taccOcO_row)); // MMA_M - if (get<1>(taccOcO_row(0)) == 0) { - #pragma unroll - for (int mi = 0; mi < size(lse); ++mi) { - const int row = get<0>(taccOcO_row(mi)); - if (row < binfo.actual_seqlen_q - m_block * kBlockM) { gLSEaccum(row) = lse(mi); } - } - } - - // Clear_OOB_K must be false since we don't want to write zeros to gmem - flash::copy_reg_to_global( - taccOrOaccum, taccOgOaccum, taccOcO, params.d_v, binfo.actual_seqlen_q - m_block * kBlockM - ); - } -} - -template -__forceinline__ __device__ void compute_attn_splitkv(const Params ¶ms, const int m_block_max) { - const int m_block = blockIdx.x; - // The block index for the batch. - const int bidb = Split ? blockIdx.z / params.h : blockIdx.y; - // The block index for the head. - const int bidh = Split ? blockIdx.z - bidb * params.h : blockIdx.z; - const int n_split_idx = Split ? blockIdx.y : 0; - const int num_n_splits = Split ? gridDim.y : 1; - compute_attn_1rowblock_splitkv_k64_mla_V1x8( - params, bidb, bidh, m_block, n_split_idx, num_n_splits); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -__forceinline__ __device__ void combine_attn_seqk_parallel(const Params ¶ms) { - using Element = typename Kernel_traits::Element; - using ElementAccum = typename Kernel_traits::ElementAccum; - using index_t = typename Kernel_traits::index_t; - constexpr int kMaxSplits = 1 << Log_max_splits; - constexpr int kHeadDimV = Kernel_traits::kHeadDimV; - constexpr int kNThreads = 256;/*Kernel_traits::kNThreads*/; - - static_assert(kMaxSplits <= 128, "kMaxSplits must be <= 128"); - static_assert(kBlockM == 4 || kBlockM == 8 || kBlockM == 16 || kBlockM == 32, "kBlockM must be 4, 8, 16 or 32"); - static_assert(kNThreads == 128 || kNThreads == 256, "We assume that each block has 128 or 256 threads"); - - // Shared memory. - // kBlockM + 1 instead of kBlockM to reduce bank conflicts. - __shared__ ElementAccum sLSE[kMaxSplits][kBlockM + 1]; - - // The thread and block index. - const int tidx = threadIdx.x; - const int bidx = blockIdx.x; - - const index_t lse_size = params.b * params.h * params.seqlen_q; - - const index_t row_offset_lse = bidx * kBlockM; - Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lseaccum_ptr) + row_offset_lse), - Shape, Int>{}, - make_stride(lse_size, _1{})); - // LSE format is different depending on params.unpadded_lse and params.seqlenq_ngroups_swapped, see comment in get_lse_tile. - // This tensor's layout maps row_offset_lse to {bidb, bidh, lse_size}. - Tensor gLSE = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lse_ptr) + row_offset_lse), - Shape>{}, Stride<_1>{}); - - // This layout maps row_offset_lse to {bidh, lse_size, bidb} or {bidh, bidb, lse_size}. - Layout flat_layout = make_layout(lse_size); - Layout orig_layout = make_layout(make_shape(params.seqlen_q, params.h, params.b)); - auto transposed_stride = make_stride(params.b, params.seqlen_q * params.b, params.seqlen_q / params.ngroups); - Layout remapped_layout = make_layout(make_shape(params.seqlen_q, params.h, params.b), transposed_stride); - Layout final_layout = cute::composition(remapped_layout, cute::composition(orig_layout, flat_layout)); - - Tensor gLSE_unpadded = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lse_ptr)), final_layout); - - constexpr int kNLsePerThread = (kMaxSplits * kBlockM + kNThreads - 1) / kNThreads; - - // Read the LSE values from gmem and store them in shared memory, then tranpose them. - constexpr int kRowsPerLoadLSE = kNThreads / kBlockM; - typedef __NATIVE_VECTOR__(1, ElementAccum) B32Type; - #pragma unroll - for (int l = 0; l < kNLsePerThread; ++l) { - const int row = l * kRowsPerLoadLSE + tidx / kBlockM; - const int col = tidx % kBlockM; - ElementAccum lse = (row < params.num_splits && col < lse_size - bidx * kBlockM) ? gLSEaccum(row, col) : -INFINITY; - if (row < kMaxSplits) { sLSE[row][col] = lse; } - } - - flash::sync_threads(); - Tensor lse_accum = make_tensor(Shape>{}); - constexpr int kRowsPerLoadTranspose = std::min(kRowsPerLoadLSE, kMaxSplits); - // To make sure that kMaxSplits is within 1 warp: we decide how many elements within kMaxSplits - // each thread should hold. If kMaxSplits = 16, then each thread holds 2 elements (128 threads, - // kBlockM rows, so each time we load we can load 128 / kBlockM rows). - // constexpr int kThreadsPerSplit = kMaxSplits / kRowsPerLoadTranspose; - // static_assert(kThreadsPerSplit <= 32); - //static_assert(kRowsPerLoadTranspose <= 32); - static_assert(kRowsPerLoadTranspose <= 64); - static_assert(kNLsePerThread * kRowsPerLoadTranspose <= kMaxSplits); - const int lse_base_row = tidx % kRowsPerLoadTranspose; - const int lse_base_col = tidx / kRowsPerLoadTranspose; - #pragma unroll - for (int l = 0; l < kNLsePerThread; ++l) { - const int row = l * kRowsPerLoadTranspose + lse_base_row; - const int col = lse_base_col; - lse_accum(l) = (row < kMaxSplits && col < kBlockM) ? sLSE[row][col] : -INFINITY; - // if (bidx == 0 && tidx < 32) { printf("tidx = %d, row = %d, col = %d, lse = %f\n", tidx, row, col, lse_accum(l)); } - } - - // Compute the logsumexp of the LSE along the split dimension. - ElementAccum lse_max = lse_accum(0); - #pragma unroll - for (int l = 1; l < kNLsePerThread; ++l) { lse_max = max(lse_max, lse_accum(l)); } - MaxOp max_op; - lse_max = Allreduce::run(lse_max, max_op); - lse_max = lse_max == -INFINITY ? 0.0f : lse_max; // In case all local LSEs are -inf - float lse_sum = __expf(lse_accum(0) - lse_max); - #pragma unroll - for (int l = 1; l < kNLsePerThread; ++l) { lse_sum += __expf(lse_accum(l) - lse_max); } - SumOp sum_op; - lse_sum = Allreduce::run(lse_sum, sum_op); - // For the case where all local lse == -INFINITY, we want to set lse_logsum to INFINITY. Otherwise - // lse_logsum is log(0.0) = -INFINITY and we get NaN when we do lse_accum(l) - lse_logsum. - ElementAccum lse_logsum = (lse_sum == 0.f || lse_sum != lse_sum) ? INFINITY : __logf(lse_sum) + lse_max; - if (tidx % kRowsPerLoadTranspose == 0 && tidx / kRowsPerLoadTranspose < kBlockM) { - if (params.unpadded_lse) { - const index_t lse_offset = row_offset_lse + tidx / kRowsPerLoadTranspose; - if (lse_offset < lse_size) { - gLSE_unpadded(lse_offset) = lse_logsum; - } - } else { - gLSE(tidx / kRowsPerLoadTranspose) = lse_logsum; - } - } - // Store the scales exp(lse - lse_logsum) in shared memory. - #pragma unroll - for (int l = 0; l < kNLsePerThread; ++l) { - const int row = l * kRowsPerLoadTranspose + lse_base_row; - const int col = lse_base_col; - if (row < params.num_splits && col < kBlockM) { sLSE[row][col] = __expf(lse_accum(l) - lse_logsum); } - } - - const index_t row_offset_oaccum = bidx * kBlockM * params.d_v; - Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.oaccum_ptr) + row_offset_oaccum), - Shape, Int>{}, - Stride, _1>{}); - constexpr int kBlockN = kNThreads / kBlockM; - using GmemLayoutAtomOaccum = Layout, Int>, Stride, _1>>; - using GmemTiledCopyOaccum = decltype( - make_tiled_copy(Copy_Atom{}, - GmemLayoutAtomOaccum{}, - Layout>{})); // Val layout, 4 vals per store - GmemTiledCopyOaccum gmem_tiled_copy_Oaccum; - auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); - Tensor tOgOaccum = gmem_thr_copy_Oaccum.partition_S(gOaccum); - Tensor tOrO = make_tensor(shape(tOgOaccum)); - Tensor tOrOaccum = make_tensor(shape(tOgOaccum)); - clear(tOrO); - flash::sync_threads(); - - typedef __NATIVE_VECTOR__(2, float) Float2; - - // Predicates - Tensor cOaccum = make_identity_tensor(Shape, Int>{}); - // Repeat the partitioning with identity layouts - Tensor tOcOaccum = gmem_thr_copy_Oaccum.partition_S(cOaccum); - static_assert(decltype(size<0>(tOrOaccum))::value % 2 == 0); - // Load Oaccum in then scale and accumulate to O - for (int split = 0; split < params.num_splits; ++split) { - flash::copy_b128( - tOgOaccum, tOrOaccum, tOcOaccum, params.d_v, lse_size - bidx * kBlockM - ); - #pragma unroll - for (int m = 0; m < size<1>(tOrOaccum); ++m) { - int row = get<0>(tOcOaccum(0, m, 0)); - ElementAccum lse_scale = sLSE[split][row]; - Float2 lse_scale_vec = {lse_scale, lse_scale}; - #pragma unroll - for (int k = 0; k < size<2>(tOrOaccum); ++k) { - #pragma unroll - for (int i = 0; i < size<0>(tOrOaccum); i += 2) { - Float2 x_vec = {tOrOaccum(i, m, k), tOrOaccum(i + 1, m, k)}; - Float2 y_vec = {tOrO(i, m, k), tOrO(i + 1, m, k)}; - y_vec = __builtin_mxc_pk_fma_f32(x_vec, lse_scale_vec, y_vec); - tOrO(i, m, k) = y_vec[0]; - tOrO(i + 1, m, k) = y_vec[1]; - } - } - } - tOgOaccum.data() = tOgOaccum.data() + lse_size * params.d_v; - } - - //Tensor rO = flash::convert_type(tOrO); - CONVERT_TENSOR_TYPE(ElementAccum, Element, tOrO, rO) - const int q_head_offset = params.h * params.seqlen_q; - // Write to gO - #pragma unroll - for (int m = 0; m < size<1>(rO); ++m) { - const int idx = bidx * kBlockM + get<0>(tOcOaccum(0, m, 0)); - const int batch_idx = idx / q_head_offset; - const int head_idx = (idx - batch_idx * q_head_offset) / params.seqlen_q; - // The index to the rows of Q - const int row = idx - batch_idx * q_head_offset - head_idx * params.seqlen_q; - auto o_ptr = reinterpret_cast(params.o_ptr) + batch_idx * params.o_batch_stride - + head_idx * params.o_head_stride + row * params.o_row_stride; - #pragma unroll - for (int k = 0; k < size<2>(rO); ++k) { - const int col = get<1>(tOcOaccum(0, m, k)); - Tensor gO = make_tensor(make_gmem_ptr(o_ptr + col), - Shape(rO))::value>>{}, Stride<_1>{}); - auto gO_ptr = reinterpret_cast(gO.data().ptr_); - auto rO_ptr = reinterpret_cast(rO(_, m, k).data().ptr_); - __builtin_mxc_stg_b64_predicator(gO_ptr, 0, rO_ptr[0], true, false, false, idx < lse_size && (Is_even_K || col < params.d_v), 1, MACA_ICMP_EQ); - } - } -} - -} // namespace flash diff --git a/csrc/flash_kernel/flash_fwd_split_kernel.h b/csrc/flash_kernel/flash_fwd_splitkv_mla_combine_kernel.h similarity index 69% rename from csrc/flash_kernel/flash_fwd_split_kernel.h rename to csrc/flash_kernel/flash_fwd_splitkv_mla_combine_kernel.h index ec16cc1..2698d15 100644 --- a/csrc/flash_kernel/flash_fwd_split_kernel.h +++ b/csrc/flash_kernel/flash_fwd_splitkv_mla_combine_kernel.h @@ -1,5 +1,3 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - #pragma once #include @@ -8,56 +6,25 @@ #include #include -#include "block_info.h" #include "kernel_traits.h" #include "utils.h" -#include "softmax.h" -#include "mask.h" -#include "rotary.h" -#include "attn_mask.h" -#include "flash_fwd_split_kernel_k64_16x16_4waves.h" -#include "flash_fwd_split_kernel_k64_32x16_4waves.h" -#include "flash_fwd_split_kernel_k64_64x16_8waves.h" namespace flash { using namespace cute; -template -__forceinline__ __device__ void compute_attn_splitkv(const Params ¶ms, const int m_block_max) { - const int m_block = blockIdx.x; - // The block index for the batch. - const int bidb = Split ? blockIdx.z / params.h : blockIdx.y; - // The block index for the head. - const int bidh = Split ? blockIdx.z - bidb * params.h : blockIdx.z; - const int n_split_idx = Split ? blockIdx.y : 0; - const int num_n_splits = Split ? gridDim.y : 1; - if constexpr (Kernel_traits::kBlockM == 32 && Kernel_traits::kBlockN == 16 && Kernel_traits::kNWarps == 4) { - compute_attn_1rowblock_splitkv_k64_mla_32x16_4waves( - params, bidb, bidh, m_block, n_split_idx, num_n_splits); - } else if constexpr (Kernel_traits::kBlockM == 64 && Kernel_traits::kBlockN == 16 && Kernel_traits::kNWarps == 8) { - compute_attn_1rowblock_splitkv_k64_mla_64x16_8waves( - params, bidb, bidh, m_block, n_split_idx, num_n_splits); - } else if constexpr (Kernel_traits::kBlockM == 16 && Kernel_traits::kBlockN == 16 && Kernel_traits::kNWarps == 4) { - compute_attn_1rowblock_splitkv_k64_mla_16x16_4waves( - params, bidb, bidh, m_block, n_split_idx, num_n_splits); - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - template -__forceinline__ __device__ void combine_attn_seqk_parallel(const Params ¶ms) { - using Element = typename Kernel_traits::Element; +__forceinline__ __device__ void combine_attn_seqk_parallel_splitkv_mla(const Params ¶ms) { + using ElementO = typename Kernel_traits::ElementO; using ElementAccum = typename Kernel_traits::ElementAccum; using index_t = typename Kernel_traits::index_t; constexpr int kMaxSplits = 1 << Log_max_splits; constexpr int kHeadDimV = Kernel_traits::kHeadDimV; - constexpr int kNThreads = 256;/*Kernel_traits::kNThreads*/; + constexpr int kNThreads = kBlockM == 1 ? 64 : (kBlockM == 2 ? 128 : 256);/*Kernel_traits::kNThreads*/; - static_assert(kMaxSplits <= 128, "kMaxSplits must be <= 128"); - static_assert(kBlockM == 4 || kBlockM == 8 || kBlockM == 16 || kBlockM == 32, "kBlockM must be 4, 8, 16 or 32"); - static_assert(kNThreads == 128 || kNThreads == 256, "We assume that each block has 128 or 256 threads"); + // static_assert(kMaxSplits <= 128, "kMaxSplits must be <= 128"); + static_assert(kBlockM == 1 || kBlockM == 2 || kBlockM == 4 || kBlockM == 8 || kBlockM == 16 || kBlockM == 32, "kBlockM must be 1, 2, 4, 8, 16 or 32"); + static_assert(kNThreads == 64 || kNThreads == 128 || kNThreads == 256, "We assume that each block has 64, 128 or 256 threads"); // Shared memory. // kBlockM + 1 instead of kBlockM to reduce bank conflicts. @@ -67,10 +34,20 @@ __forceinline__ __device__ void combine_attn_seqk_parallel(const Params ¶ms) const int tidx = threadIdx.x; const int bidx = blockIdx.x; - const index_t lse_size = params.b * params.h * params.seqlen_q; + const int hs = params.h * params.seqlen_q; + const int batch_idx = (bidx * kBlockM) / hs; + const int hs_idx = (bidx * kBlockM) % hs; + const int split_offset = params.num_splits_ptr[batch_idx]; + const int actual_num_splits = params.num_splits_ptr[batch_idx + 1] - split_offset; + FLASH_DEVICE_ASSERT(actual_num_splits <= kMaxSplits); + if (actual_num_splits == 1) return; + + const index_t lse_size = hs; + + const index_t row_offset_lseaccum = split_offset * lse_size + hs_idx; const index_t row_offset_lse = bidx * kBlockM; - Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lseaccum_ptr) + row_offset_lse), + Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lseaccum_ptr) + row_offset_lseaccum), Shape, Int>{}, make_stride(lse_size, _1{})); // LSE format is different depending on params.unpadded_lse and params.seqlenq_ngroups_swapped, see comment in get_lse_tile. @@ -96,7 +73,7 @@ __forceinline__ __device__ void combine_attn_seqk_parallel(const Params ¶ms) for (int l = 0; l < kNLsePerThread; ++l) { const int row = l * kRowsPerLoadLSE + tidx / kBlockM; const int col = tidx % kBlockM; - ElementAccum lse = (row < params.num_splits && col < lse_size - bidx * kBlockM) ? gLSEaccum(row, col) : -INFINITY; + ElementAccum lse = (row < actual_num_splits && col < lse_size - hs_idx) ? gLSEaccum(row, col) : -INFINITY; if (row < kMaxSplits) { sLSE[row][col] = lse; } } @@ -118,7 +95,6 @@ __forceinline__ __device__ void combine_attn_seqk_parallel(const Params ¶ms) const int row = l * kRowsPerLoadTranspose + lse_base_row; const int col = lse_base_col; lse_accum(l) = (row < kMaxSplits && col < kBlockM) ? sLSE[row][col] : -INFINITY; - // if (bidx == 0 && tidx < 32) { printf("tidx = %d, row = %d, col = %d, lse = %f\n", tidx, row, col, lse_accum(l)); } } // Compute the logsumexp of the LSE along the split dimension. @@ -151,10 +127,10 @@ __forceinline__ __device__ void combine_attn_seqk_parallel(const Params ¶ms) for (int l = 0; l < kNLsePerThread; ++l) { const int row = l * kRowsPerLoadTranspose + lse_base_row; const int col = lse_base_col; - if (row < params.num_splits && col < kBlockM) { sLSE[row][col] = __expf(lse_accum(l) - lse_logsum); } + if (row < actual_num_splits && col < kBlockM) { sLSE[row][col] = __expf(lse_accum(l) - lse_logsum); } } - const index_t row_offset_oaccum = bidx * kBlockM * params.d_v; + const index_t row_offset_oaccum = (split_offset * hs + hs_idx) * params.d_v; Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.oaccum_ptr) + row_offset_oaccum), Shape, Int>{}, Stride, _1>{}); @@ -170,7 +146,8 @@ __forceinline__ __device__ void combine_attn_seqk_parallel(const Params ¶ms) Tensor tOrO = make_tensor(shape(tOgOaccum)); Tensor tOrOaccum = make_tensor(shape(tOgOaccum)); clear(tOrO); - flash::sync_threads(); + flash::sync_threads(); // first barrier + flash::barrier(); // second barrier typedef __NATIVE_VECTOR__(2, float) Float2; @@ -180,9 +157,9 @@ __forceinline__ __device__ void combine_attn_seqk_parallel(const Params ¶ms) Tensor tOcOaccum = gmem_thr_copy_Oaccum.partition_S(cOaccum); static_assert(decltype(size<0>(tOrOaccum))::value % 2 == 0); // Load Oaccum in then scale and accumulate to O - for (int split = 0; split < params.num_splits; ++split) { + for (int split = 0; split < actual_num_splits; ++split) { flash::copy_b128( - tOgOaccum, tOrOaccum, tOcOaccum, params.d_v, lse_size - bidx * kBlockM + tOgOaccum, tOrOaccum, tOcOaccum, params.d_v, lse_size - hs_idx ); #pragma unroll for (int m = 0; m < size<1>(tOrOaccum); ++m) { @@ -201,21 +178,20 @@ __forceinline__ __device__ void combine_attn_seqk_parallel(const Params ¶ms) } } } - tOgOaccum.data() = tOgOaccum.data() + lse_size * params.d_v; + tOgOaccum.data() = tOgOaccum.data() + hs * params.d_v; } - //Tensor rO = flash::convert_type(tOrO); - CONVERT_TENSOR_TYPE(ElementAccum, Element, tOrO, rO) + //Tensor rO = flash::convert_type(tOrO); + CONVERT_TENSOR_TYPE(ElementAccum, ElementO, tOrO, rO) const int q_head_offset = params.h * params.seqlen_q; // Write to gO #pragma unroll for (int m = 0; m < size<1>(rO); ++m) { - const int idx = bidx * kBlockM + get<0>(tOcOaccum(0, m, 0)); - const int batch_idx = idx / q_head_offset; - const int head_idx = (idx - batch_idx * q_head_offset) / params.seqlen_q; + const int idx = hs_idx + get<0>(tOcOaccum(0, m, 0)); + const int head_idx = idx / params.seqlen_q; // The index to the rows of Q - const int row = idx - batch_idx * q_head_offset - head_idx * params.seqlen_q; - auto o_ptr = reinterpret_cast(params.o_ptr) + batch_idx * params.o_batch_stride + const int row = idx % params.seqlen_q; + auto o_ptr = reinterpret_cast(params.o_ptr) + batch_idx * params.o_batch_stride + head_idx * params.o_head_stride + row * params.o_row_stride; #pragma unroll for (int k = 0; k < size<2>(rO); ++k) { @@ -229,4 +205,4 @@ __forceinline__ __device__ void combine_attn_seqk_parallel(const Params ¶ms) } } -} // namespace flash +}// namespace flash diff --git a/csrc/flash_kernel/flash_sparse_mla_decode_kernel.h b/csrc/flash_kernel/flash_sparse_mla_decode_kernel.h new file mode 100644 index 0000000..32da5c8 --- /dev/null +++ b/csrc/flash_kernel/flash_sparse_mla_decode_kernel.h @@ -0,0 +1,63 @@ +#pragma once + +#include + +#include +#include +#include + +#include "block_info.h" +#include "kernel_traits.h" +#include "utils.h" +#include "softmax.h" +#include "mask.h" + +#include "xcore1000/flash_fwd_sparse_mla_kernel_k64_64x16_8waves_xcore1000.h" + + + +namespace flash { + +using namespace cute; + +template +__forceinline__ __device__ void compute_attn_1rowblock_splitkv_sparse_mla(const Params ¶ms, const int m_block_max) { + constexpr int kBlockN = Kernel_traits::kBlockN; + const int m_block = blockIdx.x; + const int bidh = blockIdx.y; + const int partition_idx = blockIdx.z; + + extern __shared__ char shared_memory[]; + + int *tile_scheduler_metadata_ptr = params.tile_scheduler_metadata_ptr + partition_idx * TileSchedulerMetaDataSize; + int4 tile_scheduler_metadata = *reinterpret_cast(tile_scheduler_metadata_ptr); + int begin_idx = tile_scheduler_metadata.x; + int sched_begin_block_idx = tile_scheduler_metadata.y; + int end_idx = tile_scheduler_metadata.z; + int sched_end_block_idx = tile_scheduler_metadata.w; + if (begin_idx >= params.b || begin_idx < 0) return; + int begin_n_split_idx = tile_scheduler_metadata_ptr[4]; + +#pragma unroll 1 + for (int batch_id = begin_idx; batch_id <= end_idx; ++batch_id) { + const int n_split_idx = batch_id == begin_idx ? begin_n_split_idx : 0; + const int seqlen_k = params.cu_seqlens_k[batch_id]; + const int n_block_min = batch_id == begin_idx ? sched_begin_block_idx : 0; + const int n_block_max = batch_id == end_idx ? sched_end_block_idx : cute::ceil_div(params.topk, kBlockN); + // [n_block_min, n_block_max) need be calculated in kernel + if (n_block_max <= n_block_min) continue; + const bool NoSplit = __ldg(params.num_splits_ptr + batch_id + 1) - __ldg(params.num_splits_ptr + batch_id) == 1; + if (batch_id > begin_idx) { + __syncthreads(); // Barrier between two tiles. + } + + // printf( + // "batch_id is %d, begin_idx is %d, end_idx is %d, params.b is %d, n_split_idx is %d, seqlen_k is %d, n_block_min is %d, n_block_max is %d, NoSplit is %d\n", + // batch_id, begin_idx, end_idx, params.b, n_split_idx, seqlen_k, n_block_min, n_block_max, NoSplit); + + compute_attn_1rowblock_splitkv_sparse_mla_k64_64x16_8waves_xcore1000( + params, batch_id, bidh, m_block, n_split_idx, n_block_min, n_block_max, NoSplit); + } +} + +}// namespace flash diff --git a/csrc/flash_kernel/kernel_traits.h b/csrc/flash_kernel/kernel_traits.h index 6f46b1c..b3f2ff1 100644 --- a/csrc/flash_kernel/kernel_traits.h +++ b/csrc/flash_kernel/kernel_traits.h @@ -17,33 +17,29 @@ using namespace cute; template struct Flash_kernel_traits { -#if defined(__MACA_ARCH__) using Element = elem_type; static constexpr bool Has_cp_async = false; -#else - using Element = mctlass::half_t; - static constexpr bool Has_cp_async = false; -#endif using ElementAccum = float; using index_t = int64_t; -#if defined(__MACA_ARCH__) - using MMA_Atom_Arch = std::conditional_t< - std::is_same_v, + using MMA_Atom_Arch_16x16x16_fp16 = std::conditional_t, MMA_Atom, MMA_Atom >; + using MMA_Atom_Arch_16x16x32_fp16 = std::conditional_t, + MMA_Atom, + MMA_Atom + >; + using MMA_Atom_Arch_16x16x32_i8 = MMA_Atom; using ValLayoutMNK = Layout>; -#else - using MMA_Atom_Arch = MMA_Atom; - using ValLayoutMNK = Layout>; -#endif using SmemCopyAtom = Copy_Atom; using SmemCopyAtomTransposed = Copy_Atom; - using SmemCopyB64 = Copy_Atom, elem_type>; - using UniversalCopyAtom32 = Copy_Atom, elem_type>; + using UniversalCopyAtomB32 = Copy_Atom, elem_type>; + using UniversalCopyAtomB64 = Copy_Atom, elem_type>; + using UniversalCopyAtomB128 = Copy_Atom, elem_type>; + using LDSB64Trans4x16Atom = Copy_Atom, elem_type>; }; // If Share_Q_K_smem is true, that forces Is_Q_in_regs to be true @@ -56,11 +52,17 @@ struct Flash_fwd_kernel_traits : public Base { static constexpr bool Has_cp_async = Base::Has_cp_async; using SmemCopyAtom = typename Base::SmemCopyAtom; - using SmemCopyAtomB64 = typename Base::SmemCopyB64; - using UniversalCopyAtom32 = typename Base::UniversalCopyAtom32; - + using UniversalCopyAtomB32 = typename Base::UniversalCopyAtomB32; + using UniversalCopyAtomB64 = typename Base::UniversalCopyAtomB64; + using UniversalCopyAtomB128 = typename Base::UniversalCopyAtomB128; + using LDSB64Trans4x16Atom = typename Base::LDSB64Trans4x16Atom; using SmemCopyAtomTransposed = typename Base::SmemCopyAtomTransposed; + using MMA_Atom_Arch_16x16x16_fp16 = typename Base::MMA_Atom_Arch_16x16x16_fp16; + using MMA_Atom_Arch_16x16x32_fp16 = typename Base::MMA_Atom_Arch_16x16x32_fp16; + using MMA_Atom_Arch_16x16x32_i8 = typename Base::MMA_Atom_Arch_16x16x32_i8; + + using ElementO = Element; static constexpr bool Share_Q_K_smem = Share_Q_K_smem_; static constexpr bool Is_Q_in_regs = Is_Q_in_regs_ || Share_Q_K_smem; static constexpr int Num_Stages = Num_Stages_; @@ -73,6 +75,8 @@ struct Flash_fwd_kernel_traits : public Base { static constexpr int kBlockN = kBlockN_; static constexpr int kHeadDim = kHeadDim_; static constexpr int kHeadDimV = kHeadDimV_; + static constexpr int kHeadDimNope = kHeadDimV; + static constexpr int kHeadDimRope = kHeadDim - kHeadDimV; static_assert(kHeadDim % 32 == 0); static constexpr int kBlockKSmem = kHeadDim % 64 == 0 ? 64 : 32; static constexpr int kBlockKSmemV = kHeadDimV % 64 == 0 ? 64 : 32; @@ -83,17 +87,55 @@ struct Flash_fwd_kernel_traits : public Base { static constexpr int SShift_OPT = kBlockKSmem == 32 ? 3 : 4; // for bank conflict free static constexpr int kAtomLayoutMS = std::min(kBlockM / 16, kNWarps); static constexpr int kAtomLayoutMO = kAtomLayoutMS; + static constexpr int kBlockTopK = kNThreads * (16 / sizeof(int32_t)); + + using MMA_Atom_QK = MMA_Atom_Arch_16x16x16_fp16; + using MMA_Atom_PV = MMA_Atom_Arch_16x16x16_fp16; using TiledMmaS = TiledMMA< - typename Base::MMA_Atom_Arch, + MMA_Atom_QK, Layout,_1,_1>>, // 2x1x1 or 4x1x1 typename Base::ValLayoutMNK>; + using TiledMmaS_16x16x32 = TiledMMA< + MMA_Atom_Arch_16x16x32_fp16, + Layout,_1,_1>>, // 2x1x1 or 4x1x1 + typename Base::ValLayoutMNK>; + + using TiledMmaS_16x16x32_4x2 = TiledMMA< + MMA_Atom_Arch_16x16x32_fp16, + Layout,Int,_1>>, // 4x2x1 + typename Base::ValLayoutMNK>; + using TiledMmaO = TiledMMA< - typename Base::MMA_Atom_Arch, + MMA_Atom_PV, Layout,Int,_1>>, // 2x2x1 or 4x2x1 typename Base::ValLayoutMNK>; + using SmemLayoutAtomRowMax = decltype( + composition(Swizzle<0, 0, 0>{}, + Layout>, + Stride, _1>>{})); + using SmemLayoutRowMax = decltype(tile_to_shape( + SmemLayoutAtomRowMax{}, + Shape, Int>{})); //rowmax 16 value per wave + + using SmemLayoutAtomRowSum = decltype( + composition(Swizzle<0, 0, 0>{}, + Layout>, + Stride, _1>>{})); + using SmemLayoutRowSum = decltype(tile_to_shape( + SmemLayoutAtomRowSum{}, + Shape, Int>{})); //rowsum 64 value per wave + + using SmemLayoutAtomP = decltype( + composition(Swizzle<3, 2, 4>{}, + Layout,Int>, + Stride, _1>>{})); + using SmemLayoutP = decltype(tile_to_shape( + SmemLayoutAtomP{}, + Shape,Int>{})); //rowmax_wg0 + using SmemLayoutAtomQ = decltype( composition(Swizzle{}, // This has to be kBlockKSmem, using kHeadDim gives wrong results for d=128 @@ -103,6 +145,10 @@ struct Flash_fwd_kernel_traits : public Base { composition(Swizzle<4, 2, 4>{}, Layout>, Stride, _1>>{})); + using SmemLayoutQNoSwizzle = decltype(tile_to_shape( + Layout>, + Stride, _1>>{}, + Shape, Int>{})); using SmemLayoutQ = decltype(tile_to_shape( SmemLayoutAtomQ{}, Shape, Int>{})); @@ -122,31 +168,55 @@ struct Flash_fwd_kernel_traits : public Base { SmemLayoutAtomQ{}, Shape, Int>{})); - using SmemLayoutAtomK = decltype( + using SmemLayoutAtomKNoswizzle = Layout>, + Stride, _1>>; + using SmemLayoutAtomK424 = decltype( composition(Swizzle<4, 2, 4>{}, Layout>, Stride, _1>>{})); - using SmemLayoutK = decltype(tile_to_shape( - SmemLayoutAtomK{}, + using SmemLayoutAtomK242 = decltype( + composition(Swizzle<2, 4, 2>{}, + Layout>, + Stride, _1>>{})); + using SmemLayoutAtomK333 = decltype( + composition(Swizzle<3, 3, 3>{}, + Layout>, + Stride, _1>>{})); + using SmemLayoutK424 = decltype(tile_to_shape( + SmemLayoutAtomK424{}, + Shape, Int, Int>{})); + using SmemLayoutK242 = decltype(tile_to_shape( + SmemLayoutAtomK242{}, + Shape, Int, Int>{})); + using SmemLayoutK333 = decltype(tile_to_shape( + SmemLayoutAtomK333{}, + Shape, Int, Int>{})); + using SmemLayoutKNoswizzle = decltype(tile_to_shape( + SmemLayoutAtomKNoswizzle{}, Shape, Int, Int>{})); using SmemLayoutV = decltype(tile_to_shape( SmemLayoutAtomQ{}, Shape, Int>{})); - // This has to be kBlockN and not 8, otherwise we get wrong results for d=128 - using SmemLayoutAtomVtransposedNoSwizzle = Layout, Int>, - Stride<_1, Int>>; - using SmemLayoutAtomVtransposed = decltype( + using SmemLayoutAtomVtransposedNoSwizzle = Layout, Int, Int>, + Stride<_1, Int, Int>>; + using SmemLayoutAtomVtransposed424 = decltype( composition(Swizzle<4, 2, 4>{}, SmemLayoutAtomVtransposedNoSwizzle{})); - using SmemLayoutVtransposed = decltype(tile_to_shape( - SmemLayoutAtomVtransposed{}, - Shape, Int>{})); + using SmemLayoutVtransposed424 = decltype(tile_to_shape( + SmemLayoutAtomVtransposed424{}, + Shape, Int, Int>{})); + + using SmemLayoutAtomVtransposed242 = decltype( + composition(Swizzle<2, 4, 2>{}, SmemLayoutAtomVtransposedNoSwizzle{})); + using SmemLayoutVtransposed242 = decltype(tile_to_shape( + SmemLayoutAtomVtransposed242{}, + Shape, Int, Int>{})); // Maybe the VtransposeNoSwizzle just needs to have the right shape // And the strides don't matter? using SmemLayoutVtransposedNoSwizzle = decltype(tile_to_shape( SmemLayoutAtomVtransposedNoSwizzle{}, - Shape, Int>{})); + Shape, Int, Int>{})); using SmemLayoutVtNoSwizzle = decltype(tile_to_shape( Layout>, @@ -160,15 +230,16 @@ struct Flash_fwd_kernel_traits : public Base { using SmemLayoutO = decltype(tile_to_shape( SmemLayoutAtomO{}, Shape, Int>{})); - using SmemCopyAtomO = Copy_Atom, Element>; + using SmemCopyAtomOb128 = Copy_Atom, ElementO>; + using SmemCopyAtomO = Copy_Atom, ElementO>; using SmemCopyAtomOaccum = Copy_Atom, ElementAccum>; + static constexpr int kSmemOSize = size(SmemLayoutO{}) * sizeof(ElementAccum); static constexpr int kSmemQSize = size(SmemLayoutQ{}) * sizeof(Element); - static constexpr int kSmemKSize = size(SmemLayoutK{}) * sizeof(Element); + static constexpr int kSmemKSize = size(SmemLayoutK424{}) * sizeof(Element); static constexpr int kSmemVSize = size(SmemLayoutV{}) * sizeof(Element); static constexpr int kSmemKVSize = kSmemKSize + kSmemVSize; - static constexpr int kSmemSize = Share_Q_K_smem ? std::max(std::min((Is_Splits_ ? 2 : 1) * kSmemQSize, 64 * 1024), kSmemKSize) : kSmemQSize + kSmemKSize; - static constexpr int kRegSize = kSmemSize / sizeof(uint32_t) / kNThreads; + static constexpr int kSmemSize = Share_Q_K_smem ? std::max(std::max(kSmemQSize, kSmemKSize), kSmemOSize) : std::max(kSmemQSize + kSmemKSize, kSmemOSize); static constexpr int kGmemElemsPerLoadB128 = sizeof(cute::uint128_t) / sizeof(Element); static constexpr int kGmemElemsPerLoadB64 = sizeof(cute::uint64_t) / sizeof(Element); @@ -196,30 +267,37 @@ struct Flash_fwd_kernel_traits : public Base { using GmemLayoutAtomB32 = Layout, Int>, Stride, _1>>; - static constexpr int kGmemThreadsPerRowV = kBlockKSmemV / kGmemElemsPerLoadB128; - static_assert(kNThreads % kGmemThreadsPerRowV == 0, "kNThreads must be a multiple of kGmemThreadsPerRow"); - using GmemLayoutAtomV = Layout, Int>, - Stride, _1>>; + static constexpr int kGmemElemsPerLoadB128O = sizeof(cute::uint128_t) / sizeof(ElementO); + static constexpr int kGmemThreadsPerRowO = kBlockKSmemV / kGmemElemsPerLoadB128O; + static_assert(kNThreads % kGmemThreadsPerRowO == 0, "kNThreads must be a multiple of kGmemThreadsPerRowO"); + static constexpr bool UseWarpsNx1 = kBlockM % (kNThreads / kGmemThreadsPerRowO) == 0; + using GmemLayoutAtomO = std::conditional_t< + UseWarpsNx1, + Layout, Int>, + Stride, _1>>, + Layout, Shape, Int>>, + Stride, Stride<_1, Int>>> + >; // We use CACHEGLOBAL instead of CACHEALWAYS for both Q and K/V, since we won't be reading // from the same address by the same threadblock. This is slightly faster. using GmemTiledCopyB128 = decltype( make_tiled_copy(Copy_Atom, Element>{}, GmemLayoutAtomB128{}, - Layout>{})); // Val layout, 8 vals per read + Layout>>{})); // Val layout, 8 vals per read using GmemTiledCopyB64 = decltype( make_tiled_copy(Copy_Atom, Element>{}, GmemLayoutAtomB64{}, - Layout>{})); // Val layout, 4 vals per read + Layout>>{})); // Val layout, 4 vals per read using GmemTiledCopyB32 = decltype( make_tiled_copy(Copy_Atom, Element>{}, GmemLayoutAtomB32{}, - Layout>{})); // Val layout, 2 vals per read + Layout>>{})); // Val layout, 2 vals per read using GmemTiledCopyO = decltype( - make_tiled_copy(Copy_Atom, Element>{}, - GmemLayoutAtomV{}, - Layout>{})); // Val layout, 8 vals per store + make_tiled_copy(Copy_Atom, ElementO>{}, + GmemLayoutAtomO{}, + Layout>>{})); // Val layout, 8 vals per store using GmemLayoutAtomOaccum = std::conditional_t< kBlockKSmem == 32, diff --git a/csrc/flash_kernel/xcore1000/flash_fwd_mla_kernel_k64_16x16_4waves_stage1_xcore1000.h b/csrc/flash_kernel/xcore1000/flash_fwd_mla_kernel_k64_16x16_4waves_stage1_xcore1000.h new file mode 100644 index 0000000..afed4d4 --- /dev/null +++ b/csrc/flash_kernel/xcore1000/flash_fwd_mla_kernel_k64_16x16_4waves_stage1_xcore1000.h @@ -0,0 +1,470 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#pragma once + +#include + +#include +#include +#include + +#include "block_info.h" +#include "kernel_traits.h" +#include "utils.h" +#include "softmax.h" +#include "mask.h" + +namespace flash { + +using namespace cute; +template +__forceinline__ __device__ void store_16x16(const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, char* smem_, AccO acc_o, Softmax softmax ) { + using ElementAccum = typename Kernel_traits::ElementAccum; + using Element = typename Kernel_traits::Element; + using index_t = typename Kernel_traits::index_t; + + using GmemTiledCopyO = std::conditional_t< + !Split, + typename Kernel_traits::GmemTiledCopyO, + typename Kernel_traits::GmemTiledCopyOaccum + >; + using ElementO = std::conditional_t; + + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; + + const BlockInfo binfo(params, bidb); + typename Kernel_traits::TiledMmaO tiled_mma_o; + auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); + Tensor lse = softmax.template normalize_softmax_lse(acc_o, params.scale_softmax); + Tensor acc_o_view = make_tensor(acc_o.data(), make_layout(Shape<_4, Shape<_4, _2>>{}, + Stride<_1, Shape<_4, _16>>{})); + Tensor acc_o_copy = make_fragment_like(acc_o_view); + #pragma unroll + for (int k = 0; k < size<1, 1>(acc_o_view); k++) { + #pragma unroll + for (int idx = 0; idx < 16; idx++) { + int row = idx / 4; + int col = idx % 4; + acc_o_copy(row, make_coord(col, k)) = acc_o_view(col, make_coord(row, k)); + } + } + if constexpr (!Split) { + Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) + // Partition sO to match the accumulator partitioning + using SmemTiledCopyO = typename Kernel_traits::SmemCopyAtomO; + CONVERT_TENSOR_TYPE(ElementAccum, ElementO, acc_o_copy, rO) + int warp_offset = warp_idx * 16 * 64; + int thread_offset = lane_idx % 16 * 64 + lane_idx / 16 * 16; + ElementO *Osmem_ptr_sts = reinterpret_cast(smem_) + warp_offset + thread_offset; + Tensor tOsO = make_tensor(make_smem_ptr(Osmem_ptr_sts), make_layout(Shape<_16, _2>{}, + Stride<_1, Int<16*64*kNWarps>>{})); + Tensor tOrO = make_tensor(rO.data(), make_layout(Shape<_16, _2>{}, + Stride<_1, _16>{})); + + + if constexpr (Kernel_traits::Share_Q_K_smem) { flash::sync_threads(); } + + cute::copy(tOrO, tOsO); + + const index_t row_offset_o = binfo.q_offset(params.o_batch_stride, params.o_row_stride, bidb) + + m_block * kBlockM * params.o_row_stride + bidh * params.o_head_stride; + const index_t row_offset_lseaccum = ((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; + + Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.o_ptr) + (row_offset_o)), + Shape, Int>{}, + make_stride(params.o_row_stride, _1{})); + Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lse_ptr) + row_offset_lseaccum), + Shape>{}, Stride<_1>{}); + + GmemTiledCopyO gmem_tiled_copy_Oaccum; + auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); + Tensor tOsOaccum = gmem_thr_copy_Oaccum.partition_S(sOaccum); // ((Atom,AtomNum),ATOM_M,ATOM_N) + Tensor tOgOaccum = gmem_thr_copy_Oaccum.partition_D(gOaccum); + + flash::sync_threads(); + + Tensor tOrOaccum = make_tensor(shape(tOgOaccum)); + cute::copy(gmem_tiled_copy_Oaccum, tOsOaccum, tOrOaccum); + + Tensor caccO = make_identity_tensor(Shape, Int>{}); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor taccOcO = thr_mma_o.partition_C(caccO); // (MMA,MMA_M,MMA_K) + static_assert(decltype(size<0>(taccOcO))::value == 4); + // Convert to ((2, 2), MMA_M, MMA_K) then take only the row indices. + Tensor taccOcO_row = logical_divide(taccOcO, Shape<_4>{})(make_coord(0, _), _, 0); + CUTE_STATIC_ASSERT_V(size(lse) == size(taccOcO_row)); // MMA_M + if (get<1>(taccOcO_row(0)) == 0) { + #pragma unroll + for (int mi = 0; mi < size(lse); ++mi) { + const int row = get<0>(taccOcO_row(mi)); + if (row < binfo.actual_seqlen_q - m_block * kBlockM) { gLSEaccum(row) = lse(mi); } + } + } + + // Construct identity layout for sO + Tensor cO = make_identity_tensor(make_shape(size<0>(sOaccum), size<1>(sOaccum))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + // Repeat the partitioning with identity layouts + Tensor tOcO = gmem_thr_copy_Oaccum.partition_D(cO); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + // Clear_OOB_K must be false since we don't want to write zeros to gmem + flash::copy_reg_to_global( + tOrOaccum, tOgOaccum, tOcO, params.d_v, binfo.actual_seqlen_q - m_block * kBlockM + ); + } else { + const int split_offset = params.num_splits_ptr[bidb]; + Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) + const index_t row_offset_oaccum = (((split_offset + n_split_idx) * params.h + bidh) * params.seqlen_q + + m_block * kBlockM) * params.d_v; + const index_t row_offset_lseaccum = ((split_offset + n_split_idx) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; + + Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.oaccum_ptr) + row_offset_oaccum), + Shape, Int>{}, + make_stride(kHeadDimV, _1{})); + Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lseaccum_ptr) + row_offset_lseaccum), + Shape>{}, Stride<_1>{}); + Tensor taccOrOaccum = make_tensor(acc_o_copy.data(), acc_o_copy.layout()); + + int warp_offset = warp_idx * 16 * 64; + int thread_offset = lane_idx % 16 * 64 + lane_idx / 16 * 16; + ElementO *accOsmem_ptr_sts = reinterpret_cast(smem_) + warp_offset + thread_offset; + Tensor taccOsOaccum = make_tensor(make_smem_ptr(accOsmem_ptr_sts), make_layout(Shape<_4, _4>{}, + Stride<_1, _4>{})); + + if constexpr (Kernel_traits::Share_Q_K_smem) { flash::sync_threads(); } + int O_swizzle_row_sts = tidx % 4; + #pragma unroll + for (int i = 0; i < 4; i++) { + cute::copy(taccOrOaccum(_, make_coord(i, 0)), taccOsOaccum(_, O_swizzle_row_sts ^ i)); + } + GmemTiledCopyO gmem_tiled_copy_Oaccum; + auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); + // Tensor tOsOaccum = gmem_thr_copy_Oaccum.partition_S(sOaccum); // ((Atom,AtomNum),ATOM_M,ATOM_N) + int O_swizzle_row_lds = tidx / 16 % 4; + int O_swizzle_col_lds = tidx % 16 % 4; + int O_swizzle_col_lds_new = O_swizzle_col_lds ^ O_swizzle_row_lds; + ElementO *accOsmem_ptr_lds = reinterpret_cast(smem_) + (tidx + O_swizzle_col_lds_new - O_swizzle_col_lds) * 4; + + Tensor tOsOaccum = make_tensor(make_smem_ptr(accOsmem_ptr_lds), make_layout(Shape<_4, _1, Int>{}, + Stride<_1, _0, Int<16*64>>{})); + + Tensor tOgOaccum = gmem_thr_copy_Oaccum.partition_D(gOaccum); + + flash::sync_threads(); + + Tensor tOrOaccum = make_tensor(shape(tOgOaccum)); + cute::copy(gmem_tiled_copy_Oaccum, tOsOaccum, tOrOaccum); + flash::sync_threads(); + Tensor caccO = make_identity_tensor(Shape, Int>{}); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor taccOcO = thr_mma_o.partition_C(caccO); // (MMA,MMA_M,MMA_K) + static_assert(decltype(size<0>(taccOcO))::value == 4); + // Convert to ((2, 2), MMA_M, MMA_K) then take only the row indices. + Tensor taccOcO_row = logical_divide(taccOcO, Shape<_4>{})(make_coord(0, _), _, 0); + CUTE_STATIC_ASSERT_V(size(lse) == size(taccOcO_row)); // MMA_M + if (get<1>(taccOcO_row(0)) == 0) { + #pragma unroll + for (int mi = 0; mi < size(lse); ++mi) { + const int row = get<0>(taccOcO_row(mi)); + if (row < binfo.actual_seqlen_q - m_block * kBlockM) { gLSEaccum(row) = lse(mi); } + } + } + + Tensor cO = make_identity_tensor(make_shape(size<0>(sOaccum), size<1>(sOaccum))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + // Repeat the partitioning with identity layouts + Tensor tOcaccO = gmem_thr_copy_Oaccum.partition_D(cO); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + // Clear_OOB_K must be false since we don't want to write zeros to gmem + flash::copy_reg_to_global( + tOrOaccum, tOgOaccum, tOcaccO, params.d_v, binfo.actual_seqlen_q - m_block * kBlockM + ); + #pragma unroll + for (int i = 0; i < 4; i++) { + cute::copy(taccOrOaccum(_, make_coord(i, 1)), taccOsOaccum(_, O_swizzle_row_sts ^ i)); + } + flash::sync_threads(); + cute::copy(gmem_tiled_copy_Oaccum, tOsOaccum, tOrOaccum); + tOgOaccum.data() = tOgOaccum.data() + (kHeadDimV/2); + + + flash::copy_reg_to_global( + tOrOaccum, tOgOaccum, tOcaccO, params.d_v, binfo.actual_seqlen_q - m_block * kBlockM + ); + } +} +template +__forceinline__ __device__ void compute_attn_1rowblock_splitkv_mla_k64_16x16_4waves_xcore1000( + const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int n_block_min, const int n_block_max, const bool NoSplit) { + + using Element = typename Kernel_traits::Element; + using ElementAccum = typename Kernel_traits::ElementAccum; + using index_t = typename Kernel_traits::index_t; + + // Shared memory. + extern __shared__ char smem_[]; + + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kHeadDim = Kernel_traits::kHeadDim; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kBlockKSmem = Kernel_traits::kBlockKSmem; + constexpr int kAtomLayoutMS = Kernel_traits::kAtomLayoutMS; + constexpr int kAtomLayoutMO = Kernel_traits::kAtomLayoutMO; + constexpr int Num_Stages = Kernel_traits::Num_Stages; + constexpr int kSmemSize = Kernel_traits::kSmemSize; + static_assert(kBlockKSmem == 64); + + + const BlockInfo binfo(params, bidb); + if (m_block * kBlockM >= binfo.actual_seqlen_q) return; + const int actual_seqlen_k = params.cu_seqlens_k[bidb]; + + // We iterate over the blocks in reverse order. This is because the last block is the only one + // that needs masking when we read K and V from global memory. Moreover, iterating in reverse + // might save us 1 register (we just need n_block instead of both n_block and n_block_max). + const index_t row_offset_q = binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb) + + m_block * kBlockM * params.q_row_stride + bidh * params.q_head_stride; + // We move K and V to the last block. + const int bidb_cache = params.cache_batch_idx == nullptr ? bidb : params.cache_batch_idx[bidb]; + const int *block_table = params.block_table == nullptr ? nullptr : params.block_table + bidb * params.block_table_batch_stride; + const int block_table_idx = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN / params.page_block_size; + const int block_table_offset = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN - block_table_idx * params.page_block_size; + const index_t row_offset_k = block_table == nullptr + ? binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb_cache) + + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride + : (bidh / params.h_h_k_ratio) * params.k_head_stride; + const index_t row_offset_v = block_table == nullptr + ? binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb_cache) + + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride + : (bidh / params.h_h_k_ratio) * params.v_head_stride; + + Tensor gQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), + Shape, Int>{}, + make_stride(params.q_row_stride, _1{})); + Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + row_offset_k), + Shape, Int>{}, + make_stride(params.k_row_stride, _1{})); + Tensor gV = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_ptr) + row_offset_v), + Shape, Int>{}, + make_stride(params.v_row_stride, _1{})); + + Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutQ424{}); + Tensor sK = make_tensor(sQ.data() + (Kernel_traits::Share_Q_K_smem ? 0 : size(sQ)), + typename Kernel_traits::SmemLayoutK424{}); + Tensor sV = make_tensor(sK.data(), typename Kernel_traits::SmemLayoutVtNoSwizzle{}); + Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed424{}); + Tensor sVtNoSwizzle = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{}); + + typename Kernel_traits::GmemTiledCopyB64 gmem_tiled_copy_Q; + auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx); + + Tensor tQgQ = gmem_thr_copy_Q.partition_S(gQ); + Tensor tQsQ = gmem_thr_copy_Q.partition_D(sQ); + + typename Kernel_traits::GmemTiledCopyB64 gmem_tiled_copy_KV; + auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx); + Tensor tKgK = gmem_thr_copy_KV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K) + Tensor tKsK = gmem_thr_copy_KV.partition_D(sK); + // S is only 16x16 size, so all 4 waves compute the same S + int tidx_mma_s = tidx & 0x3F; + typename Kernel_traits::TiledMmaS tiled_mma_s; + auto thr_mma_s = tiled_mma_s.get_thread_slice(tidx_mma_s); + Tensor tSrQ = thr_mma_s.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K) + Tensor tSrK = thr_mma_s.partition_fragment_B(sK(_, _, 0)); // (MMA,MMA_N,MMA_K) + typename Kernel_traits::TiledMmaO tiled_mma_o; + auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); + // Tensor tOrVt = thr_mma_o.partition_fragment_B(sVt); // (MMA, MMA_K,MMA_N) + Tensor tOrVt = make_tensor(Shape<_4, Shape<_4, _2>, _1>{}); + + Tensor acc_o = partition_fragment_C(tiled_mma_o, Shape, Int>{}); // MMA, MMA_M, MMA_K + + // + // Copy Atom retiling + // + + auto smem_tiled_copy_Q = make_tiled_copy_A(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx_mma_s); + Tensor tSsQ = smem_thr_copy_Q.partition_S(sQ); + + auto smem_tiled_copy_K = make_tiled_copy_B(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx_mma_s); + Tensor tSsK = smem_thr_copy_K.partition_S(sK); + auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed{}, tiled_mma_o); + auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx); + int warp_offset = warp_idx / kAtomLayoutMO * 16 * 64; + int thread_offset = lane_idx / 16 * 4 * 64; + Element *Vtsmem_ptr_lds = reinterpret_cast(sVt.data().get()) + warp_offset + thread_offset; + Tensor tOsVt = make_tensor(make_smem_ptr(Vtsmem_ptr_lds), make_layout(Shape<_4, _2, Int>{}, // MMA MMA_N NUM_STAGES + Stride<_1, Int<16*256>, Int>{})); + + // PREDICATES + + // Construct identity layout for sQ and sK + Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + + // Repeat the partitioning with identity layouts + Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tKVcKV = gmem_thr_copy_KV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + + + // Prologue + + + // Read Q from gmem to smem, optionally apply rotary embedding. + Tensor tQrQ = make_fragment_like(tQgQ); + // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs + flash::copy_b64(tQgQ, tQrQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); + cute::copy(tQrQ, tQsQ); + + if constexpr (Kernel_traits::Is_Q_in_regs) { + flash::sync_threads(); + cute::copy(smem_tiled_copy_Q, tSsQ, tSrQ); + flash::sync_threads(); + } + + + int n_block = n_block_max - 1; + // We don't need to clear the sK smem tiles since we'll mask out the scores anyway. + Tensor tKrK = make_fragment_like(tKgK); + int row_offset = tidx / 16 + n_block * kBlockN; + int virtual_page_idx = row_offset / params.page_block_size; + int page_offset = row_offset - virtual_page_idx * params.page_block_size; + int32_t page_idx = block_table[virtual_page_idx]; + flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, page_idx, page_offset, actual_seqlen_k - n_block * kBlockN); + + clear(acc_o); + + flash::Softmax(acc_o)> softmax; + + flash::Mask mask(actual_seqlen_k, binfo.actual_seqlen_q, params.ngroups, binfo.tot_seqlen_k, params.cp_world_size, params.cp_rank); + + // For performance reason, we separate out two kinds of iterations: + // those that need masking on S, and those that don't. + // We need masking on S for the very last block when K and V has length not multiple of kBlockN. + // We also need masking on S if it's causal, for the last ceil_div(kBlockM, kBlockN) blocks. + // We will have at least 1 "masking" iteration. + + // If not even_N, then seqlen_k might end in the middle of a block. In that case we need to + // mask 2 blocks (e.g. when kBlockM == kBlockN), not just 1. + constexpr int n_masking_steps = (!Is_causal) + ? 1 + : ((Is_even_MN && Is_causal) ? cute::ceil_div(kBlockM, kBlockN) : cute::ceil_div(kBlockM, kBlockN) + 1); + #pragma unroll + for (int masking_step = 0; masking_step < n_masking_steps; ++masking_step, --n_block) { + if (n_block > n_block_min) { + // prefetch load page index + row_offset -= kBlockN; + virtual_page_idx = row_offset / params.page_block_size; + page_offset = row_offset - virtual_page_idx * params.page_block_size; + page_idx = block_table[virtual_page_idx]; + } + Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) + flash::sync_threads(); + cute::copy(tKrK, tKsK(_, _, _, 0)); + clear(acc_s); + + flash::sync_threads(); + if (n_block > n_block_min) { + flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, page_idx, page_offset); + } + flash::gemm( + acc_s, tSrQ, tSrK, tSsQ, tSsK(_, _, _, 0), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, + smem_thr_copy_Q, smem_thr_copy_K + ); + + mask.template apply_mask( + acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 + ); + + // We have key_padding_mask so we'll need to Check_inf + masking_step == 0 + ? softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2) + : softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); + // Convert acc_s from fp32 to fp16/bf16 + //Tensor rP = flash::convert_type(acc_s); + CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) + // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. + //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); + + // flash::sync_threads(); + lds4x4_with_swizzle424(tOsVt(_, _, 0), tOrVt); + CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 + Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); + permute_4x4_b16(tOrVt_permute_view); + Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); + flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); + + // This check is at the end of the loop since we always have at least 1 iteration + if (n_masking_steps > 1 && n_block <= n_block_min) { + --n_block; + break; + } + } + + // These are the iterations where we don't need masking on S + for (; n_block >= n_block_min; --n_block) { + if (n_block > n_block_min) { + // prefetch load page index + row_offset -= kBlockN; + virtual_page_idx = row_offset / params.page_block_size; + page_offset = row_offset - virtual_page_idx * params.page_block_size; + page_idx = block_table[virtual_page_idx]; + } + Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) + flash::sync_threads(); + cute::copy(tKrK, tKsK(_, _, _, 0)); + clear(acc_s); + flash::sync_threads(); + if (n_block > n_block_min) { + // Advance gK + flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, page_idx, page_offset); + } + + flash::gemm( + acc_s, tSrQ, tSrK, tSsQ, tSsK(_, _, _, 0), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, + smem_thr_copy_Q, smem_thr_copy_K + ); + + mask.template apply_mask( + acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 + ); + softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); + + //Tensor rP = flash::convert_type(acc_s); + CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) + // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. + //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); + // flash::sync_threads(); + lds4x4_with_swizzle424(tOsVt(_, _, 0), tOrVt); + CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 + Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); + permute_4x4_b16(tOrVt_permute_view); + Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); + + flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); + } + // Epilogue + if (NoSplit) { + store_16x16(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax); + }else{ + store_16x16(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax); + } +} + +} // namespace flash diff --git a/csrc/flash_kernel/flash_fwd_split_kernel_k64_16x16_4waves.h b/csrc/flash_kernel/xcore1000/flash_fwd_mla_kernel_k64_16x16_4waves_xcore1000.h similarity index 76% rename from csrc/flash_kernel/flash_fwd_split_kernel_k64_16x16_4waves.h rename to csrc/flash_kernel/xcore1000/flash_fwd_mla_kernel_k64_16x16_4waves_xcore1000.h index 6d208b4..e06803d 100644 --- a/csrc/flash_kernel/flash_fwd_split_kernel_k64_16x16_4waves.h +++ b/csrc/flash_kernel/xcore1000/flash_fwd_mla_kernel_k64_16x16_4waves_xcore1000.h @@ -13,40 +13,16 @@ #include "utils.h" #include "softmax.h" #include "mask.h" -#include "rotary.h" -#include "attn_mask.h" namespace flash { using namespace cute; - -template -__forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_16x16_4waves(const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int num_n_splits) { - +template +__forceinline__ __device__ void store_16x16(const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, char* smem_, AccO acc_o, Softmax softmax ) { + using ElementAccum = typename Kernel_traits::ElementAccum; using Element = typename Kernel_traits::Element; - using ElementAccum = typename Kernel_traits::ElementAccum; using index_t = typename Kernel_traits::index_t; - // Shared memory. - extern __shared__ char smem_[]; - - // The thread index. - const int tidx = threadIdx.x; - const int warp_idx = tidx / 64; - const int lane_idx = tidx % 64; - - constexpr int kBlockM = Kernel_traits::kBlockM; - constexpr int kBlockN = Kernel_traits::kBlockN; - constexpr int kHeadDim = Kernel_traits::kHeadDim; - constexpr int kHeadDimV = Kernel_traits::kHeadDimV; - constexpr int kNWarps = Kernel_traits::kNWarps; - constexpr int kBlockKSmem = Kernel_traits::kBlockKSmem; - constexpr int kAtomLayoutMS = Kernel_traits::kAtomLayoutMS; - constexpr int kAtomLayoutMO = Kernel_traits::kAtomLayoutMO; - constexpr int Num_Stages = Kernel_traits::Num_Stages; - - static_assert(kBlockKSmem == 64); - using GmemTiledCopyO = std::conditional_t< !Split, typename Kernel_traits::GmemTiledCopyO, @@ -54,300 +30,19 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_16x16_4wa >; using ElementO = std::conditional_t; + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; + const BlockInfo binfo(params, bidb); - // if (threadIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { printf("Is_even_MN = %d, is_cumulativ = %d, seqlen_k_cache = %d, actual_seqlen_k = %d\n", Is_even_MN, params.is_seqlens_k_cumulative, binfo.seqlen_k_cache, binfo.actual_seqlen_k); } - // if (threadIdx.x == 0 && blockIdx.y == 1 && blockIdx.z == 0) { printf("params.knew_ptr = %p, seqlen_k_cache + seqlen_knew = %d\n", params.knew_ptr, binfo.seqlen_k_cache + (params.knew_ptr == nullptr ? 0 : params.seqlen_knew)); } - if (m_block * kBlockM >= binfo.actual_seqlen_q) return; - - const int n_blocks_per_split = ((binfo.actual_seqlen_k + kBlockN - 1) / kBlockN + num_n_splits - 1) / num_n_splits; - const int n_block_min = !Is_local - ? n_split_idx * n_blocks_per_split - : std::max(n_split_idx * n_blocks_per_split, (m_block * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q - params.window_size_left) / kBlockN); - int n_block_max = std::min(cute::ceil_div(binfo.actual_seqlen_k, kBlockN), (n_split_idx + 1) * n_blocks_per_split); - if (Is_causal || Is_local) { - n_block_max = std::min(n_block_max, - cute::ceil_div((m_block + 1) * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q / params.ngroups + params.window_size_right, kBlockN)); - } - - if (n_block_min >= n_block_max) { // This also covers the case where n_block_max <= 0 - // We exit early and write 0 to gOaccum and -inf to gLSEaccum. - // Otherwise we might read OOB elements from gK and gV, - // or get wrong results when we combine gOaccum from different blocks. - const index_t row_offset_o = binfo.q_offset(params.o_batch_stride, params.o_row_stride, bidb) - + m_block * kBlockM * params.o_row_stride + bidh * params.o_head_stride; - const index_t row_offset_oaccum = (((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q - + m_block * kBlockM) * params.d_v; - const index_t row_offset_lseaccum = ((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; - Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.oaccum_ptr : params.o_ptr) + (Split ? row_offset_oaccum : row_offset_o)), - Shape, Int>{}, - make_stride(Split ? kHeadDimV : params.o_row_stride, _1{})); - Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.softmax_lseaccum_ptr : params.softmax_lse_ptr) + row_offset_lseaccum), - Shape>{}, Stride<_1>{}); - - GmemTiledCopyO gmem_tiled_copy_Oaccum; - auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); - Tensor tOgOaccum = gmem_thr_copy_Oaccum.partition_D(gOaccum); - Tensor tOrOaccum = make_tensor(shape(tOgOaccum)); - clear(tOrOaccum); - // Construct identity layout for sO - Tensor cO = make_identity_tensor(make_shape(size<0>(gOaccum), size<1>(gOaccum))); // (BLK_M,BLK_K) -> (blk_m,blk_k) - // Repeat the partitioning with identity layouts - Tensor tOcO = gmem_thr_copy_Oaccum.partition_D(cO); - Tensor tOpO = make_tensor(make_shape(size<2>(tOgOaccum))); - if (!Is_even_K) { - #pragma unroll - for (int k = 0; k < size(tOpO); ++k) { tOpO(k) = get<1>(tOcO(0, 0, k)) < params.d_v; } - } - // Clear_OOB_K must be false since we don't want to write zeros to gmem - flash::copy( - gmem_tiled_copy_Oaccum, tOrOaccum, tOgOaccum, tOcO, tOpO, binfo.actual_seqlen_q - m_block * kBlockM - ); - #pragma unroll - for (int m = 0; m < size<1>(tOgOaccum); ++m) { - const int row = get<0>(tOcO(0, m, 0)); - if (row < binfo.actual_seqlen_q - m_block * kBlockM && get<1>(tOcO(0, m, 0)) == 0) { gLSEaccum(row) = Split ? -INFINITY : INFINITY; } - } - return; - } - - // We iterate over the blocks in reverse order. This is because the last block is the only one - // that needs masking when we read K and V from global memory. Moreover, iterating in reverse - // might save us 1 register (we just need n_block instead of both n_block and n_block_max). - const index_t row_offset_q = binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb) - + m_block * kBlockM * params.q_row_stride + bidh * params.q_head_stride; - // We move K and V to the last block. - const int bidb_cache = params.cache_batch_idx == nullptr ? bidb : params.cache_batch_idx[bidb]; - const int *block_table = params.block_table == nullptr ? nullptr : params.block_table + bidb * params.block_table_batch_stride; - const int block_table_idx = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN / params.page_block_size; - const int block_table_offset = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN - block_table_idx * params.page_block_size; - const index_t row_offset_k = block_table == nullptr - ? binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb_cache) - + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride - : (bidh / params.h_h_k_ratio) * params.k_head_stride; - const index_t row_offset_v = block_table == nullptr - ? binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb_cache) - + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride - : (bidh / params.h_h_k_ratio) * params.v_head_stride; - - Tensor gQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), - Shape, Int>{}, - make_stride(params.q_row_stride, _1{})); - Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + row_offset_k), - Shape, Int>{}, - make_stride(params.k_row_stride, _1{})); - // if (threadIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { printf("k_ptr = %p, row_offset_k = %d, gK_ptr = %p\n", params.k_ptr, row_offset_k, gK.data()); } - Tensor gV = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_ptr) + row_offset_v), - Shape, Int>{}, - make_stride(params.v_row_stride, _1{})); - - Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), - typename Kernel_traits::SmemLayoutQ424{}); - Tensor sK = make_tensor(sQ.data() + (Kernel_traits::Share_Q_K_smem ? 0 : size(sQ)), - typename Kernel_traits::SmemLayoutK{}); - Tensor sV = make_tensor(sK.data(), typename Kernel_traits::SmemLayoutVtNoSwizzle{}); - Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed{}); - Tensor sVtNoSwizzle = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{}); - - typename Kernel_traits::GmemTiledCopyB64 gmem_tiled_copy_Q; - auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx); - - Tensor tQgQ = gmem_thr_copy_Q.partition_S(gQ); - Tensor tQsQ = gmem_thr_copy_Q.partition_D(sQ); - - typename Kernel_traits::GmemTiledCopyB64 gmem_tiled_copy_KV; - auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx); - Tensor tKgK = gmem_thr_copy_KV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K) - Tensor tKsK = gmem_thr_copy_KV.partition_D(sK); - // S is only 16x16 size, so all 4 waves compute the same S - int tidx_mma_s = tidx & 0x3F; - typename Kernel_traits::TiledMmaS tiled_mma_s; - auto thr_mma_s = tiled_mma_s.get_thread_slice(tidx_mma_s); - Tensor tSrQ = thr_mma_s.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K) - Tensor tSrK = thr_mma_s.partition_fragment_B(sK(_, _, 0)); // (MMA,MMA_N,MMA_K) typename Kernel_traits::TiledMmaO tiled_mma_o; auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); - // Tensor tOrVt = thr_mma_o.partition_fragment_B(sVt); // (MMA, MMA_K,MMA_N) - Tensor tOrVt = make_tensor(Shape<_4, Shape<_4, _2>, _1>{}); - - Tensor acc_o = partition_fragment_C(tiled_mma_o, Shape, Int>{}); // MMA, MMA_M, MMA_K - - // - // Copy Atom retiling - // - - auto smem_tiled_copy_Q = make_tiled_copy_A(typename Kernel_traits::SmemCopyAtomB64{}, tiled_mma_s); - auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx_mma_s); - Tensor tSsQ = smem_thr_copy_Q.partition_S(sQ); - - auto smem_tiled_copy_K = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomB64{}, tiled_mma_s); - auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx_mma_s); - Tensor tSsK = smem_thr_copy_K.partition_S(sK); - auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed{}, tiled_mma_o); - auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx); - int warp_offset = warp_idx / kAtomLayoutMO * 16 * 64; - int thread_offset = lane_idx / 16 * 4 * 64; - Element *Vtsmem_ptr_lds = reinterpret_cast(sVt.data().get()) + warp_offset + thread_offset; - Tensor tOsVt = make_tensor(make_smem_ptr(Vtsmem_ptr_lds), make_layout(Shape<_4, _2, Int>{}, // MMA MMA_N NUM_STAGES - Stride<_1, Int<16*256>, Int>{})); - - // PREDICATES - - // Construct identity layout for sQ and sK - Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) - Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) - - // Repeat the partitioning with identity layouts - Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) - Tensor tKVcKV = gmem_thr_copy_KV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) - - - // Prologue - - - // Read Q from gmem to smem, optionally apply rotary embedding. - Tensor tQrQ = make_fragment_like(tQgQ); - // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs - flash::copy_b64(tQgQ, tQrQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); - cute::copy(tQrQ, tQsQ); - - if constexpr (Kernel_traits::Is_Q_in_regs) { - flash::sync_threads(); - cute::copy(smem_tiled_copy_Q, tSsQ, tSrQ); - flash::sync_threads(); - } - - - int n_block = n_block_max - 1; - int Ksmem_read_index = 0; - int Ksmem_write_index = 0; - // We don't need to clear the sK smem tiles since we'll mask out the scores anyway. - Tensor tKrK = make_fragment_like(tKgK); - flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block, - block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, binfo.actual_seqlen_k - n_block * kBlockN); - - - // flash::cp_async_wait<0>(); - // __syncthreads(); - // if (tidx == 0 && blockIdx.y == 0 && blockIdx.z == 0) { print(tKsK); } - // __syncthreads(); - - clear(acc_o); - - flash::Softmax(acc_o)> softmax; - - const float alibi_slope = !Has_alibi ? 0.0f : reinterpret_cast(params.alibi_slopes_ptr)[bidb * params.alibi_slopes_batch_stride + bidh] / params.scale_softmax; - flash::Mask mask(binfo.actual_seqlen_k, binfo.actual_seqlen_q, params.ngroups, params.window_size_left, params.window_size_right, alibi_slope); - - // For performance reason, we separate out two kinds of iterations: - // those that need masking on S, and those that don't. - // We need masking on S for the very last block when K and V has length not multiple of kBlockN. - // We also need masking on S if it's causal, for the last ceil_div(kBlockM, kBlockN) blocks. - // We will have at least 1 "masking" iteration. - - // If not even_N, then seqlen_k might end in the middle of a block. In that case we need to - // mask 2 blocks (e.g. when kBlockM == kBlockN), not just 1. - constexpr int n_masking_steps = (!Is_causal && !Is_local) - ? 1 - : ((Is_even_MN && Is_causal) ? cute::ceil_div(kBlockM, kBlockN) : cute::ceil_div(kBlockM, kBlockN) + 1); - #pragma unroll - for (int masking_step = 0; masking_step < n_masking_steps; ++masking_step, --n_block) { - Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) - cute::copy(tKrK, tKsK(_, _, _, Ksmem_write_index)); - Ksmem_write_index ^= 1; - clear(acc_s); - - flash::sync_threads(); - if (n_block > n_block_min) { - flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, - block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size); - } - - flash::gemm_opt( - acc_s, tSrQ, tSrK, tSsQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, - smem_thr_copy_Q, smem_thr_copy_K - ); - // if (cute::thread0()) { print(acc_s); } - if constexpr (Is_softcap){ - flash::apply_softcap(acc_s, params.softcap); - } - - mask.template apply_mask( - acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 - ); - - // We have key_padding_mask so we'll need to Check_inf - masking_step == 0 - ? softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2) - : softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); - // if (cute::thread0()) { print(scores_max); print(scores_sum); print(scores); } - - // Convert acc_s from fp32 to fp16/bf16 - //Tensor rP = flash::convert_type(acc_s); - CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) - // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) - // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. - //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); - lds4x4_with_swizzle424(tOsVt(_, _, Ksmem_read_index), tOrVt); - CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 - Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); - permute_4x4_b16(tOrVt_permute_view); - Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); - flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); - Ksmem_read_index ^= 1; - - // This check is at the end of the loop since we always have at least 1 iteration - if (n_masking_steps > 1 && n_block <= n_block_min) { - --n_block; - break; - } - } - - // These are the iterations where we don't need masking on S - for (; n_block >= n_block_min; --n_block) { - Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) - cute::copy(tKrK, tKsK(_, _, _, Ksmem_write_index)); - Ksmem_write_index ^= 1; - clear(acc_s); - flash::sync_threads(); - if (n_block > n_block_min) { - // Advance gK - flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, - block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size); - } - - flash::gemm_opt( - acc_s, tSrQ, tSrK, tSsQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, - smem_thr_copy_Q, smem_thr_copy_K - ); - - - if constexpr (Is_softcap){ - flash::apply_softcap(acc_s, params.softcap); - } - - mask.template apply_mask( - acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 - ); - softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); - - //Tensor rP = flash::convert_type(acc_s); - CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) - // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) - // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. - //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); - lds4x4_with_swizzle424(tOsVt(_, _, Ksmem_read_index), tOrVt); - CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 - Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); - permute_4x4_b16(tOrVt_permute_view); - Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); - - flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); - Ksmem_read_index ^= 1; - } - - // Epilogue - - Tensor lse = softmax.template normalize_softmax_lse(acc_o, params.scale_softmax); Tensor acc_o_view = make_tensor(acc_o.data(), make_layout(Shape<_4, Shape<_4, _2>>{}, Stride<_1, Shape<_4, _16>>{})); @@ -361,7 +56,6 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_16x16_4wa acc_o_copy(row, make_coord(col, k)) = acc_o_view(col, make_coord(row, k)); } } - // if (cute::thread0()) { print(lse); } if constexpr (!Split) { Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) // Partition sO to match the accumulator partitioning @@ -376,8 +70,6 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_16x16_4wa Stride<_1, _16>{})); - // sOaccum is larger than sQ, so we need to syncthreads here - // TODO: allocate enough smem for sOaccum if constexpr (Kernel_traits::Share_Q_K_smem) { flash::sync_threads(); } cute::copy(tOrO, tOsO); @@ -391,7 +83,6 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_16x16_4wa make_stride(params.o_row_stride, _1{})); Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lse_ptr) + row_offset_lseaccum), Shape>{}, Stride<_1>{}); - // if (tidx == 0) { printf("row_offset_o = %d, bidh = %d, gOaccum = %p\n", row_offset_o, bidh, gOaccum.data()); } GmemTiledCopyO gmem_tiled_copy_Oaccum; auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); @@ -426,17 +117,17 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_16x16_4wa tOrOaccum, tOgOaccum, tOcO, params.d_v, binfo.actual_seqlen_q - m_block * kBlockM ); } else { + const int split_offset = params.num_splits_ptr[bidb]; Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) - const index_t row_offset_oaccum = (((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + const index_t row_offset_oaccum = (((split_offset + n_split_idx) * params.h + bidh) * params.seqlen_q + m_block * kBlockM) * params.d_v; - const index_t row_offset_lseaccum = ((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; + const index_t row_offset_lseaccum = ((split_offset + n_split_idx) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.oaccum_ptr) + row_offset_oaccum), Shape, Int>{}, make_stride(kHeadDimV, _1{})); Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.softmax_lseaccum_ptr) + row_offset_lseaccum), Shape>{}, Stride<_1>{}); - // if (tidx == 0) { printf("row_offset_o = %d, bidh = %d, gOaccum = %p\n", row_offset_o, bidh, gOaccum.data()); } Tensor taccOrOaccum = make_tensor(acc_o_copy.data(), acc_o_copy.layout()); int warp_offset = warp_idx * 16 * 64; @@ -491,5 +182,262 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_16x16_4wa ); } } +template +__forceinline__ __device__ void compute_attn_1rowblock_splitkv_mla_k64_16x16_4waves_xcore1000( + const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int n_block_min, const int n_block_max, const bool NoSplit) { + + using Element = typename Kernel_traits::Element; + using ElementAccum = typename Kernel_traits::ElementAccum; + using index_t = typename Kernel_traits::index_t; + + // Shared memory. + extern __shared__ char smem_[]; + + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kHeadDim = Kernel_traits::kHeadDim; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kBlockKSmem = Kernel_traits::kBlockKSmem; + constexpr int kAtomLayoutMS = Kernel_traits::kAtomLayoutMS; + constexpr int kAtomLayoutMO = Kernel_traits::kAtomLayoutMO; + constexpr int Num_Stages = Kernel_traits::Num_Stages; + + static_assert(kBlockKSmem == 64); + + + const BlockInfo binfo(params, bidb); + if (m_block * kBlockM >= binfo.actual_seqlen_q) return; + const int actual_seqlen_k = params.cu_seqlens_k[bidb]; + + // We iterate over the blocks in reverse order. This is because the last block is the only one + // that needs masking when we read K and V from global memory. Moreover, iterating in reverse + // might save us 1 register (we just need n_block instead of both n_block and n_block_max). + const index_t row_offset_q = binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb) + + m_block * kBlockM * params.q_row_stride + bidh * params.q_head_stride; + // We move K and V to the last block. + const int bidb_cache = params.cache_batch_idx == nullptr ? bidb : params.cache_batch_idx[bidb]; + const int *block_table = params.block_table == nullptr ? nullptr : params.block_table + bidb * params.block_table_batch_stride; + const int block_table_idx = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN / params.page_block_size; + const int block_table_offset = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN - block_table_idx * params.page_block_size; + const index_t row_offset_k = block_table == nullptr + ? binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb_cache) + + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride + : (bidh / params.h_h_k_ratio) * params.k_head_stride; + const index_t row_offset_v = block_table == nullptr + ? binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb_cache) + + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride + : (bidh / params.h_h_k_ratio) * params.v_head_stride; + + Tensor gQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), + Shape, Int>{}, + make_stride(params.q_row_stride, _1{})); + Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + row_offset_k), + Shape, Int>{}, + make_stride(params.k_row_stride, _1{})); + Tensor gV = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_ptr) + row_offset_v), + Shape, Int>{}, + make_stride(params.v_row_stride, _1{})); + + Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutQ424{}); + Tensor sK = make_tensor(sQ.data() + (Kernel_traits::Share_Q_K_smem ? 0 : size(sQ)), + typename Kernel_traits::SmemLayoutK424{}); + Tensor sV = make_tensor(sK.data(), typename Kernel_traits::SmemLayoutVtNoSwizzle{}); + Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed424{}); + Tensor sVtNoSwizzle = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{}); + + typename Kernel_traits::GmemTiledCopyB64 gmem_tiled_copy_Q; + auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx); + + Tensor tQgQ = gmem_thr_copy_Q.partition_S(gQ); + Tensor tQsQ = gmem_thr_copy_Q.partition_D(sQ); + + typename Kernel_traits::GmemTiledCopyB64 gmem_tiled_copy_KV; + auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx); + Tensor tKgK = gmem_thr_copy_KV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K) + Tensor tKsK = gmem_thr_copy_KV.partition_D(sK); + // S is only 16x16 size, so all 4 waves compute the same S + int tidx_mma_s = tidx & 0x3F; + typename Kernel_traits::TiledMmaS tiled_mma_s; + auto thr_mma_s = tiled_mma_s.get_thread_slice(tidx_mma_s); + Tensor tSrQ = thr_mma_s.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K) + Tensor tSrK = thr_mma_s.partition_fragment_B(sK(_, _, 0)); // (MMA,MMA_N,MMA_K) + typename Kernel_traits::TiledMmaO tiled_mma_o; + auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); + // Tensor tOrVt = thr_mma_o.partition_fragment_B(sVt); // (MMA, MMA_K,MMA_N) + Tensor tOrVt = make_tensor(Shape<_4, Shape<_4, _2>, _1>{}); + + Tensor acc_o = partition_fragment_C(tiled_mma_o, Shape, Int>{}); // MMA, MMA_M, MMA_K + + // + // Copy Atom retiling + // + + auto smem_tiled_copy_Q = make_tiled_copy_A(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx_mma_s); + Tensor tSsQ = smem_thr_copy_Q.partition_S(sQ); + + auto smem_tiled_copy_K = make_tiled_copy_B(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx_mma_s); + Tensor tSsK = smem_thr_copy_K.partition_S(sK); + auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed{}, tiled_mma_o); + auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx); + int warp_offset = warp_idx / kAtomLayoutMO * 16 * 64; + int thread_offset = lane_idx / 16 * 4 * 64; + Element *Vtsmem_ptr_lds = reinterpret_cast(sVt.data().get()) + warp_offset + thread_offset; + Tensor tOsVt = make_tensor(make_smem_ptr(Vtsmem_ptr_lds), make_layout(Shape<_4, _2, Int>{}, // MMA MMA_N NUM_STAGES + Stride<_1, Int<16*256>, Int>{})); + + // PREDICATES + + // Construct identity layout for sQ and sK + Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + + // Repeat the partitioning with identity layouts + Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tKVcKV = gmem_thr_copy_KV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + + + // Prologue + + + // Read Q from gmem to smem, optionally apply rotary embedding. + Tensor tQrQ = make_fragment_like(tQgQ); + // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs + flash::copy_b64(tQgQ, tQrQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); + cute::copy(tQrQ, tQsQ); + + if constexpr (Kernel_traits::Is_Q_in_regs) { + flash::sync_threads(); + cute::copy(smem_tiled_copy_Q, tSsQ, tSrQ); + flash::sync_threads(); + } + + + int n_block = n_block_max - 1; + int Ksmem_read_index = 0; + int Ksmem_write_index = 0; + // We don't need to clear the sK smem tiles since we'll mask out the scores anyway. + Tensor tKrK = make_fragment_like(tKgK); + flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, binfo.actual_seqlen_k - n_block * kBlockN); + + + clear(acc_o); + + flash::Softmax(acc_o)> softmax; + + flash::Mask mask(binfo.actual_seqlen_k, binfo.actual_seqlen_q, params.ngroups, binfo.tot_seqlen_k, params.cp_world_size, params.cp_rank); + + // For performance reason, we separate out two kinds of iterations: + // those that need masking on S, and those that don't. + // We need masking on S for the very last block when K and V has length not multiple of kBlockN. + // We also need masking on S if it's causal, for the last ceil_div(kBlockM, kBlockN) blocks. + // We will have at least 1 "masking" iteration. + + // If not even_N, then seqlen_k might end in the middle of a block. In that case we need to + // mask 2 blocks (e.g. when kBlockM == kBlockN), not just 1. + constexpr int n_masking_steps = (!Is_causal) + ? 1 + : ((Is_even_MN && Is_causal) ? cute::ceil_div(kBlockM, kBlockN) : cute::ceil_div(kBlockM, kBlockN) + 1); + #pragma unroll + for (int masking_step = 0; masking_step < n_masking_steps; ++masking_step, --n_block) { + Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) + cute::copy(tKrK, tKsK(_, _, _, Ksmem_write_index)); + Ksmem_write_index ^= 1; + clear(acc_s); + + flash::sync_threads(); + if (n_block > n_block_min) { + flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size); + } + + flash::gemm( + acc_s, tSrQ, tSrK, tSsQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, + smem_thr_copy_Q, smem_thr_copy_K + ); + + mask.template apply_mask( + acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 + ); + + // We have key_padding_mask so we'll need to Check_inf + masking_step == 0 + ? softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2) + : softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); + // Convert acc_s from fp32 to fp16/bf16 + //Tensor rP = flash::convert_type(acc_s); + CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) + // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. + //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); + lds4x4_with_swizzle424(tOsVt(_, _, Ksmem_read_index), tOrVt); + CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 + Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); + permute_4x4_b16(tOrVt_permute_view); + Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); + flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); + Ksmem_read_index ^= 1; + + // This check is at the end of the loop since we always have at least 1 iteration + if (n_masking_steps > 1 && n_block <= n_block_min) { + --n_block; + break; + } + } + + // These are the iterations where we don't need masking on S + for (; n_block >= n_block_min; --n_block) { + Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) + cute::copy(tKrK, tKsK(_, _, _, Ksmem_write_index)); + Ksmem_write_index ^= 1; + clear(acc_s); + flash::sync_threads(); + if (n_block > n_block_min) { + // Advance gK + flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size); + } + + flash::gemm( + acc_s, tSrQ, tSrK, tSsQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, + smem_thr_copy_Q, smem_thr_copy_K + ); + + mask.template apply_mask( + acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 + ); + softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); + + //Tensor rP = flash::convert_type(acc_s); + CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) + // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. + //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); + lds4x4_with_swizzle424(tOsVt(_, _, Ksmem_read_index), tOrVt); + CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 + Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); + permute_4x4_b16(tOrVt_permute_view); + Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); + + flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); + Ksmem_read_index ^= 1; + } + + // Epilogue + if (NoSplit) { + store_16x16(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax); + }else{ + store_16x16(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax); + } +} } // namespace flash diff --git a/csrc/flash_kernel/flash_fwd_split_kernel_k64_32x16_4waves.h b/csrc/flash_kernel/xcore1000/flash_fwd_mla_kernel_k64_32x16_4waves_xcore1000.h similarity index 80% rename from csrc/flash_kernel/flash_fwd_split_kernel_k64_32x16_4waves.h rename to csrc/flash_kernel/xcore1000/flash_fwd_mla_kernel_k64_32x16_4waves_xcore1000.h index b5f5e57..20e7dcb 100644 --- a/csrc/flash_kernel/flash_fwd_split_kernel_k64_32x16_4waves.h +++ b/csrc/flash_kernel/xcore1000/flash_fwd_mla_kernel_k64_32x16_4waves_xcore1000.h @@ -13,40 +13,16 @@ #include "utils.h" #include "softmax.h" #include "mask.h" -#include "rotary.h" -#include "attn_mask.h" namespace flash { using namespace cute; - -template -__forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_32x16_4waves(const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int num_n_splits) { - - using Element = typename Kernel_traits::Element; +template +__forceinline__ __device__ void store_32x16(const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, char* smem_, AccO acc_o, Softmax softmax ) { using ElementAccum = typename Kernel_traits::ElementAccum; + using Element = typename Kernel_traits::Element; using index_t = typename Kernel_traits::index_t; - // Shared memory. - extern __shared__ char smem_[]; - - // The thread index. - const int tidx = threadIdx.x; - const int warp_idx = tidx / 64; - const int lane_idx = tidx % 64; - - constexpr int kBlockM = Kernel_traits::kBlockM; - constexpr int kBlockN = Kernel_traits::kBlockN; - constexpr int kHeadDim = Kernel_traits::kHeadDim; - constexpr int kHeadDimV = Kernel_traits::kHeadDimV; - constexpr int kNWarps = Kernel_traits::kNWarps; - constexpr int kBlockKSmem = Kernel_traits::kBlockKSmem; - constexpr int kAtomLayoutMS = Kernel_traits::kAtomLayoutMS; - constexpr int kAtomLayoutMO = Kernel_traits::kAtomLayoutMO; - constexpr int Num_Stages = Kernel_traits::Num_Stages; - - static_assert(kBlockKSmem == 64); - using GmemTiledCopyO = std::conditional_t< !Split, typename Kernel_traits::GmemTiledCopyO, @@ -54,300 +30,19 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_32x16_4wa >; using ElementO = std::conditional_t; + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; + const BlockInfo binfo(params, bidb); - // if (threadIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { printf("Is_even_MN = %d, is_cumulativ = %d, seqlen_k_cache = %d, actual_seqlen_k = %d\n", Is_even_MN, params.is_seqlens_k_cumulative, binfo.seqlen_k_cache, binfo.actual_seqlen_k); } - // if (threadIdx.x == 0 && blockIdx.y == 1 && blockIdx.z == 0) { printf("params.knew_ptr = %p, seqlen_k_cache + seqlen_knew = %d\n", params.knew_ptr, binfo.seqlen_k_cache + (params.knew_ptr == nullptr ? 0 : params.seqlen_knew)); } - if (m_block * kBlockM >= binfo.actual_seqlen_q) return; - - const int n_blocks_per_split = ((binfo.actual_seqlen_k + kBlockN - 1) / kBlockN + num_n_splits - 1) / num_n_splits; - const int n_block_min = !Is_local - ? n_split_idx * n_blocks_per_split - : std::max(n_split_idx * n_blocks_per_split, (m_block * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q - params.window_size_left) / kBlockN); - int n_block_max = std::min(cute::ceil_div(binfo.actual_seqlen_k, kBlockN), (n_split_idx + 1) * n_blocks_per_split); - if (Is_causal || Is_local) { - n_block_max = std::min(n_block_max, - cute::ceil_div((m_block + 1) * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q / params.ngroups + params.window_size_right, kBlockN)); - } - - if (n_block_min >= n_block_max) { // This also covers the case where n_block_max <= 0 - // We exit early and write 0 to gOaccum and -inf to gLSEaccum. - // Otherwise we might read OOB elements from gK and gV, - // or get wrong results when we combine gOaccum from different blocks. - const index_t row_offset_o = binfo.q_offset(params.o_batch_stride, params.o_row_stride, bidb) - + m_block * kBlockM * params.o_row_stride + bidh * params.o_head_stride; - const index_t row_offset_oaccum = (((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q - + m_block * kBlockM) * params.d_v; - const index_t row_offset_lseaccum = ((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; - Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.oaccum_ptr : params.o_ptr) + (Split ? row_offset_oaccum : row_offset_o)), - Shape, Int>{}, - make_stride(Split ? kHeadDimV : params.o_row_stride, _1{})); - Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.softmax_lseaccum_ptr : params.softmax_lse_ptr) + row_offset_lseaccum), - Shape>{}, Stride<_1>{}); - - GmemTiledCopyO gmem_tiled_copy_Oaccum; - auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); - Tensor tOgOaccum = gmem_thr_copy_Oaccum.partition_D(gOaccum); - Tensor tOrOaccum = make_tensor(shape(tOgOaccum)); - clear(tOrOaccum); - // Construct identity layout for sO - Tensor cO = make_identity_tensor(make_shape(size<0>(gOaccum), size<1>(gOaccum))); // (BLK_M,BLK_K) -> (blk_m,blk_k) - // Repeat the partitioning with identity layouts - Tensor tOcO = gmem_thr_copy_Oaccum.partition_D(cO); - Tensor tOpO = make_tensor(make_shape(size<2>(tOgOaccum))); - if (!Is_even_K) { - #pragma unroll - for (int k = 0; k < size(tOpO); ++k) { tOpO(k) = get<1>(tOcO(0, 0, k)) < params.d_v; } - } - // Clear_OOB_K must be false since we don't want to write zeros to gmem - flash::copy( - gmem_tiled_copy_Oaccum, tOrOaccum, tOgOaccum, tOcO, tOpO, binfo.actual_seqlen_q - m_block * kBlockM - ); - #pragma unroll - for (int m = 0; m < size<1>(tOgOaccum); ++m) { - const int row = get<0>(tOcO(0, m, 0)); - if (row < binfo.actual_seqlen_q - m_block * kBlockM && get<1>(tOcO(0, m, 0)) == 0) { gLSEaccum(row) = Split ? -INFINITY : INFINITY; } - } - return; - } - - // We iterate over the blocks in reverse order. This is because the last block is the only one - // that needs masking when we read K and V from global memory. Moreover, iterating in reverse - // might save us 1 register (we just need n_block instead of both n_block and n_block_max). - const index_t row_offset_q = binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb) - + m_block * kBlockM * params.q_row_stride + bidh * params.q_head_stride; - // We move K and V to the last block. - const int bidb_cache = params.cache_batch_idx == nullptr ? bidb : params.cache_batch_idx[bidb]; - const int *block_table = params.block_table == nullptr ? nullptr : params.block_table + bidb * params.block_table_batch_stride; - const int block_table_idx = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN / params.page_block_size; - const int block_table_offset = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN - block_table_idx * params.page_block_size; - const index_t row_offset_k = block_table == nullptr - ? binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb_cache) - + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride - : (bidh / params.h_h_k_ratio) * params.k_head_stride; - const index_t row_offset_v = block_table == nullptr - ? binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb_cache) - + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride - : (bidh / params.h_h_k_ratio) * params.v_head_stride; - - Tensor gQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), - Shape, Int>{}, - make_stride(params.q_row_stride, _1{})); - Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + row_offset_k), - Shape, Int>{}, - make_stride(params.k_row_stride, _1{})); - // if (threadIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { printf("k_ptr = %p, row_offset_k = %d, gK_ptr = %p\n", params.k_ptr, row_offset_k, gK.data()); } - Tensor gV = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_ptr) + row_offset_v), - Shape, Int>{}, - make_stride(params.v_row_stride, _1{})); - - Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), - typename Kernel_traits::SmemLayoutQ{}); - Tensor sK = make_tensor(sQ.data() + (Kernel_traits::Share_Q_K_smem ? 0 : size(sQ)), - typename Kernel_traits::SmemLayoutK{}); - Tensor sV = make_tensor(sK.data(), typename Kernel_traits::SmemLayoutVtNoSwizzle{}); - Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed{}); - Tensor sVtNoSwizzle = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{}); - - typename Kernel_traits::GmemTiledCopyB128 gmem_tiled_copy_Q; - auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx); - - Tensor tQgQ = gmem_thr_copy_Q.partition_S(gQ); - Tensor tQsQ = gmem_thr_copy_Q.partition_D(sQ); - - typename Kernel_traits::GmemTiledCopyB64 gmem_tiled_copy_KV; - auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx); - Tensor tKgK = gmem_thr_copy_KV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K) - Tensor tKsK = gmem_thr_copy_KV.partition_D(sK); - // wave0 and wave2 compute the same S, wave1 and wave3 compute the same S - int tidx_mma_s = tidx & 0x7F; - typename Kernel_traits::TiledMmaS tiled_mma_s; - auto thr_mma_s = tiled_mma_s.get_thread_slice(tidx_mma_s); - Tensor tSrQ = thr_mma_s.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K) - Tensor tSrK = thr_mma_s.partition_fragment_B(sK(_, _, 0)); // (MMA,MMA_N,MMA_K) typename Kernel_traits::TiledMmaO tiled_mma_o; auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); - // Tensor tOrVt = thr_mma_o.partition_fragment_B(sVt); // (MMA, MMA_K,MMA_N) - Tensor tOrVt = make_tensor(Shape<_4, Shape<_4, _4>, _1>{}); - - Tensor acc_o = partition_fragment_C(tiled_mma_o, Shape, Int>{}); // MMA, MMA_M, MMA_K - - // - // Copy Atom retiling - // - - auto smem_tiled_copy_Q = make_tiled_copy_A(typename Kernel_traits::SmemCopyAtomB64{}, tiled_mma_s); - auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx_mma_s); - Tensor tSsQ = smem_thr_copy_Q.partition_S(sQ); - - auto smem_tiled_copy_K = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomB64{}, tiled_mma_s); - auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx_mma_s); - Tensor tSsK = smem_thr_copy_K.partition_S(sK); - auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed{}, tiled_mma_o); - auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx); - int warp_offset = warp_idx / kAtomLayoutMO * 16 * 64; - int thread_offset = lane_idx / 16 * 4 * 64; - Element *Vtsmem_ptr_lds = reinterpret_cast(sVt.data().get()) + warp_offset + thread_offset; - Tensor tOsVt = make_tensor(make_smem_ptr(Vtsmem_ptr_lds), make_layout(Shape<_4, _4, Int>{}, // MMA MMA_N NUM_STAGES - Stride<_1, Int<16*128>, Int>{})); - - // PREDICATES - - // Construct identity layout for sQ and sK - Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) - Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) - - // Repeat the partitioning with identity layouts - Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) - Tensor tKVcKV = gmem_thr_copy_KV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) - - - // Prologue - - - // Read Q from gmem to smem, optionally apply rotary embedding. - Tensor tQrQ = make_fragment_like(tQgQ); - // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs - flash::copy_b128(tQgQ, tQrQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); - cute::copy(tQrQ, tQsQ); - - if constexpr (Kernel_traits::Is_Q_in_regs) { - flash::sync_threads(); - cute::copy(smem_tiled_copy_Q, tSsQ, tSrQ); - flash::sync_threads(); - } - - - int n_block = n_block_max - 1; - int Ksmem_read_index = 0; - int Ksmem_write_index = 0; - // We don't need to clear the sK smem tiles since we'll mask out the scores anyway. - Tensor tKrK = make_fragment_like(tKgK); - flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block, - block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, binfo.actual_seqlen_k - n_block * kBlockN); - - - // flash::cp_async_wait<0>(); - // __syncthreads(); - // if (tidx == 0 && blockIdx.y == 0 && blockIdx.z == 0) { print(tKsK); } - // __syncthreads(); - - clear(acc_o); - - flash::Softmax(acc_o)> softmax; - - const float alibi_slope = !Has_alibi ? 0.0f : reinterpret_cast(params.alibi_slopes_ptr)[bidb * params.alibi_slopes_batch_stride + bidh] / params.scale_softmax; - flash::Mask mask(binfo.actual_seqlen_k, binfo.actual_seqlen_q, params.ngroups, params.window_size_left, params.window_size_right, alibi_slope); - - // For performance reason, we separate out two kinds of iterations: - // those that need masking on S, and those that don't. - // We need masking on S for the very last block when K and V has length not multiple of kBlockN. - // We also need masking on S if it's causal, for the last ceil_div(kBlockM, kBlockN) blocks. - // We will have at least 1 "masking" iteration. - - // If not even_N, then seqlen_k might end in the middle of a block. In that case we need to - // mask 2 blocks (e.g. when kBlockM == kBlockN), not just 1. - constexpr int n_masking_steps = (!Is_causal && !Is_local) - ? 1 - : ((Is_even_MN && Is_causal) ? cute::ceil_div(kBlockM, kBlockN) : cute::ceil_div(kBlockM, kBlockN) + 1); - #pragma unroll - for (int masking_step = 0; masking_step < n_masking_steps; ++masking_step, --n_block) { - Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) - cute::copy(tKrK, tKsK(_, _, _, Ksmem_write_index)); - Ksmem_write_index ^= 1; - clear(acc_s); - - flash::sync_threads(); - if (n_block > n_block_min) { - flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, - block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size); - } - - flash::gemm_opt( - acc_s, tSrQ, tSrK, tSsQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, - smem_thr_copy_Q, smem_thr_copy_K - ); - // if (cute::thread0()) { print(acc_s); } - if constexpr (Is_softcap){ - flash::apply_softcap(acc_s, params.softcap); - } - - mask.template apply_mask( - acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 - ); - - // We have key_padding_mask so we'll need to Check_inf - masking_step == 0 - ? softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2) - : softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); - // if (cute::thread0()) { print(scores_max); print(scores_sum); print(scores); } - - // Convert acc_s from fp32 to fp16/bf16 - //Tensor rP = flash::convert_type(acc_s); - CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) - // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) - // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. - //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); - lds4x4_with_swizzle424(tOsVt(_, _, Ksmem_read_index), tOrVt); - CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 - Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); - permute_4x4_b16(tOrVt_permute_view); - Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); - flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); - Ksmem_read_index ^= 1; - - // This check is at the end of the loop since we always have at least 1 iteration - if (n_masking_steps > 1 && n_block <= n_block_min) { - --n_block; - break; - } - } - - // These are the iterations where we don't need masking on S - for (; n_block >= n_block_min; --n_block) { - Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) - cute::copy(tKrK, tKsK(_, _, _, Ksmem_write_index)); - Ksmem_write_index ^= 1; - clear(acc_s); - flash::sync_threads(); - if (n_block > n_block_min) { - // Advance gK - flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, - block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size); - } - - flash::gemm_opt( - acc_s, tSrQ, tSrK, tSsQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, - smem_thr_copy_Q, smem_thr_copy_K - ); - - - if constexpr (Is_softcap){ - flash::apply_softcap(acc_s, params.softcap); - } - - mask.template apply_mask( - acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 - ); - softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); - - //Tensor rP = flash::convert_type(acc_s); - CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) - // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) - // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. - //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); - lds4x4_with_swizzle424(tOsVt(_, _, Ksmem_read_index), tOrVt); - CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 - Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); - permute_4x4_b16(tOrVt_permute_view); - Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); - - flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); - Ksmem_read_index ^= 1; - } - - // Epilogue - - Tensor lse = softmax.template normalize_softmax_lse(acc_o, params.scale_softmax); Tensor acc_o_view = make_tensor(acc_o.data(), make_layout(Shape<_4, Shape<_4, _4>>{}, Stride<_1, Shape<_4, _16>>{})); @@ -361,7 +56,6 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_32x16_4wa acc_o_copy(row, make_coord(col, k)) = acc_o_view(col, make_coord(row, k)); } } - // if (cute::thread0()) { print(lse); } if constexpr (!Split) { Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) // Partition sO to match the accumulator partitioning @@ -376,8 +70,6 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_32x16_4wa Stride<_1, _16>{})); - // sOaccum is larger than sQ, so we need to syncthreads here - // TODO: allocate enough smem for sOaccum if constexpr (Kernel_traits::Share_Q_K_smem) { flash::sync_threads(); } cute::copy(tOrO, tOsO); @@ -426,10 +118,11 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_32x16_4wa tOrOaccum, tOgOaccum, tOcO, params.d_v, binfo.actual_seqlen_q - m_block * kBlockM ); } else { + const int split_offset = params.num_splits_ptr[bidb]; Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) - const index_t row_offset_oaccum = (((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + const index_t row_offset_oaccum = (((split_offset + n_split_idx) * params.h + bidh) * params.seqlen_q + m_block * kBlockM) * params.d_v; - const index_t row_offset_lseaccum = ((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; + const index_t row_offset_lseaccum = ((split_offset + n_split_idx) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.oaccum_ptr) + row_offset_oaccum), Shape, Int>{}, @@ -491,5 +184,292 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_32x16_4wa ); } } +template +__forceinline__ __device__ void compute_attn_1rowblock_splitkv_mla_k64_32x16_4waves_xcore1000( + const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int n_block_min, const int n_block_max, const bool NoSplit) { + using Element = typename Kernel_traits::Element; + using ElementAccum = typename Kernel_traits::ElementAccum; + using index_t = typename Kernel_traits::index_t; + + // Shared memory. + extern __shared__ char smem_[]; + + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kHeadDim = Kernel_traits::kHeadDim; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kBlockKSmem = Kernel_traits::kBlockKSmem; + constexpr int kAtomLayoutMS = Kernel_traits::kAtomLayoutMS; + constexpr int kAtomLayoutMO = Kernel_traits::kAtomLayoutMO; + constexpr int Num_Stages = Kernel_traits::Num_Stages; + + static_assert(kBlockKSmem == 64); + + const BlockInfo binfo(params, bidb); + // if (threadIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { printf("Is_even_MN = %d, is_cumulativ = %d, seqlen_k_cache = %d, actual_seqlen_k = %d\n", Is_even_MN, params.is_seqlens_k_cumulative, binfo.seqlen_k_cache, binfo.actual_seqlen_k); } + // if (threadIdx.x == 0 && blockIdx.y == 1 && blockIdx.z == 0) { printf("params.knew_ptr = %p, seqlen_k_cache + seqlen_knew = %d\n", params.knew_ptr, binfo.seqlen_k_cache + (params.knew_ptr == nullptr ? 0 : params.seqlen_knew)); } + if (m_block * kBlockM >= binfo.actual_seqlen_q) return; + + // We iterate over the blocks in reverse order. This is because the last block is the only one + // that needs masking when we read K and V from global memory. Moreover, iterating in reverse + // might save us 1 register (we just need n_block instead of both n_block and n_block_max). + const index_t row_offset_q = binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb) + + m_block * kBlockM * params.q_row_stride + bidh * params.q_head_stride; + // We move K and V to the last block. + const int bidb_cache = params.cache_batch_idx == nullptr ? bidb : params.cache_batch_idx[bidb]; + const int *block_table = params.block_table == nullptr ? nullptr : params.block_table + bidb * params.block_table_batch_stride; + const int block_table_idx = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN / params.page_block_size; + const int block_table_offset = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN - block_table_idx * params.page_block_size; + const index_t row_offset_k = block_table == nullptr + ? binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb_cache) + + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride + : (bidh / params.h_h_k_ratio) * params.k_head_stride; + const index_t row_offset_v = block_table == nullptr + ? binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb_cache) + + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride + : (bidh / params.h_h_k_ratio) * params.v_head_stride; + + Tensor gQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), + Shape, Int>{}, + make_stride(params.q_row_stride, _1{})); + Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + row_offset_k), + Shape, Int>{}, + make_stride(params.k_row_stride, _1{})); + // if (threadIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { printf("k_ptr = %p, row_offset_k = %d, gK_ptr = %p\n", params.k_ptr, row_offset_k, gK.data()); } + Tensor gV = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_ptr) + row_offset_v), + Shape, Int>{}, + make_stride(params.v_row_stride, _1{})); + + Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutQ{}); + Tensor sK = make_tensor(sQ.data() + (Kernel_traits::Share_Q_K_smem ? 0 : size(sQ)), + typename Kernel_traits::SmemLayoutK424{}); + Tensor sV = make_tensor(sK.data(), typename Kernel_traits::SmemLayoutVtNoSwizzle{}); + Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed424{}); + Tensor sVtNoSwizzle = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{}); + + typename Kernel_traits::GmemTiledCopyB128 gmem_tiled_copy_Q; + auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx); + + Tensor tQgQ = gmem_thr_copy_Q.partition_S(gQ); + Tensor tQsQ = gmem_thr_copy_Q.partition_D(sQ); + + typename Kernel_traits::GmemTiledCopyB64 gmem_tiled_copy_KV; + auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx); + Tensor tKgK = gmem_thr_copy_KV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K) + Tensor tKsK = gmem_thr_copy_KV.partition_D(sK); + // wave0 and wave2 compute the same S, wave1 and wave3 compute the same S + int tidx_mma_s = tidx & 0x7F; + typename Kernel_traits::TiledMmaS tiled_mma_s; + auto thr_mma_s = tiled_mma_s.get_thread_slice(tidx_mma_s); + Tensor tSrQ = thr_mma_s.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K) + Tensor tSrK = thr_mma_s.partition_fragment_B(sK(_, _, 0)); // (MMA,MMA_N,MMA_K) + typename Kernel_traits::TiledMmaO tiled_mma_o; + auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); + // Tensor tOrVt = thr_mma_o.partition_fragment_B(sVt); // (MMA, MMA_K,MMA_N) + Tensor tOrVt = make_tensor(Shape<_4, Shape<_4, _4>, _1>{}); + + Tensor acc_o = partition_fragment_C(tiled_mma_o, Shape, Int>{}); // MMA, MMA_M, MMA_K + + // + // Copy Atom retiling + // + + auto smem_tiled_copy_Q = make_tiled_copy_A(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx_mma_s); + Tensor tSsQ = smem_thr_copy_Q.partition_S(sQ); + + auto smem_tiled_copy_K = make_tiled_copy_B(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx_mma_s); + Tensor tSsK = smem_thr_copy_K.partition_S(sK); + auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed{}, tiled_mma_o); + auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx); + int warp_offset = warp_idx / kAtomLayoutMO * 16 * 64; + int thread_offset = lane_idx / 16 * 4 * 64; + Element *Vtsmem_ptr_lds = reinterpret_cast(sVt.data().get()) + warp_offset + thread_offset; + Tensor tOsVt = make_tensor(make_smem_ptr(Vtsmem_ptr_lds), make_layout(Shape<_4, _4, Int>{}, // MMA MMA_N NUM_STAGES + Stride<_1, Int<16*128>, Int>{})); + + // PREDICATES + + // Construct identity layout for sQ and sK + Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + + // Repeat the partitioning with identity layouts + Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tKVcKV = gmem_thr_copy_KV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + + + // Prologue + + + // Read Q from gmem to smem, optionally apply rotary embedding. + Tensor tQrQ = make_fragment_like(tQgQ); + // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs + flash::copy_b128(tQgQ, tQrQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); + cute::copy(tQrQ, tQsQ); + + if constexpr (Kernel_traits::Is_Q_in_regs) { + flash::sync_threads(); + cute::copy(smem_tiled_copy_Q, tSsQ, tSrQ); + flash::sync_threads(); + } + + + int n_block = n_block_max - 1; + int Ksmem_read_index = 0; + int Ksmem_write_index = 0; + // We don't need to clear the sK smem tiles since we'll mask out the scores anyway. + Tensor tKrK = make_fragment_like(tKgK); + int row_offset = tidx / 16 + n_block * kBlockN; + int virtual_page_idx = row_offset / params.page_block_size; + int page_offset = row_offset - virtual_page_idx * params.page_block_size; + int32_t page_idx = block_table[virtual_page_idx]; + flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, page_idx, page_offset, binfo.actual_seqlen_k - n_block * kBlockN); + + + // flash::cp_async_wait<0>(); + // __syncthreads(); + // if (tidx == 0 && blockIdx.y == 0 && blockIdx.z == 0) { print(tKsK); } + // __syncthreads(); + + clear(acc_o); + + flash::Softmax(acc_o)> softmax; + + flash::Mask mask(binfo.actual_seqlen_k, binfo.actual_seqlen_q, params.ngroups, binfo.tot_seqlen_k, params.cp_world_size, params.cp_rank); + + // For performance reason, we separate out two kinds of iterations: + // those that need masking on S, and those that don't. + // We need masking on S for the very last block when K and V has length not multiple of kBlockN. + // We also need masking on S if it's causal, for the last ceil_div(kBlockM, kBlockN) blocks. + // We will have at least 1 "masking" iteration. + + // If not even_N, then seqlen_k might end in the middle of a block. In that case we need to + // mask 2 blocks (e.g. when kBlockM == kBlockN), not just 1. + constexpr int n_masking_steps = (!Is_causal) + ? 1 + : ((Is_even_MN && Is_causal) ? cute::ceil_div(kBlockM, kBlockN) : cute::ceil_div(kBlockM, kBlockN) + 1); + #pragma unroll + for (int masking_step = 0; masking_step < n_masking_steps; ++masking_step, --n_block) { + if (n_block > n_block_min) { + // prefetch load page index + row_offset -= kBlockN; + virtual_page_idx = row_offset / params.page_block_size; + page_offset = row_offset - virtual_page_idx * params.page_block_size; + page_idx = block_table[virtual_page_idx]; + } + Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) + cute::copy(tKrK, tKsK(_, _, _, Ksmem_write_index)); + Ksmem_write_index ^= 1; + clear(acc_s); + + flash::sync_threads(); + if (n_block > n_block_min) { + flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, page_idx, page_offset); + } + + flash::gemm( + acc_s, tSrQ, tSrK, tSsQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, + smem_thr_copy_Q, smem_thr_copy_K + ); + // if (cute::thread0()) { print(acc_s); } + + mask.template apply_mask( + acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 + ); + + // We have key_padding_mask so we'll need to Check_inf + masking_step == 0 + ? softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2) + : softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); + // if (cute::thread0()) { print(scores_max); print(scores_sum); print(scores); } + + // Convert acc_s from fp32 to fp16/bf16 + //Tensor rP = flash::convert_type(acc_s); + CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) + // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. + //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); + lds4x4_with_swizzle424(tOsVt(_, _, Ksmem_read_index), tOrVt); + CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 + Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); + permute_4x4_b16(tOrVt_permute_view); + Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); + flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); + Ksmem_read_index ^= 1; + + // This check is at the end of the loop since we always have at least 1 iteration + if (n_masking_steps > 1 && n_block <= n_block_min) { + --n_block; + break; + } + } + + // These are the iterations where we don't need masking on S + for (; n_block >= n_block_min; --n_block) { + if (n_block > n_block_min) { + // prefetch load page index + row_offset -= kBlockN; + virtual_page_idx = row_offset / params.page_block_size; + page_offset = row_offset - virtual_page_idx * params.page_block_size; + page_idx = block_table[virtual_page_idx]; + } + Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) + cute::copy(tKrK, tKsK(_, _, _, Ksmem_write_index)); + Ksmem_write_index ^= 1; + clear(acc_s); + flash::sync_threads(); + if (n_block > n_block_min) { + // Advance gK + flash::copy_b64_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, page_idx, page_offset); + } + + flash::gemm( + acc_s, tSrQ, tSrK, tSsQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, + smem_thr_copy_Q, smem_thr_copy_K + ); + + mask.template apply_mask( + acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 + ); + softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); + + //Tensor rP = flash::convert_type(acc_s); + CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) + // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. + //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); + lds4x4_with_swizzle424(tOsVt(_, _, Ksmem_read_index), tOrVt); + CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 + Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); + permute_4x4_b16(tOrVt_permute_view); + Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); + + flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); + Ksmem_read_index ^= 1; + } + + // Epilogue + + + + // if (cute::thread0()) { print(lse); } + if (NoSplit) { + store_32x16(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax); + }else{ + store_32x16(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax); + } +} } // namespace flash diff --git a/csrc/flash_kernel/flash_fwd_split_kernel_k64_64x16_8waves.h b/csrc/flash_kernel/xcore1000/flash_fwd_mla_kernel_k64_64x16_8waves_xcore1000.h similarity index 76% rename from csrc/flash_kernel/flash_fwd_split_kernel_k64_64x16_8waves.h rename to csrc/flash_kernel/xcore1000/flash_fwd_mla_kernel_k64_64x16_8waves_xcore1000.h index 702569c..62ddf70 100644 --- a/csrc/flash_kernel/flash_fwd_split_kernel_k64_64x16_8waves.h +++ b/csrc/flash_kernel/xcore1000/flash_fwd_mla_kernel_k64_64x16_8waves_xcore1000.h @@ -13,297 +13,35 @@ #include "utils.h" #include "softmax.h" #include "mask.h" -#include "rotary.h" -#include "attn_mask.h" namespace flash { using namespace cute; -template -__forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_64x16_8waves(const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int num_n_splits) { - - using Element = typename Kernel_traits::Element; +template +__forceinline__ __device__ void store_64x16_xcore1000(const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, char* smem_, AccO acc_o, Softmax softmax){ using ElementAccum = typename Kernel_traits::ElementAccum; + using Element = typename Kernel_traits::Element; using index_t = typename Kernel_traits::index_t; - - // Shared memory. - extern __shared__ char smem_[]; - - // The thread index. - const int tidx = threadIdx.x; - const int warp_idx = tidx / 64; - const int lane_idx = tidx % 64; - - constexpr int kBlockM = Kernel_traits::kBlockM; - constexpr int kBlockN = Kernel_traits::kBlockN; - constexpr int kHeadDim = Kernel_traits::kHeadDim; - constexpr int kHeadDimV = Kernel_traits::kHeadDimV; - constexpr int kNWarps = Kernel_traits::kNWarps; - constexpr int kBlockKSmem = Kernel_traits::kBlockKSmem; - constexpr int kAtomLayoutMS = Kernel_traits::kAtomLayoutMS; - constexpr int kAtomLayoutMO = Kernel_traits::kAtomLayoutMO; - constexpr int Num_Stages = Kernel_traits::Num_Stages; - constexpr int kHeadDimNope = kHeadDimV; - constexpr int kHeadDimRope = kHeadDim - kHeadDimV; - - static_assert(kBlockKSmem == 64); - static_assert(Kernel_traits::Share_Q_K_smem && Kernel_traits::Is_Q_in_regs); - using GmemTiledCopyO = std::conditional_t< !Split, typename Kernel_traits::GmemTiledCopyO, typename Kernel_traits::GmemTiledCopyOaccum >; using ElementO = std::conditional_t; + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; const BlockInfo binfo(params, bidb); - // if (threadIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { printf("Is_even_MN = %d, is_cumulativ = %d, seqlen_k_cache = %d, actual_seqlen_k = %d\n", Is_even_MN, params.is_seqlens_k_cumulative, binfo.seqlen_k_cache, binfo.actual_seqlen_k); } - // if (threadIdx.x == 0 && blockIdx.y == 1 && blockIdx.z == 0) { printf("params.knew_ptr = %p, seqlen_k_cache + seqlen_knew = %d\n", params.knew_ptr, binfo.seqlen_k_cache + (params.knew_ptr == nullptr ? 0 : params.seqlen_knew)); } - if (m_block * kBlockM >= binfo.actual_seqlen_q) return; - - const int n_blocks_per_split = ((binfo.actual_seqlen_k + kBlockN - 1) / kBlockN + num_n_splits - 1) / num_n_splits; - const int n_block_min = !Is_local - ? n_split_idx * n_blocks_per_split - : std::max(n_split_idx * n_blocks_per_split, (m_block * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q - params.window_size_left) / kBlockN); - int n_block_max = std::min(cute::ceil_div(binfo.actual_seqlen_k, kBlockN), (n_split_idx + 1) * n_blocks_per_split); - if (Is_causal || Is_local) { - n_block_max = std::min(n_block_max, - cute::ceil_div((m_block + 1) * kBlockM + binfo.actual_seqlen_k - binfo.actual_seqlen_q / params.ngroups + params.window_size_right, kBlockN)); - } - - if (n_block_min >= n_block_max) { // This also covers the case where n_block_max <= 0 - // We exit early and write 0 to gOaccum and -inf to gLSEaccum. - // Otherwise we might read OOB elements from gK and gV, - // or get wrong results when we combine gOaccum from different blocks. - const index_t row_offset_o = binfo.q_offset(params.o_batch_stride, params.o_row_stride, bidb) - + m_block * kBlockM * params.o_row_stride + bidh * params.o_head_stride; - const index_t row_offset_oaccum = (((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q - + m_block * kBlockM) * params.d_v; - const index_t row_offset_lseaccum = ((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; - Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.oaccum_ptr : params.o_ptr) + (Split ? row_offset_oaccum : row_offset_o)), - Shape, Int>{}, - make_stride(Split ? kHeadDimV : params.o_row_stride, _1{})); - Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.softmax_lseaccum_ptr : params.softmax_lse_ptr) + row_offset_lseaccum), - Shape>{}, Stride<_1>{}); - - GmemTiledCopyO gmem_tiled_copy_Oaccum; - auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); - Tensor tOgOaccum = gmem_thr_copy_Oaccum.partition_D(gOaccum); - Tensor tOrOaccum = make_tensor(shape(tOgOaccum)); - clear(tOrOaccum); - // Construct identity layout for sO - Tensor cO = make_identity_tensor(make_shape(size<0>(gOaccum), size<1>(gOaccum))); // (BLK_M,BLK_K) -> (blk_m,blk_k) - // Repeat the partitioning with identity layouts - Tensor tOcO = gmem_thr_copy_Oaccum.partition_D(cO); - Tensor tOpO = make_tensor(make_shape(size<2>(tOgOaccum))); - if (!Is_even_K) { - #pragma unroll - for (int k = 0; k < size(tOpO); ++k) { tOpO(k) = get<1>(tOcO(0, 0, k)) < params.d_v; } - } - // Clear_OOB_K must be false since we don't want to write zeros to gmem - flash::copy( - gmem_tiled_copy_Oaccum, tOrOaccum, tOgOaccum, tOcO, tOpO, binfo.actual_seqlen_q - m_block * kBlockM - ); - #pragma unroll - for (int m = 0; m < size<1>(tOgOaccum); ++m) { - const int row = get<0>(tOcO(0, m, 0)); - if (row < binfo.actual_seqlen_q - m_block * kBlockM && get<1>(tOcO(0, m, 0)) == 0) { gLSEaccum(row) = Split ? -INFINITY : INFINITY; } - } - return; - } - - // We iterate over the blocks in reverse order. This is because the last block is the only one - // that needs masking when we read K and V from global memory. Moreover, iterating in reverse - // might save us 1 register (we just need n_block instead of both n_block and n_block_max). - const index_t row_offset_q = binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb) - + m_block * kBlockM * params.q_row_stride + bidh * params.q_head_stride; - // We move K and V to the last block. - const int bidb_cache = params.cache_batch_idx == nullptr ? bidb : params.cache_batch_idx[bidb]; - const int *block_table = params.block_table == nullptr ? nullptr : params.block_table + bidb * params.block_table_batch_stride; - const int block_table_idx = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN / params.page_block_size; - const int block_table_offset = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN - block_table_idx * params.page_block_size; - const index_t row_offset_k = block_table == nullptr - ? binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb_cache) - + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride - : (bidh / params.h_h_k_ratio) * params.k_head_stride; - const index_t row_offset_v = block_table == nullptr - ? binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb_cache) - + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride - : (bidh / params.h_h_k_ratio) * params.v_head_stride; - - Tensor gNopeQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), - Shape, Int>{}, - make_stride(params.q_row_stride, _1{})); - Tensor gRopeQ = make_tensor(make_gmem_ptr(gNopeQ.data().get() + kHeadDimNope), - Shape, Int>{}, - make_stride(params.q_row_stride, _1{})); - Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + row_offset_k), - Shape, Int>{}, - make_stride(params.k_row_stride, _1{})); - // if (threadIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { printf("k_ptr = %p, row_offset_k = %d, gK_ptr = %p\n", params.k_ptr, row_offset_k, gK.data()); } - Tensor gV = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_ptr) + row_offset_v), - Shape, Int>{}, - make_stride(params.v_row_stride, _1{})); - - Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), - typename Kernel_traits::SmemLayoutQ{}); - Tensor sNopeQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), - typename Kernel_traits::SmemLayoutNopeQ{}); - Tensor sRopeQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), - typename Kernel_traits::SmemLayoutRopeQ{}); - Tensor sK = make_tensor(sNopeQ.data(), - typename Kernel_traits::SmemLayoutK{}); - Tensor sV = make_tensor(sK.data(), typename Kernel_traits::SmemLayoutVtNoSwizzle{}); - Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed{}); - Tensor sVtNoSwizzle = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{}); - - typename Kernel_traits::GmemTiledCopyB128 gmem_tiled_copy_Q; - auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx); - - Tensor tNopeQgNopeQ = gmem_thr_copy_Q.partition_S(gNopeQ); - Tensor tNopeQsNopeQ = gmem_thr_copy_Q.partition_D(sNopeQ); - Tensor tRopeQgRopeQ = gmem_thr_copy_Q.partition_S(gRopeQ); - Tensor tRopeQsRopeQ = gmem_thr_copy_Q.partition_D(sRopeQ); - - typename Kernel_traits::GmemTiledCopyB32 gmem_tiled_copy_KV; - auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx); - Tensor tKgK = gmem_thr_copy_KV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K) - Tensor tKsK = gmem_thr_copy_KV.partition_D(sK); - // gemm S is 4x1 wave layout, wave(n) compute the same S with wave(n + 4) - int tidx_mma_s = tidx & 0xFF; - typename Kernel_traits::TiledMmaS tiled_mma_s; - auto thr_mma_s = tiled_mma_s.get_thread_slice(tidx_mma_s); - Tensor tSrQ = thr_mma_s.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K) - Tensor tSrNopeQ = thr_mma_s.partition_fragment_A(sNopeQ); // (MMA,MMA_M,MMA_K) - Tensor tSrRopeQ = thr_mma_s.partition_fragment_A(sRopeQ); // (MMA,MMA_M,MMA_K) - Tensor tSrK = thr_mma_s.partition_fragment_B(sK(_, _, 0)); // (MMA,MMA_N,MMA_K) typename Kernel_traits::TiledMmaO tiled_mma_o; auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); - // Tensor tOrVt = thr_mma_o.partition_fragment_B(sVt); // (MMA, MMA_K,MMA_N) - Tensor tOrVt = make_tensor(Shape<_4, Shape<_4, _4>, _1>{}); - - Tensor acc_o = partition_fragment_C(tiled_mma_o, Shape, Int>{}); // MMA, MMA_M, MMA_K - - // - // Copy Atom retiling - // - - auto smem_tiled_copy_Q = make_tiled_copy_A(typename Kernel_traits::SmemCopyAtomB64{}, tiled_mma_s); - auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx_mma_s); - Tensor tSsNopeQ = smem_thr_copy_Q.partition_S(sNopeQ); - Tensor tSsRopeQ = smem_thr_copy_Q.partition_S(sRopeQ); - - auto smem_tiled_copy_K = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomB64{}, tiled_mma_s); - auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx_mma_s); - Tensor tSsK = smem_thr_copy_K.partition_S(sK); - auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed{}, tiled_mma_o); - auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx); - int warp_offset = warp_idx / kAtomLayoutMO * 16 * 64; - int thread_offset = lane_idx / 16 * 4 * 64; - Element *Vtsmem_ptr_lds = reinterpret_cast(sVt.data().get()) + warp_offset + thread_offset; - Tensor tOsVt = make_tensor(make_smem_ptr(Vtsmem_ptr_lds), make_layout(Shape<_4, _4, Int>{}, // MMA MMA_N NUM_STAGES - Stride<_1, Int<16*128>, Int>{})); - - // PREDICATES - - // Construct identity layout for sQ and sK - Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) - Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) - - // Repeat the partitioning with identity layouts - Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) - Tensor tKVcKV = gmem_thr_copy_KV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) - - - // Prologue - - - // Read Q from gmem to smem, optionally apply rotary embedding. - Tensor tNopeQrNopeQ = make_fragment_like(tNopeQgNopeQ); - Tensor tRopeQrRopeQ = make_fragment_like(tRopeQgRopeQ); - // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs - flash::copy_b128(tNopeQgNopeQ, tNopeQrNopeQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); - flash::copy_b128(tRopeQgRopeQ, tRopeQrRopeQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); - - cute::copy(tNopeQrNopeQ, tNopeQsNopeQ); - flash::sync_threads(); - cute::copy(smem_tiled_copy_Q, tSsNopeQ, tSrNopeQ); - flash::sync_threads(); - - cute::copy(tRopeQrRopeQ, tRopeQsRopeQ); - flash::sync_threads(); - cute::copy(smem_tiled_copy_Q, tSsRopeQ, tSrRopeQ); - flash::sync_threads(); - flash::concat(tSrNopeQ, tSrRopeQ, tSrQ); - - - int n_block = n_block_max - 1; - int Ksmem_read_index = 0; - int Ksmem_write_index = 0; - // We don't need to clear the sK smem tiles since we'll mask out the scores anyway. - Tensor tKrK = make_fragment_like(tKgK); - flash::copy_b32_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block, - block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, binfo.actual_seqlen_k - n_block * kBlockN); - - - // flash::cp_async_wait<0>(); - // __syncthreads(); - // if (tidx == 0 && blockIdx.y == 0 && blockIdx.z == 0) { print(tKsK); } - // __syncthreads(); - - clear(acc_o); - - flash::Softmax(acc_o)> softmax; - - const float alibi_slope = !Has_alibi ? 0.0f : reinterpret_cast(params.alibi_slopes_ptr)[bidb * params.alibi_slopes_batch_stride + bidh] / params.scale_softmax; - flash::Mask mask(binfo.actual_seqlen_k, binfo.actual_seqlen_q, params.ngroups, params.window_size_left, params.window_size_right, alibi_slope); - - for (; n_block >= n_block_min; --n_block) { - Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) - cute::copy(tKrK, tKsK(_, _, _, Ksmem_write_index)); - Ksmem_write_index ^= 1; - clear(acc_s); - flash::sync_threads(); - if (n_block > n_block_min) { - // Advance gK - flash::copy_b32_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, - block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size); - } - - flash::gemm_opt( - acc_s, tSrQ, tSrK, tSsNopeQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, - smem_thr_copy_Q, smem_thr_copy_K - ); - - - if constexpr (Is_softcap){ - flash::apply_softcap(acc_s, params.softcap); - } - - mask.template apply_mask( - acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 - ); - softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); - - //Tensor rP = flash::convert_type(acc_s); - CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) - // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) - // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. - //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); - lds4x4_with_swizzle424(tOsVt(_, _, Ksmem_read_index), tOrVt); - CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 - Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); - permute_4x4_b16(tOrVt_permute_view); - Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); - - flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); - Ksmem_read_index ^= 1; - } - - // Epilogue - - Tensor lse = softmax.template normalize_softmax_lse(acc_o, params.scale_softmax); Tensor acc_o_view = make_tensor(acc_o.data(), make_layout(Shape<_4, Shape<_4, _4>>{}, Stride<_1, Shape<_4, _16>>{})); @@ -317,7 +55,6 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_64x16_8wa acc_o_copy(row, make_coord(col, k)) = acc_o_view(col, make_coord(row, k)); } } - // if (cute::thread0()) { print(lse); } if constexpr (!Split) { Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) // Partition sO to match the accumulator partitioning @@ -332,8 +69,6 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_64x16_8wa Stride<_1, _16>{})); - // sOaccum is larger than sQ, so we need to syncthreads here - // TODO: allocate enough smem for sOaccum if constexpr (Kernel_traits::Share_Q_K_smem) { flash::sync_threads(); } cute::copy(tOrO, tOsO); @@ -382,10 +117,11 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_64x16_8wa tOrOaccum, tOgOaccum, tOcO, params.d_v, binfo.actual_seqlen_q - m_block * kBlockM ); } else { + const int split_offset = params.num_splits_ptr[bidb]; Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) - const index_t row_offset_oaccum = (((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + const index_t row_offset_oaccum = (((split_offset + n_split_idx) * params.h + bidh) * params.seqlen_q + m_block * kBlockM) * params.d_v; - const index_t row_offset_lseaccum = ((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; + const index_t row_offset_lseaccum = ((split_offset + n_split_idx) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.oaccum_ptr) + row_offset_oaccum), Shape, Int>{}, @@ -471,4 +207,237 @@ __forceinline__ __device__ void compute_attn_1rowblock_splitkv_k64_mla_64x16_8wa } } +template +__forceinline__ __device__ void compute_attn_1rowblock_splitkv_mla_k64_64x16_8waves_xcore1000( + const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int n_block_min, const int n_block_max, const bool NoSplit) { + using Element = typename Kernel_traits::Element; + using ElementAccum = typename Kernel_traits::ElementAccum; + using index_t = typename Kernel_traits::index_t; + + // Shared memory. + extern __shared__ char smem_[]; + + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kHeadDim = Kernel_traits::kHeadDim; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kBlockKSmem = Kernel_traits::kBlockKSmem; + constexpr int kAtomLayoutMS = Kernel_traits::kAtomLayoutMS; + constexpr int kAtomLayoutMO = Kernel_traits::kAtomLayoutMO; + constexpr int Num_Stages = Kernel_traits::Num_Stages; + constexpr int kHeadDimNope = kHeadDimV; + constexpr int kHeadDimRope = kHeadDim - kHeadDimV; + + static_assert(kBlockKSmem == 64); + static_assert(Kernel_traits::Share_Q_K_smem && Kernel_traits::Is_Q_in_regs); + + + const BlockInfo binfo(params, bidb); + if (m_block * kBlockM >= binfo.actual_seqlen_q) return; + + // We iterate over the blocks in reverse order. This is because the last block is the only one + // that needs masking when we read K and V from global memory. Moreover, iterating in reverse + // might save us 1 register (we just need n_block instead of both n_block and n_block_max). + const index_t row_offset_q = binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb) + + m_block * kBlockM * params.q_row_stride + bidh * params.q_head_stride; + // We move K and V to the last block. + const int bidb_cache = params.cache_batch_idx == nullptr ? bidb : params.cache_batch_idx[bidb]; + const int *block_table = params.block_table == nullptr ? nullptr : params.block_table + bidb * params.block_table_batch_stride; + const int block_table_idx = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN / params.page_block_size; + const int block_table_offset = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN - block_table_idx * params.page_block_size; + const index_t row_offset_k = block_table == nullptr + ? binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb_cache) + + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride + : (bidh / params.h_h_k_ratio) * params.k_head_stride; + const index_t row_offset_v = block_table == nullptr + ? binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb_cache) + + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride + : (bidh / params.h_h_k_ratio) * params.v_head_stride; + + Tensor gNopeQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), + Shape, Int>{}, + make_stride(params.q_row_stride, _1{})); + Tensor gRopeQ = make_tensor(make_gmem_ptr(gNopeQ.data().get() + kHeadDimNope), + Shape, Int>{}, + make_stride(params.q_row_stride, _1{})); + Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + row_offset_k), + Shape, Int>{}, + make_stride(params.k_row_stride, _1{})); + // if (threadIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { printf("k_ptr = %p, row_offset_k = %d, gK_ptr = %p\n", params.k_ptr, row_offset_k, gK.data()); } + Tensor gV = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_ptr) + row_offset_v), + Shape, Int>{}, + make_stride(params.v_row_stride, _1{})); + + Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutQ{}); + Tensor sNopeQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutNopeQ{}); + Tensor sRopeQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutRopeQ{}); + Tensor sK = make_tensor(sNopeQ.data(), typename Kernel_traits::SmemLayoutK424{}); + Tensor sV = make_tensor(sK.data(), typename Kernel_traits::SmemLayoutVtNoSwizzle{}); + Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed424{}); + Tensor sVtNoSwizzle = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{}); + + + typename Kernel_traits::GmemTiledCopyB128 gmem_tiled_copy_Q; + auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx); + + Tensor tNopeQgNopeQ = gmem_thr_copy_Q.partition_S(gNopeQ); + Tensor tNopeQsNopeQ = gmem_thr_copy_Q.partition_D(sNopeQ); + Tensor tRopeQgRopeQ = gmem_thr_copy_Q.partition_S(gRopeQ); + Tensor tRopeQsRopeQ = gmem_thr_copy_Q.partition_D(sRopeQ); + + typename Kernel_traits::GmemTiledCopyB32 gmem_tiled_copy_KV; + auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx); + Tensor tKgK = gmem_thr_copy_KV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K) + Tensor tKsK = gmem_thr_copy_KV.partition_D(sK); + // gemm S is 4x1 wave layout, wave(n) compute the same S with wave(n + 4) + int tidx_mma_s = tidx & 0xFF; + typename Kernel_traits::TiledMmaS tiled_mma_s; + auto thr_mma_s = tiled_mma_s.get_thread_slice(tidx_mma_s); + Tensor tSrQ = thr_mma_s.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K) + Tensor tSrNopeQ = thr_mma_s.partition_fragment_A(sNopeQ); // (MMA,MMA_M,MMA_K) + Tensor tSrRopeQ = thr_mma_s.partition_fragment_A(sRopeQ); // (MMA,MMA_M,MMA_K) + Tensor tSrK = thr_mma_s.partition_fragment_B(sK(_, _, 0)); // (MMA,MMA_N,MMA_K) + typename Kernel_traits::TiledMmaO tiled_mma_o; + auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); + Tensor tOrVt = make_tensor(Shape<_4, Shape<_4, _4>, _1>{}); + + + Tensor acc_o = partition_fragment_C(tiled_mma_o, Shape, Int>{}); // MMA, MMA_M, MMA_K + + // + // Copy Atom retiling + // + + auto smem_tiled_copy_Q = make_tiled_copy_A(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx_mma_s); + Tensor tSsNopeQ = smem_thr_copy_Q.partition_S(sNopeQ); + Tensor tSsRopeQ = smem_thr_copy_Q.partition_S(sRopeQ); + + auto smem_tiled_copy_K = make_tiled_copy_B(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx_mma_s); + Tensor tSsK = smem_thr_copy_K.partition_S(sK); + auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed{}, tiled_mma_o); + auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx); + int warp_offset = warp_idx / kAtomLayoutMO * 16 * 64; + int thread_offset = lane_idx / 16 * 4 * 64; + Element *Vtsmem_ptr_lds = reinterpret_cast(sVt.data().get()) + warp_offset + thread_offset; + Tensor tOsVt = make_tensor(make_smem_ptr(Vtsmem_ptr_lds), make_layout(Shape<_4, _4, Int>{}, // MMA MMA_N NUM_STAGES + Stride<_1, Int<16*128>, Int>{})); + + + // PREDICATES + + // Construct identity layout for sQ and sK + Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + + // Repeat the partitioning with identity layouts + Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tKVcKV = gmem_thr_copy_KV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + + + // Prologue + + + // Read Q from gmem to smem, optionally apply rotary embedding. + Tensor tNopeQrNopeQ = make_fragment_like(tNopeQgNopeQ); + Tensor tRopeQrRopeQ = make_fragment_like(tRopeQgRopeQ); + // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs + flash::copy_b128(tNopeQgNopeQ, tNopeQrNopeQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); + flash::copy_b128(tRopeQgRopeQ, tRopeQrRopeQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); + + cute::copy(tNopeQrNopeQ, tNopeQsNopeQ); + flash::sync_threads(); + cute::copy(smem_tiled_copy_Q, tSsNopeQ, tSrNopeQ); + flash::sync_threads(); + + cute::copy(tRopeQrRopeQ, tRopeQsRopeQ); + flash::sync_threads(); + cute::copy(smem_tiled_copy_Q, tSsRopeQ, tSrRopeQ); + flash::sync_threads(); + flash::concat(tSrNopeQ, tSrRopeQ, tSrQ); + + + int n_block = n_block_max - 1; + int Ksmem_read_index = 0; + int Ksmem_write_index = 0; + // We don't need to clear the sK smem tiles since we'll mask out the scores anyway. + Tensor tKrK = make_fragment_like(tKgK); + int row_offset = tidx / 32 + n_block * kBlockN; + int virtual_page_idx = row_offset / params.page_block_size; + int page_offset = row_offset - virtual_page_idx * params.page_block_size; + int32_t page_idx = block_table[virtual_page_idx]; + flash::copy_b32_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, page_idx, page_offset, binfo.actual_seqlen_k - n_block * kBlockN); + + + // flash::cp_async_wait<0>(); + // __syncthreads(); + // if (tidx == 0 && blockIdx.y == 0 && blockIdx.z == 0) { print(tKsK); } + // __syncthreads(); + + clear(acc_o); + + flash::Softmax(acc_o)> softmax; + + flash::Mask mask(binfo.actual_seqlen_k, binfo.actual_seqlen_q, params.ngroups, binfo.tot_seqlen_k, params.cp_world_size, params.cp_rank); + + for (; n_block >= n_block_min; --n_block) { + if (n_block > n_block_min) { + // prefetch load page index + row_offset -= kBlockN; + virtual_page_idx = row_offset / params.page_block_size; + page_offset = row_offset - virtual_page_idx * params.page_block_size; + page_idx = block_table[virtual_page_idx]; + } + Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) + cute::copy(tKrK, tKsK(_, _, _, Ksmem_write_index)); + Ksmem_write_index ^= 1; + clear(acc_s); + flash::sync_threads(); + if (n_block > n_block_min) { + // Advance gK + flash::copy_b32_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, page_idx, page_offset); + } + + flash::gemm( + acc_s, tSrQ, tSrK, tSsNopeQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, + smem_thr_copy_Q, smem_thr_copy_K + ); + + mask.template apply_mask( + acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 + ); + softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); + + lds4x4_with_swizzle424(tOsVt(_, _, Ksmem_read_index), tOrVt); + CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) + CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 + Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); + permute_4x4_b16(tOrVt_permute_view); + Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); + + flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); + Ksmem_read_index ^= 1; + } + + // Epilogue + if (NoSplit) { + store_64x16_xcore1000(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax); + }else{ + store_64x16_xcore1000(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax); + } + +} + } // namespace flash diff --git a/csrc/flash_kernel/xcore1000/flash_fwd_sparse_mla_kernel_k64_64x16_8waves_xcore1000.h b/csrc/flash_kernel/xcore1000/flash_fwd_sparse_mla_kernel_k64_64x16_8waves_xcore1000.h new file mode 100644 index 0000000..c1866f2 --- /dev/null +++ b/csrc/flash_kernel/xcore1000/flash_fwd_sparse_mla_kernel_k64_64x16_8waves_xcore1000.h @@ -0,0 +1,275 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#pragma once + +#include + +#include +#include +#include + +#include "block_info.h" +#include "kernel_traits.h" +#include "utils.h" +#include "softmax.h" +#include "mask.h" +#include "flash_fwd_mla_kernel_k64_64x16_8waves_xcore1000.h" + +namespace flash { + +using namespace cute; + + +template +__forceinline__ __device__ void compute_attn_1rowblock_splitkv_sparse_mla_k64_64x16_8waves_xcore1000( + const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int n_block_min, const int n_block_max, const bool NoSplit) { + using Element = typename Kernel_traits::Element; + using ElementAccum = typename Kernel_traits::ElementAccum; + using index_t = typename Kernel_traits::index_t; + + // Shared memory. + extern __shared__ char smem_[]; + + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kHeadDim = Kernel_traits::kHeadDim; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kBlockKSmem = Kernel_traits::kBlockKSmem; + constexpr int kAtomLayoutMS = Kernel_traits::kAtomLayoutMS; + constexpr int kAtomLayoutMO = Kernel_traits::kAtomLayoutMO; + constexpr int Num_Stages = Kernel_traits::Num_Stages; + constexpr int kBlockTopK = Kernel_traits::kBlockTopK; + constexpr int kHeadDimNope = kHeadDimV; + constexpr int kHeadDimRope = kHeadDim - kHeadDimV; + + static_assert(kBlockKSmem == 64); + static_assert(Kernel_traits::Share_Q_K_smem && Kernel_traits::Is_Q_in_regs); + + const BlockInfo binfo(params, bidb); + if (m_block * kBlockM >= binfo.actual_seqlen_q) return; + + // We iterate over the blocks in reverse order. This is because the last block is the only one + // that needs masking when we read K and V from global memory. Moreover, iterating in reverse + // might save us 1 register (we just need n_block instead of both n_block and n_block_max). + const index_t row_offset_q = binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb) + + m_block * kBlockM * params.q_row_stride + bidh * params.q_head_stride; + // We move K and V to the last block. + const int bidb_cache = params.cache_batch_idx == nullptr ? bidb : params.cache_batch_idx[bidb]; + + const int s_q_idx = params.ngroups >= kBlockM ? m_block / (params.ngroups / kBlockM) : 0; //s_q_idx, ngroups = head_q_ori + const index_t offset_indices = bidb * params.indices_batch_stride + s_q_idx * params.indices_row_stride; + + const int* gIndices = params.indices_ptr + offset_indices; // top_k values + const bool indices_all_valid_per_q = params.indices_all_valid_per_q_ptr[bidb * params.indices_all_valid_per_q_batch_stride + s_q_idx]; + + // will calculate row_offset_k later + + const index_t row_offset_k = (bidh / params.h_h_k_ratio) * params.k_head_stride; + const index_t row_offset_v = (bidh / params.h_h_k_ratio) * params.v_head_stride; + + Tensor gNopeQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), + Shape, Int>{}, + make_stride(params.q_row_stride, _1{})); + Tensor gRopeQ = make_tensor(make_gmem_ptr(gNopeQ.data().get() + kHeadDimNope), + Shape, Int>{}, + make_stride(params.q_row_stride, _1{})); + Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + row_offset_k), + Shape, Int>{}, + make_stride(params.k_row_stride, _1{})); + + Tensor gV = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_ptr) + row_offset_v), + Shape, Int>{}, + make_stride(params.v_row_stride, _1{})); + + Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutQ{}); + Tensor sNopeQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutNopeQ{}); + Tensor sRopeQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutRopeQ{}); + Tensor sK = make_tensor(sNopeQ.data(), typename Kernel_traits::SmemLayoutK424{}); // kBlockN * kheadDim * NumStages + Tensor sV = make_tensor(sK.data(), typename Kernel_traits::SmemLayoutVtNoSwizzle{}); // kBlockN * kHeadDimV + Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed424{}); // kBlockN * kheadDimV * NumStages + Tensor sVtNoSwizzle = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{}); // kBlockN * kheadDimV * NumStages + + int32_t* indices_smem_ptr = reinterpret_cast(reinterpret_cast(smem_) + (size(sK) * sizeof(Element) + 3) / 4 * 4); + + typename Kernel_traits::GmemTiledCopyB128 gmem_tiled_copy_Q; + auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx); + + Tensor tNopeQgNopeQ = gmem_thr_copy_Q.partition_S(gNopeQ); + Tensor tNopeQsNopeQ = gmem_thr_copy_Q.partition_D(sNopeQ); + Tensor tRopeQgRopeQ = gmem_thr_copy_Q.partition_S(gRopeQ); + Tensor tRopeQsRopeQ = gmem_thr_copy_Q.partition_D(sRopeQ); + + typename Kernel_traits::GmemTiledCopyB32 gmem_tiled_copy_KV; + auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx); + Tensor tKgK = gmem_thr_copy_KV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K) + Tensor tKsK = gmem_thr_copy_KV.partition_D(sK); + // gemm S is 4x1 wave layout, wave(n) compute the same S with wave(n + 4) + int tidx_mma_s = tidx & 0xFF; + typename Kernel_traits::TiledMmaS tiled_mma_s; + auto thr_mma_s = tiled_mma_s.get_thread_slice(tidx_mma_s); + Tensor tSrQ = thr_mma_s.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K) + Tensor tSrNopeQ = thr_mma_s.partition_fragment_A(sNopeQ); // (MMA,MMA_M,MMA_K) + Tensor tSrRopeQ = thr_mma_s.partition_fragment_A(sRopeQ); // (MMA,MMA_M,MMA_K) + Tensor tSrK = thr_mma_s.partition_fragment_B(sK(_, _, 0)); // (MMA,MMA_N,MMA_K) + typename Kernel_traits::TiledMmaO tiled_mma_o; + auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); + Tensor tOrVt = make_tensor(Shape<_4, Shape<_4, _4>, _1>{}); + + + Tensor acc_o = partition_fragment_C(tiled_mma_o, Shape, Int>{}); // MMA, MMA_M, MMA_K + + // + // Copy Atom retiling + // + + auto smem_tiled_copy_Q = make_tiled_copy_A(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx_mma_s); + Tensor tSsNopeQ = smem_thr_copy_Q.partition_S(sNopeQ); + Tensor tSsRopeQ = smem_thr_copy_Q.partition_S(sRopeQ); + + auto smem_tiled_copy_K = make_tiled_copy_B(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx_mma_s); + Tensor tSsK = smem_thr_copy_K.partition_S(sK); + auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed{}, tiled_mma_o); + auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx); + int warp_offset = warp_idx / kAtomLayoutMO * 16 * 64; + int thread_offset = lane_idx / 16 * 4 * 64; + Element *Vtsmem_ptr_lds = reinterpret_cast(sVt.data().get()) + warp_offset + thread_offset; + Tensor tOsVt = make_tensor(make_smem_ptr(Vtsmem_ptr_lds), make_layout(Shape<_4, _4, Int>{}, // MMA MMA_N NUM_STAGES + Stride<_1, Int<16*128>, Int>{})); + + + // PREDICATES + + // Construct identity layout for sQ and sK + Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + + // Repeat the partitioning with identity layouts + Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tKVcKV = gmem_thr_copy_KV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + + + // Prologue + // Read Q from gmem to smem, optionally apply rotary embedding. + Tensor tNopeQrNopeQ = make_fragment_like(tNopeQgNopeQ); + Tensor tRopeQrRopeQ = make_fragment_like(tRopeQgRopeQ); + // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs + flash::copy_b128(tNopeQgNopeQ, tNopeQrNopeQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); + flash::copy_b128(tRopeQgRopeQ, tRopeQrRopeQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); + + cute::copy(tNopeQrNopeQ, tNopeQsNopeQ); + flash::sync_threads(); + cute::copy(smem_tiled_copy_Q, tSsNopeQ, tSrNopeQ); + flash::sync_threads(); + + cute::copy(tRopeQrRopeQ, tRopeQsRopeQ); + flash::sync_threads(); + cute::copy(smem_tiled_copy_Q, tSsRopeQ, tSrRopeQ); + flash::sync_threads(); + flash::concat(tSrNopeQ, tSrRopeQ, tSrQ); + + + int n_block = n_block_max - 1; + int Ksmem_read_index = 0; + int Ksmem_write_index = 0; + // We don't need to clear the sK smem tiles since we'll mask out the scores anyway. + Tensor tKrK = make_fragment_like(tKgK); + + int indices_start_block = (n_block * kBlockN) / kBlockTopK * kBlockTopK; + int offset_indices_per_q = indices_start_block + (tidx << 2); + int4 indices_vec = {-1, -1, -1, -1}; + + if (offset_indices_per_q + 4 <= params.topk) { + indices_vec = __ldg(reinterpret_cast(&gIndices[offset_indices_per_q])); + } + else if (offset_indices_per_q < params.topk) { + indices_vec.x = offset_indices_per_q + 0 < params.topk ? gIndices[offset_indices_per_q + 0] : -1; + indices_vec.y = offset_indices_per_q + 1 < params.topk ? gIndices[offset_indices_per_q + 1] : -1; + indices_vec.z = offset_indices_per_q + 2 < params.topk ? gIndices[offset_indices_per_q + 2] : -1; + indices_vec.w = offset_indices_per_q + 3 < params.topk ? gIndices[offset_indices_per_q + 3] : -1; + } + + *((int4*)(&indices_smem_ptr[offset_indices_per_q % kBlockTopK])) = indices_vec; + flash::sync_threads(); + + uint32_t row_offset = tidx / (kNWarps * 64 / kBlockN) + n_block * kBlockN; + int32_t topk_sparse_idx = indices_smem_ptr[row_offset % kBlockTopK]; + topk_sparse_idx = topk_sparse_idx < 0 ? 0 : topk_sparse_idx; + + flash::copy_b32_sparse(gK, tKgK, tKrK, tKVcKV, params.d, n_block, + params.d, topk_sparse_idx, + params.topk - n_block * kBlockN); + + flash::cp_async_wait<0>(); + + clear(acc_o); + + flash::Softmax(acc_o)> softmax; + + flash::Mask mask(params.topk, binfo.actual_seqlen_q, params.ngroups); + + for (; n_block >= n_block_min; --n_block) { + Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) + cute::copy(tKrK, tKsK(_, _, _, Ksmem_write_index)); + Ksmem_write_index ^= 1; + clear(acc_s); + flash::sync_threads(); + + flash::gemm( + acc_s, tSrQ, tSrK, tSsNopeQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, + smem_thr_copy_Q, smem_thr_copy_K + ); + + mask.template apply_sparse_attn_mask(acc_s, n_block * kBlockN, indices_smem_ptr, indices_all_valid_per_q); + + if ((n_block * kBlockN) % kBlockTopK == 0 && n_block > n_block_min) { + flash::barrier(); + indices_start_block = ((n_block - 1) * kBlockN) / kBlockTopK * kBlockTopK; + offset_indices_per_q = indices_start_block + (tidx << 2); + indices_vec = __ldg(reinterpret_cast(&gIndices[offset_indices_per_q])); + *((int4*)(&indices_smem_ptr[offset_indices_per_q % kBlockTopK])) = indices_vec; + flash::sync_threads(); + } + + n_block == n_block_max - 1 + ? softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2) + : softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); + + if (n_block > n_block_min) { + row_offset -= kBlockN; + topk_sparse_idx = indices_smem_ptr[row_offset % kBlockTopK]; + topk_sparse_idx = topk_sparse_idx < 0 ? 0 : topk_sparse_idx; + flash::copy_b32_sparse(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, params.d, topk_sparse_idx); + } + + CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) + lds4x4_with_swizzle424(tOsVt(_, _, Ksmem_read_index), tOrVt); + CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 + Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); + permute_4x4_b16(tOrVt_permute_view); + Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); + + flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); + + Ksmem_read_index ^= 1; + } + + // Epilogue + if (NoSplit) { + store_64x16_xcore1000(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax); + }else{ + store_64x16_xcore1000(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax); + } +} + +} // namespace flash \ No newline at end of file diff --git a/csrc/flash_kernel/xcore1000/sparse_prefill_kernel_64x16_8waves_xcore1000.h b/csrc/flash_kernel/xcore1000/sparse_prefill_kernel_64x16_8waves_xcore1000.h new file mode 100644 index 0000000..d6eba85 --- /dev/null +++ b/csrc/flash_kernel/xcore1000/sparse_prefill_kernel_64x16_8waves_xcore1000.h @@ -0,0 +1,347 @@ +#pragma once + +#include + +#include +#include +#include + +#include "block_info.h" +#include "kernel_traits.h" +#include "utils.h" +#include "softmax.h" +#include "mask.h" + +namespace flash { + +using namespace cute; + +template +__forceinline__ __device__ void sparse_attn_fwd_kernel(const Params ¶ms) { + constexpr bool Is_local = false; + + using Element = typename Kernel_traits::Element; + using ElementAccum = typename Kernel_traits::ElementAccum; + using index_t = typename Kernel_traits::index_t; + + // Shared memory. + extern __shared__ char smem_[]; + + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kHeadDim = Kernel_traits::kHeadDim; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kBlockKSmem = Kernel_traits::kBlockKSmem; + constexpr int kAtomLayoutMS = Kernel_traits::kAtomLayoutMS; + constexpr int kAtomLayoutMO = Kernel_traits::kAtomLayoutMO; + constexpr int Num_Stages = Kernel_traits::Num_Stages; + constexpr int kBlockTopK = Kernel_traits::kBlockTopK; + constexpr int kHeadDimNope = kHeadDimV; + constexpr int kHeadDimRope = kHeadDim - kHeadDimV; + + static_assert(kBlockKSmem == 64); + static_assert(Kernel_traits::Share_Q_K_smem && Kernel_traits::Is_Q_in_regs); + + using GmemTiledCopyO = typename Kernel_traits::GmemTiledCopyO; + + using ElementO = Element; + + const int bidb = 0; + const int h_q_idx = blockIdx.x % (params.h_q / kBlockM); //q_h_idx + const int s_q_idx = blockIdx.x / (params.h_q / kBlockM); //s_q_idx + const int q_block_idx = h_q_idx * kBlockM; + // const BlockInfo binfo(params, bidb); + if (q_block_idx >= params.h_q) return; + + const int n_block_min = 0; + int n_block_max = cute::ceil_div(params.topk, kBlockN); + // if (Is_causal || Is_local) { + // n_block_max = std::min(n_block_max, + // cute::ceil_div((m_block + 1) * kBlockM + params.s_kv - params.s_q / params.ngroups + params.window_size_right, kBlockN)); + // + + // We iterate over the blocks in reverse order. This is because the last block is the only one + // that needs masking when we read K and V from global memory. Moreover, iterating in reverse + // might save us 1 register (we just need n_block instead of both n_block and n_block_max). + const index_t row_offset_q = q_block_idx * params.q_head_stride + s_q_idx * params.q_row_stride; + const index_t row_offset_k = (n_block_max - 1) * kBlockN * params.k_row_stride; + const index_t offset_indices = s_q_idx * params.stride_indices_s_q; + // We move K and V to the last block. + // const int bidb_cache = bidb; + // const int *block_table = nullptr; + const int* gIndices = params.indices_ptr + offset_indices; + const bool indices_all_valid_per_q = params.indices_all_valid_per_q_ptr[s_q_idx]; + + Tensor gNopeQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), + Shape, Int>{}, + make_stride(params.q_head_stride, _1{})); + Tensor gRopeQ = make_tensor(make_gmem_ptr(gNopeQ.data().get() + kHeadDimNope), + Shape, Int>{}, + make_stride(params.q_head_stride, _1{})); + Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.kv_ptr)), + Shape, Int>{}, + make_stride(params.k_row_stride, _1{})); + + Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutQ{}); + Tensor sNopeQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutNopeQ{}); + Tensor sRopeQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutRopeQ{}); + Tensor sK = make_tensor(sNopeQ.data(), + typename Kernel_traits::SmemLayoutK424{}); + Tensor sV = make_tensor(sK.data(), typename Kernel_traits::SmemLayoutVtNoSwizzle{}); + Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed424{}); + + int32_t* indices_smem_ptr = reinterpret_cast(reinterpret_cast(smem_) + (size(sK) * sizeof(Element) + 3) / 4 * 4); + + typename Kernel_traits::GmemTiledCopyB128 gmem_tiled_copy_Q; + auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx); + + Tensor tNopeQgNopeQ = gmem_thr_copy_Q.partition_S(gNopeQ); + Tensor tNopeQsNopeQ = gmem_thr_copy_Q.partition_D(sNopeQ); + Tensor tRopeQgRopeQ = gmem_thr_copy_Q.partition_S(gRopeQ); + Tensor tRopeQsRopeQ = gmem_thr_copy_Q.partition_D(sRopeQ); + + typename Kernel_traits::GmemTiledCopyB32 gmem_tiled_copy_KV; + auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx); + Tensor tKgK = gmem_thr_copy_KV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K) + Tensor tKsK = gmem_thr_copy_KV.partition_D(sK); + // gemm S is 4x1 wave layout, wave(n) compute the same S with wave(n + 4) + int tidx_mma_s = tidx & 0xFF; + typename Kernel_traits::TiledMmaS tiled_mma_s; + auto thr_mma_s = tiled_mma_s.get_thread_slice(tidx_mma_s); + Tensor tSrQ = thr_mma_s.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K) + Tensor tSrNopeQ = thr_mma_s.partition_fragment_A(sNopeQ); // (MMA,MMA_M,MMA_K) + Tensor tSrRopeQ = thr_mma_s.partition_fragment_A(sRopeQ); // (MMA,MMA_M,MMA_K) + Tensor tSrK = thr_mma_s.partition_fragment_B(sK(_, _, 0)); // (MMA,MMA_N,MMA_K) + typename Kernel_traits::TiledMmaO tiled_mma_o; + auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); + // Tensor tOrVt = thr_mma_o.partition_fragment_B(sVt); // (MMA, MMA_K,MMA_N) + Tensor tOrVt = make_tensor(Shape<_4, Shape<_4, _4>, _1>{}); + + Tensor acc_o = partition_fragment_C(tiled_mma_o, Shape, Int>{}); // MMA, MMA_M, MMA_K + + // + // Copy Atom retiling + // + + auto smem_tiled_copy_Q = make_tiled_copy_A(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx_mma_s); + Tensor tSsNopeQ = smem_thr_copy_Q.partition_S(sNopeQ); + Tensor tSsRopeQ = smem_thr_copy_Q.partition_S(sRopeQ); + + auto smem_tiled_copy_K = make_tiled_copy_B(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx_mma_s); + Tensor tSsK = smem_thr_copy_K.partition_S(sK); + auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed{}, tiled_mma_o); + auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx); + int warp_offset = warp_idx / kAtomLayoutMO * 16 * 64; + int thread_offset = lane_idx / 16 * 4 * 64; + Element *Vtsmem_ptr_lds = reinterpret_cast(sVt.data().get()) + warp_offset + thread_offset; + Tensor tOsVt = make_tensor(make_smem_ptr(Vtsmem_ptr_lds), make_layout(Shape<_4, _4, Int>{}, // MMA MMA_N NUM_STAGES + Stride<_1, Int<16*128>, Int>{})); + + // PREDICATES + + // Construct identity layout for sQ and sK + Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + + // Repeat the partitioning with identity layouts + Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tKVcKV = gmem_thr_copy_KV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + + + // Prologue + + // Read Q from gmem to smem, optionally apply rotary embedding. + Tensor tNopeQrNopeQ = make_fragment_like(tNopeQgNopeQ); + Tensor tRopeQrRopeQ = make_fragment_like(tRopeQgRopeQ); + // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs + flash::copy_b128(tNopeQgNopeQ, tNopeQrNopeQ, tQcQ, params.d_qk, params.h_q - q_block_idx); + flash::copy_b128(tRopeQgRopeQ, tRopeQrRopeQ, tQcQ, params.d_qk, params.h_q - q_block_idx); + + cute::copy(tNopeQrNopeQ, tNopeQsNopeQ); + flash::sync_threads(); + cute::copy(smem_tiled_copy_Q, tSsNopeQ, tSrNopeQ); + flash::sync_threads(); + + cute::copy(tRopeQrRopeQ, tRopeQsRopeQ); + flash::sync_threads(); + + cute::copy(smem_tiled_copy_Q, tSsRopeQ, tSrRopeQ); + flash::sync_threads(); + + flash::concat(tSrNopeQ, tSrRopeQ, tSrQ); + + int n_block = n_block_max - 1; + int Ksmem_read_index = 0; + int Ksmem_write_index = 0; + // We don't need to clear the sK smem tiles since we'll mask out the scores anyway. + Tensor tKrK = make_fragment_like(tKgK); + + int indices_start_block = (n_block * kBlockN) / kBlockTopK * kBlockTopK; + int offset_indices_per_q = indices_start_block + (tidx << 2); + int4 indices_vec = {-1, -1, -1, -1}; + + if (offset_indices_per_q + 4 <= params.topk) { + indices_vec = __ldg(reinterpret_cast(&gIndices[offset_indices_per_q])); + } + else if (offset_indices_per_q < params.topk) { + indices_vec.x = offset_indices_per_q + 0 < params.topk ? gIndices[offset_indices_per_q + 0] : -1; + indices_vec.y = offset_indices_per_q + 1 < params.topk ? gIndices[offset_indices_per_q + 1] : -1; + indices_vec.z = offset_indices_per_q + 2 < params.topk ? gIndices[offset_indices_per_q + 2] : -1; + indices_vec.w = offset_indices_per_q + 3 < params.topk ? gIndices[offset_indices_per_q + 3] : -1; + } + + *((int4*)(&indices_smem_ptr[offset_indices_per_q % kBlockTopK])) = indices_vec; + flash::sync_threads(); + + uint32_t row_offset = tidx / 32 + n_block * kBlockN; + int32_t topk_sparse_idx = indices_smem_ptr[row_offset % kBlockTopK]; + topk_sparse_idx = topk_sparse_idx < 0 ? 0 : topk_sparse_idx; + + flash::copy_b32_sparse(gK, tKgK, tKrK, tKVcKV, params.d_qk, n_block, params.k_row_stride, topk_sparse_idx, params.topk - n_block * kBlockN); + + clear(acc_o); + + flash::Softmax(acc_o)> softmax; + + flash::Mask mask(params.topk, params.s_q, params.h_q); + for (; n_block >= n_block_min; --n_block) { + Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) + cute::copy(tKrK, tKsK(_, _, _, Ksmem_write_index)); + Ksmem_write_index ^= 1; + clear(acc_s); + flash::sync_threads(); + flash::gemm( + acc_s, tSrQ, tSrK, tSsNopeQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, + smem_thr_copy_Q, smem_thr_copy_K + ); + + mask.template apply_sparse_attn_mask(acc_s, n_block * kBlockN, indices_smem_ptr, indices_all_valid_per_q); + + if ((n_block * kBlockN) % kBlockTopK == 0 && n_block > n_block_min) { + flash::barrier(); + indices_start_block = ((n_block - 1) * kBlockN) / kBlockTopK * kBlockTopK; + offset_indices_per_q = indices_start_block + (tidx << 2); + int4 indices_vec = __ldg(reinterpret_cast(&gIndices[offset_indices_per_q])); + *((int4*)(&indices_smem_ptr[offset_indices_per_q % kBlockTopK])) = indices_vec; + flash::sync_threads(); + } + n_block == n_block_max - 1 + ? softmax.template softmax_rescale_o(acc_s, acc_o, params.sm_scale_div_log2) + : softmax.template softmax_rescale_o(acc_s, acc_o, params.sm_scale_div_log2); + + if (n_block > n_block_min) { + row_offset -= kBlockN; + topk_sparse_idx = indices_smem_ptr[row_offset % kBlockTopK]; + topk_sparse_idx = topk_sparse_idx < 0 ? 0 : topk_sparse_idx; + flash::copy_b32_sparse(gK, tKgK, tKrK, tKVcKV, params.d_qk, n_block - 1, params.k_row_stride, topk_sparse_idx); + } + //Tensor rP = flash::convert_type(acc_s); + CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) + // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. + //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); + lds4x4_with_swizzle424(tOsVt(_, _, Ksmem_read_index), tOrVt); + CUTE_STATIC_ASSERT_V(size<2>(tOrVt) == _1{}); // only support MMA_K = 1 + Tensor tOrVt_permute_view = make_tensor(tOrVt.data(), make_layout(make_shape(size<0>(tOrVt), size<1, 0>(tOrVt), size<1, 1>(tOrVt)))); + permute_4x4_b16(tOrVt_permute_view); + Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); + + flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); + Ksmem_read_index ^= 1; + } + // if(thread0()){print(acc_o);} + // Epilogue + + Tensor lse = softmax.template normalize_softmax_lse(acc_o, params.sm_scale); + Tensor acc_o_view = make_tensor(acc_o.data(), make_layout(Shape<_4, Shape<_4, _4>>{}, + Stride<_1, Shape<_4, _16>>{})); + Tensor acc_o_copy = make_fragment_like(acc_o_view); + #pragma unroll + for (int k = 0; k < size<1, 1>(acc_o_view); k++) { + #pragma unroll + for (int idx = 0; idx < 16; idx++) { + int row = idx / 4; + int col = idx % 4; + acc_o_copy(row, make_coord(col, k)) = acc_o_view(col, make_coord(row, k)); + } + } + // if (cute::thread0()) { print(lse); } + + Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) + // Partition sO to match the accumulator partitioning + using SmemTiledCopyO = typename Kernel_traits::SmemCopyAtomO; + CONVERT_TENSOR_TYPE(ElementAccum, ElementO, acc_o_copy, rO) + warp_offset = warp_idx * 16 * 64; + thread_offset = lane_idx % 16 * 64 + lane_idx / 16 * 16; + Element *Osmem_ptr_sts = reinterpret_cast(smem_) + warp_offset + thread_offset; + Tensor tOsO = make_tensor(make_smem_ptr(Osmem_ptr_sts), make_layout(Shape<_16, _4>{}, + Stride<_1, Int<16*64*kNWarps>>{})); + Tensor tOrO = make_tensor(rO.data(), make_layout(Shape<_16, _4>{}, + Stride<_1, _16>{})); + + + if constexpr (Kernel_traits::Share_Q_K_smem) { flash::sync_threads(); } + + cute::copy(tOrO, tOsO); + + const index_t row_offset_o = q_block_idx * params.o_head_stride + s_q_idx * params.o_row_stride; + const index_t row_offset_lseaccum = s_q_idx * params.h_q + q_block_idx; + const index_t row_offset_max_logits = row_offset_lseaccum; + + Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.out_ptr) + (row_offset_o)), + Shape, Int>{}, + make_stride(params.o_head_stride, _1{})); + Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(params.lse_ptr) + row_offset_lseaccum), + Shape>{}, Stride<_1>{}); + Tensor gMaxLogits = make_tensor(make_gmem_ptr(reinterpret_cast(params.max_logits) + row_offset_max_logits), + Shape>{}, Stride<_1>{}); + // if (tidx == 0) { printf("row_offset_o = %d, s_q_idx = %d, gOaccum = %p\n", row_offset_o, s_q_idx, gOaccum.data()); } + GmemTiledCopyO gmem_tiled_copy_Oaccum; + auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); + Tensor tOsOaccum = gmem_thr_copy_Oaccum.partition_S(sOaccum); // ((Atom,AtomNum),ATOM_M,ATOM_N) + Tensor tOgOaccum = gmem_thr_copy_Oaccum.partition_D(gOaccum); + + flash::sync_threads(); + + Tensor tOrOaccum = make_tensor(shape(tOgOaccum)); + cute::copy(gmem_tiled_copy_Oaccum, tOsOaccum, tOrOaccum); + + Tensor caccO = make_identity_tensor(Shape, Int>{}); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor taccOcO = thr_mma_o.partition_C(caccO); // (MMA,MMA_M,MMA_K) + static_assert(decltype(size<0>(taccOcO))::value == 4); + // Convert to ((2, 2), MMA_M, MMA_K) then take only the row indices. + Tensor taccOcO_row = logical_divide(taccOcO, Shape<_4>{})(make_coord(0, _), _, 0); + CUTE_STATIC_ASSERT_V(size(lse) == size(taccOcO_row)); // MMA_M + if (get<1>(taccOcO_row(0)) == 0) { + #pragma unroll + for (int mi = 0; mi < size(lse); ++mi) { + const int row = get<0>(taccOcO_row(mi)); + if (row < params.h_q - q_block_idx) { + gMaxLogits(row) = softmax.row_max(mi) * params.sm_scale * M_LOG2E; + gLSEaccum(row) = lse(mi); + } + } + } + + // Construct identity layout for sO + Tensor cO = make_identity_tensor(make_shape(size<0>(sOaccum), size<1>(sOaccum))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + // Repeat the partitioning with identity layouts + Tensor tOcO = gmem_thr_copy_Oaccum.partition_D(cO); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + // Clear_OOB_K must be false since we don't want to write zeros to gmem + flash::copy_reg_to_global( + tOrOaccum, tOgOaccum, tOcO, params.d_v, params.h_q - q_block_idx + ); +} + +} // namespace flash \ No newline at end of file diff --git a/csrc/flash_kernel/xcore1500/flash_fwd_mla_kernel_k64_64x16_8waves_xcore1500.h b/csrc/flash_kernel/xcore1500/flash_fwd_mla_kernel_k64_64x16_8waves_xcore1500.h new file mode 100644 index 0000000..2965b50 --- /dev/null +++ b/csrc/flash_kernel/xcore1500/flash_fwd_mla_kernel_k64_64x16_8waves_xcore1500.h @@ -0,0 +1,341 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#pragma once + +#include + +#include +#include +#include + +#include "block_info.h" +#include "kernel_traits.h" +#include "utils.h" +#include "softmax.h" +#include "mask.h" + +namespace flash { + +using namespace cute; + +template +__forceinline__ __device__ void store_64x16_xcore1500(const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, char* smem_, AccO acc_o, Softmax softmax){ + using ElementAccum = typename Kernel_traits::ElementAccum; + using Element = typename Kernel_traits::Element; + using index_t = typename Kernel_traits::index_t; + using GmemTiledCopyO = std::conditional_t< + !Split, + typename Kernel_traits::GmemTiledCopyO, + typename Kernel_traits::GmemTiledCopyOaccum + >; + using ElementO = std::conditional_t; + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; + + const BlockInfo binfo(params, bidb); + typename Kernel_traits::TiledMmaO tiled_mma_o; + auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); + Tensor lse = softmax.template normalize_softmax_lse(acc_o, params.scale_softmax); + + const int split_offset = params.num_splits_ptr[bidb]; + Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) + // Partition sO to match the accumulator partitioning + using SmemTiledCopyO = std::conditional_t< + !Split, + typename Kernel_traits::SmemCopyAtomO, + typename Kernel_traits::SmemCopyAtomOaccum + >; + CONVERT_TENSOR_TYPE(ElementAccum, ElementO, acc_o, rO) + auto smem_tiled_copy_Oaccum = make_tiled_copy_C(SmemTiledCopyO{}, tiled_mma_o); + auto smem_thr_copy_Oaccum = smem_tiled_copy_Oaccum.get_thread_slice(tidx); + Tensor taccOrOaccum = smem_thr_copy_Oaccum.retile_S(rO); // ((Atom,AtomNum), MMA_M, MMA_N) + Tensor taccOsOaccum = smem_thr_copy_Oaccum.partition_D(sOaccum); // ((Atom,AtomNum),PIPE_M,PIPE_N) + + + if constexpr (Split || Kernel_traits::Share_Q_K_smem) { flash::sync_threads(); } + + cute::copy(smem_tiled_copy_Oaccum, taccOrOaccum, taccOsOaccum); + + const index_t row_offset_o = binfo.q_offset(params.o_batch_stride, params.o_row_stride, bidb) + + m_block * kBlockM * params.o_row_stride + bidh * params.o_head_stride; + const index_t row_offset_oaccum = (((split_offset + n_split_idx) * params.h + bidh) * params.seqlen_q + + m_block * kBlockM) * params.d_v; + const index_t row_offset_lseaccum = ((split_offset + n_split_idx) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; + + Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.oaccum_ptr : params.o_ptr) + (Split ? row_offset_oaccum : row_offset_o)), + Shape, Int>{}, + make_stride(Split ? kHeadDimV : params.o_row_stride, _1{})); + Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.softmax_lseaccum_ptr : params.softmax_lse_ptr) + row_offset_lseaccum), + Shape>{}, Stride<_1>{}); + // if (tidx == 0) { printf("row_offset_o = %d, bidh = %d, gOaccum = %p\n", row_offset_o, bidh, gOaccum.data()); } + + GmemTiledCopyO gmem_tiled_copy_Oaccum; + auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); + Tensor tOsOaccum = gmem_thr_copy_Oaccum.partition_S(sOaccum); // ((Atom,AtomNum),ATOM_M,ATOM_N) + Tensor tOgOaccum = gmem_thr_copy_Oaccum.partition_D(gOaccum); + + flash::sync_threads(); + + Tensor tOrOaccum = make_tensor(shape(tOgOaccum)); + cute::copy(gmem_tiled_copy_Oaccum, tOsOaccum, tOrOaccum); + + Tensor caccO = make_identity_tensor(Shape, Int>{}); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor taccOcO = thr_mma_o.partition_C(caccO); // (MMA,MMA_M,MMA_K) + static_assert(decltype(size<0>(taccOcO))::value == 4); + // Convert to ((2, 2), MMA_M, MMA_K) then take only the row indices. + Tensor taccOcO_row = logical_divide(taccOcO, Shape<_4>{})(make_coord(0, _), _, 0); + CUTE_STATIC_ASSERT_V(size(lse) == size(taccOcO_row)); // MMA_M + if (get<1>(taccOcO_row(0)) == 0) { + #pragma unroll + for (int mi = 0; mi < size(lse); ++mi) { + const int row = get<0>(taccOcO_row(mi)); + if (row < binfo.actual_seqlen_q - m_block * kBlockM) { gLSEaccum(row) = lse(mi); } + } + } + + // Construct identity layout for sO + Tensor cO = make_identity_tensor(make_shape(size<0>(sOaccum), size<1>(sOaccum))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + // Repeat the partitioning with identity layouts + Tensor tOcO = gmem_thr_copy_Oaccum.partition_D(cO); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + // Clear_OOB_K must be false since we don't want to write zeros to gmem + flash::copy_reg_to_global( + tOrOaccum, tOgOaccum, tOcO, params.d_v, binfo.actual_seqlen_q - m_block * kBlockM + ); +} + +template +__forceinline__ __device__ void compute_attn_1rowblock_splitkv_mla_k64_64x16_8waves_xcore1500( + const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int n_block_min, const int n_block_max, const bool NoSplit) { + using Element = typename Kernel_traits::Element; + using ElementAccum = typename Kernel_traits::ElementAccum; + using index_t = typename Kernel_traits::index_t; + + // Shared memory. + extern __shared__ char smem_[]; + + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kHeadDim = Kernel_traits::kHeadDim; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kBlockKSmem = Kernel_traits::kBlockKSmem; + constexpr int kAtomLayoutMS = Kernel_traits::kAtomLayoutMS; + constexpr int kAtomLayoutMO = Kernel_traits::kAtomLayoutMO; + constexpr int Num_Stages = Kernel_traits::Num_Stages; + constexpr int kHeadDimNope = kHeadDimV; + constexpr int kHeadDimRope = kHeadDim - kHeadDimV; + + static_assert(kBlockKSmem == 64); + static_assert(Kernel_traits::Share_Q_K_smem && Kernel_traits::Is_Q_in_regs); + + + const BlockInfo binfo(params, bidb); + if (m_block * kBlockM >= binfo.actual_seqlen_q) return; + + // We iterate over the blocks in reverse order. This is because the last block is the only one + // that needs masking when we read K and V from global memory. Moreover, iterating in reverse + // might save us 1 register (we just need n_block instead of both n_block and n_block_max). + const index_t row_offset_q = binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb) + + m_block * kBlockM * params.q_row_stride + bidh * params.q_head_stride; + // We move K and V to the last block. + const int bidb_cache = params.cache_batch_idx == nullptr ? bidb : params.cache_batch_idx[bidb]; + const int *block_table = params.block_table == nullptr ? nullptr : params.block_table + bidb * params.block_table_batch_stride; + const int block_table_idx = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN / params.page_block_size; + const int block_table_offset = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN - block_table_idx * params.page_block_size; + const index_t row_offset_k = block_table == nullptr + ? binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb_cache) + + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride + : (bidh / params.h_h_k_ratio) * params.k_head_stride; + const index_t row_offset_v = block_table == nullptr + ? binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb_cache) + + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride + : (bidh / params.h_h_k_ratio) * params.v_head_stride; + + Tensor gNopeQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), + Shape, Int>{}, + make_stride(params.q_row_stride, _1{})); + Tensor gRopeQ = make_tensor(make_gmem_ptr(gNopeQ.data().get() + kHeadDimNope), + Shape, Int>{}, + make_stride(params.q_row_stride, _1{})); + Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + row_offset_k), + Shape, Int>{}, + make_stride(params.k_row_stride, _1{})); + // if (threadIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { printf("k_ptr = %p, row_offset_k = %d, gK_ptr = %p\n", params.k_ptr, row_offset_k, gK.data()); } + Tensor gV = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_ptr) + row_offset_v), + Shape, Int>{}, + make_stride(params.v_row_stride, _1{})); + + Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutQ{}); + Tensor sNopeQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutNopeQ{}); + Tensor sRopeQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutRopeQ{}); + + Tensor sK = make_tensor(sNopeQ.data(), typename Kernel_traits::SmemLayoutK242{}); + Tensor sV = make_tensor(sK.data(), typename Kernel_traits::SmemLayoutVtNoSwizzle{}); + Tensor sVt = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposed242{}); + Tensor sVtNoSwizzle = make_tensor(sV.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{}); + + typename Kernel_traits::GmemTiledCopyB128 gmem_tiled_copy_Q; + auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx); + + Tensor tNopeQgNopeQ = gmem_thr_copy_Q.partition_S(gNopeQ); + Tensor tNopeQsNopeQ = gmem_thr_copy_Q.partition_D(sNopeQ); + Tensor tRopeQgRopeQ = gmem_thr_copy_Q.partition_S(gRopeQ); + Tensor tRopeQsRopeQ = gmem_thr_copy_Q.partition_D(sRopeQ); + + typename Kernel_traits::GmemTiledCopyB32 gmem_tiled_copy_KV; + auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx); + Tensor tKgK = gmem_thr_copy_KV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K) + Tensor tKsK = gmem_thr_copy_KV.partition_D(sK); + // gemm S is 4x1 wave layout, wave(n) compute the same S with wave(n + 4) + int tidx_mma_s = tidx & 0xFF; + typename Kernel_traits::TiledMmaS tiled_mma_s; + auto thr_mma_s = tiled_mma_s.get_thread_slice(tidx_mma_s); + Tensor tSrQ = thr_mma_s.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K) + Tensor tSrNopeQ = thr_mma_s.partition_fragment_A(sNopeQ); // (MMA,MMA_M,MMA_K) + Tensor tSrRopeQ = thr_mma_s.partition_fragment_A(sRopeQ); // (MMA,MMA_M,MMA_K) + Tensor tSrK = thr_mma_s.partition_fragment_B(sK(_, _, 0)); // (MMA,MMA_N,MMA_K) + typename Kernel_traits::TiledMmaO tiled_mma_o; + auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); + Tensor tOrVt = thr_mma_o.partition_fragment_B(sVtNoSwizzle); // (MMA, MMA_K,MMA_N) + + Tensor acc_o = partition_fragment_C(tiled_mma_o, Shape, Int>{}); // MMA, MMA_M, MMA_K + + // + // Copy Atom retiling + // + + auto smem_tiled_copy_Q = make_tiled_copy_A(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx_mma_s); + Tensor tSsNopeQ = smem_thr_copy_Q.partition_S(sNopeQ); + Tensor tSsRopeQ = smem_thr_copy_Q.partition_S(sRopeQ); + + auto smem_tiled_copy_K = make_tiled_copy_B(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx_mma_s); + Tensor tSsK = smem_thr_copy_K.partition_S(sK); + + auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::LDSB64Trans4x16Atom{}, tiled_mma_o); + int warp_offset = warp_idx / 4 * 16; + int thread_offset = lane_idx % 4 * 4 + + lane_idx % 16 / 4 * 64 + + lane_idx / 16 * 4 * 64; + + Element *Vtsmem_ptr_lds = reinterpret_cast(sVt.data().get()) + warp_offset + thread_offset; + Tensor tOsVt = make_tensor(make_smem_ptr(Vtsmem_ptr_lds), make_layout(Shape<_4, Shape<_2, _8>, Int, Int>{}, // MMA MMA_N NUM_STAGES + Stride<_1, Stride<_32, Int>,Int<16*64>, Int>{})); + + // PREDICATES + + // Construct identity layout for sQ and sK + Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + + // Repeat the partitioning with identity layouts + Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tKVcKV = gmem_thr_copy_KV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + + + // Prologue + + + // Read Q from gmem to smem, optionally apply rotary embedding. + Tensor tNopeQrNopeQ = make_fragment_like(tNopeQgNopeQ); + Tensor tRopeQrRopeQ = make_fragment_like(tRopeQgRopeQ); + // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs + flash::copy_b128(tNopeQgNopeQ, tNopeQrNopeQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); + flash::copy_b128(tRopeQgRopeQ, tRopeQrRopeQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); + + cute::copy(tNopeQrNopeQ, tNopeQsNopeQ); + flash::sync_threads(); + cute::copy(smem_tiled_copy_Q, tSsNopeQ, tSrNopeQ); + flash::sync_threads(); + + cute::copy(tRopeQrRopeQ, tRopeQsRopeQ); + flash::sync_threads(); + cute::copy(smem_tiled_copy_Q, tSsRopeQ, tSrRopeQ); + flash::sync_threads(); + flash::concat(tSrNopeQ, tSrRopeQ, tSrQ); + + + int n_block = n_block_max - 1; + int Ksmem_read_index = 0; + int Ksmem_write_index = 0; + // We don't need to clear the sK smem tiles since we'll mask out the scores anyway. + Tensor tKrK = make_fragment_like(tKgK); + int row_offset = tidx / 32 + n_block * kBlockN; + int virtual_page_idx = row_offset / params.page_block_size; + int page_offset = row_offset - virtual_page_idx * params.page_block_size; + int32_t page_idx = block_table[virtual_page_idx]; + flash::copy_b32_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, page_idx, page_offset, binfo.actual_seqlen_k - n_block * kBlockN); + + clear(acc_o); + + flash::Softmax(acc_o)> softmax; + + flash::Mask mask(binfo.actual_seqlen_k, binfo.actual_seqlen_q, params.ngroups, binfo.tot_seqlen_k, params.cp_world_size, params.cp_rank); + + for (; n_block >= n_block_min; --n_block) { + if (n_block > n_block_min) { + // prefetch load page index + row_offset -= kBlockN; + virtual_page_idx = row_offset / params.page_block_size; + page_offset = row_offset - virtual_page_idx * params.page_block_size; + page_idx = block_table[virtual_page_idx]; + } + Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) + cute::copy(tKrK, tKsK(_, _, _, Ksmem_write_index)); + Ksmem_write_index ^= 1; + clear(acc_s); + flash::sync_threads(); + if (n_block > n_block_min) { + // Advance gK + flash::copy_b32_page_one(gK, tKgK, tKrK, tKVcKV, params.d, n_block - 1, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, page_idx, page_offset); + } + + flash::gemm( + acc_s, tSrQ, tSrK, tSsNopeQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, + smem_thr_copy_Q, smem_thr_copy_K + ); + + mask.template apply_mask( + acc_s, n_block * kBlockN, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 + ); + softmax.template softmax_rescale_o(acc_s, acc_o, params.scale_softmax_log2); + + //Tensor rP = flash::convert_type(acc_s); + // Reshape rP from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + // if using m16n8k16 or (4, MMA_M, MMA_N) if using m16n8k8. + //Tensor tOrP = make_tensor(rP.data(), flash::convert_layout_acc_Aregs(rP.layout())); + cute::copy(smem_tiled_copy_V, tOsVt(_, _, _, Ksmem_read_index), tOrVt); + CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) + Tensor tOrP = make_tensor(rP.data(), acc_s.layout()); + + flash::gemm_rr(acc_o, tOrP, tOrVt, tiled_mma_o); + Ksmem_read_index ^= 1; + } + + // Epilogue + if (NoSplit) { + store_64x16_xcore1500(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax); + }else{ + store_64x16_xcore1500(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax); + } + +} +} // namespace flash diff --git a/csrc/flash_kernel/xcore1500/flash_fwd_mla_kernel_k64_64x32_8waves_xcore1500.h b/csrc/flash_kernel/xcore1500/flash_fwd_mla_kernel_k64_64x32_8waves_xcore1500.h new file mode 100644 index 0000000..c7365e7 --- /dev/null +++ b/csrc/flash_kernel/xcore1500/flash_fwd_mla_kernel_k64_64x32_8waves_xcore1500.h @@ -0,0 +1,336 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#pragma once + +#include + +#include +#include +#include + +#include "block_info.h" +#include "kernel_traits.h" +#include "utils.h" +#include "softmax.h" +#include "mask.h" + +namespace flash { + +using namespace cute; + +template +__forceinline__ __device__ void store_64x32_xcore1500(const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, char* smem_, AccO acc_o, Softmax softmax, Tensor0& sRowmax) { + using ElementAccum = typename Kernel_traits::ElementAccum; + using Element = typename Kernel_traits::Element; + using index_t = typename Kernel_traits::index_t; + using GmemTiledCopyO = std::conditional_t< + !Split, + typename Kernel_traits::GmemTiledCopyO, + typename Kernel_traits::GmemTiledCopyOaccum + >; + using ElementO = std::conditional_t; + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; + + const BlockInfo binfo(params, bidb); + typename Kernel_traits::TiledMmaO tiled_mma_o; + auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); + Tensor lse = softmax.template normalize_softmax_lse(acc_o, sRowmax, params.scale_softmax); + + const int split_offset = params.num_splits_ptr[bidb]; + Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) + // Partition sO to match the accumulator partitioning + using SmemTiledCopyO = std::conditional_t< + !Split, + typename Kernel_traits::SmemCopyAtomO, + typename Kernel_traits::SmemCopyAtomOaccum + >; + CONVERT_TENSOR_TYPE(ElementAccum, ElementO, acc_o, rO) + auto smem_tiled_copy_Oaccum = make_tiled_copy_C(SmemTiledCopyO{}, tiled_mma_o); + auto smem_thr_copy_Oaccum = smem_tiled_copy_Oaccum.get_thread_slice(tidx); + Tensor taccOrOaccum = smem_thr_copy_Oaccum.retile_S(rO); // ((Atom,AtomNum), MMA_M, MMA_N) + Tensor taccOsOaccum = smem_thr_copy_Oaccum.partition_D(sOaccum); // ((Atom,AtomNum),PIPE_M,PIPE_N) + + + if constexpr (Split || Kernel_traits::Share_Q_K_smem) { flash::sync_threads(); } + + cute::copy(smem_tiled_copy_Oaccum, taccOrOaccum, taccOsOaccum); + + const index_t row_offset_o = binfo.q_offset(params.o_batch_stride, params.o_row_stride, bidb) + + m_block * kBlockM * params.o_row_stride + bidh * params.o_head_stride; + const index_t row_offset_oaccum = (((split_offset + n_split_idx) * params.h + bidh) * params.seqlen_q + + m_block * kBlockM) * params.d_v; + const index_t row_offset_lseaccum = Split ? ((split_offset + n_split_idx) * params.h + bidh) * params.seqlen_q + m_block * kBlockM : ((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + m_block * kBlockM; + + Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.oaccum_ptr : params.o_ptr) + (Split ? row_offset_oaccum : row_offset_o)), + Shape, Int>{}, + make_stride(Split ? kHeadDimV : params.o_row_stride, _1{})); + Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.softmax_lseaccum_ptr : params.softmax_lse_ptr) + row_offset_lseaccum), + Shape>{}, Stride<_1>{}); + // if (tidx == 0) { printf("row_offset_o = %d, bidh = %d, gOaccum = %p\n", row_offset_o, bidh, gOaccum.data()); } + + GmemTiledCopyO gmem_tiled_copy_Oaccum; + auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); + Tensor tOsOaccum = gmem_thr_copy_Oaccum.partition_S(sOaccum); // ((Atom,AtomNum),ATOM_M,ATOM_N) + Tensor tOgOaccum = gmem_thr_copy_Oaccum.partition_D(gOaccum); + + flash::sync_threads(); + + Tensor tOrOaccum = make_tensor(shape(tOgOaccum)); + cute::copy(gmem_tiled_copy_Oaccum, tOsOaccum, tOrOaccum); + + Tensor caccO = make_identity_tensor(Shape, Int>{}); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor taccOcO = thr_mma_o.partition_C(caccO); // (MMA,MMA_M,MMA_K) + static_assert(decltype(size<0>(taccOcO))::value == 4); + // Convert to ((2, 2), MMA_M, MMA_K) then take only the row indices. + Tensor taccOcO_row = logical_divide(taccOcO, Shape<_4>{})(make_coord(0, _), _, 0); + CUTE_STATIC_ASSERT_V(size(lse) == size(taccOcO_row)); // MMA_M + if (get<1>(taccOcO_row(0)) == 0) { + #pragma unroll + for (int mi = 0; mi < size(lse); ++mi) { + const int row = get<0>(taccOcO_row(mi)); + if (row < binfo.actual_seqlen_q - m_block * kBlockM) { gLSEaccum(row) = lse(mi); } + } + } + + // Construct identity layout for sO + Tensor cO = make_identity_tensor(make_shape(size<0>(sOaccum), size<1>(sOaccum))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + // Repeat the partitioning with identity layouts + Tensor tOcO = gmem_thr_copy_Oaccum.partition_D(cO); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + // Clear_OOB_K must be false since we don't want to write zeros to gmem + flash::copy_reg_to_global( + tOrOaccum, tOgOaccum, tOcO, params.d_v, binfo.actual_seqlen_q - m_block * kBlockM + ); +} + +template +__forceinline__ __device__ void compute_attn_1rowblock_splitkv_mla_k64_64x32_8waves_xcore1500(const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int n_block_min, const int n_block_max, const bool NoSplit) { + using Element = typename Kernel_traits::Element; + using ElementAccum = typename Kernel_traits::ElementAccum; + using index_t = typename Kernel_traits::index_t; + + // Shared memory. + extern __shared__ char smem_[]; + + // The thread index. + const int tidx = threadIdx.x; + const int warp_idx = tidx / 64; + const int lane_idx = tidx % 64; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kHeadDim = Kernel_traits::kHeadDim; + constexpr int kHeadDimV = Kernel_traits::kHeadDimV; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int kBlockKSmem = Kernel_traits::kBlockKSmem; + constexpr int kAtomLayoutMS = Kernel_traits::kAtomLayoutMS; + constexpr int kAtomLayoutMO = Kernel_traits::kAtomLayoutMO; + constexpr int Num_Stages = Kernel_traits::Num_Stages; + constexpr int kHeadDimNope = kHeadDimV; + constexpr int kHeadDimRope = kHeadDim - kHeadDimV; + + static_assert(kBlockKSmem == 64); + static_assert(Kernel_traits::Share_Q_K_smem && Kernel_traits::Is_Q_in_regs); + + + const BlockInfo binfo(params, bidb); + if (m_block * kBlockM >= binfo.actual_seqlen_q) return; + + // We iterate over the blocks in reverse order. This is because the last block is the only one + // that needs masking when we read K and V from global memory. Moreover, iterating in reverse + // might save us 1 register (we just need n_block instead of both n_block and n_block_max). + const index_t row_offset_q = binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb) + + m_block * kBlockM * params.q_row_stride + bidh * params.q_head_stride; + // We move K and V to the last block. + const int bidb_cache = params.cache_batch_idx == nullptr ? bidb : params.cache_batch_idx[bidb]; + const int *block_table = params.block_table == nullptr ? nullptr : params.block_table + bidb * params.block_table_batch_stride; + const int block_table_idx = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN / params.page_block_size; + const int block_table_offset = block_table == nullptr ? 0 : (n_block_max - 1) * kBlockN - block_table_idx * params.page_block_size; + const index_t row_offset_k = block_table == nullptr + ? binfo.k_offset(params.k_batch_stride, params.k_row_stride, bidb_cache) + + (n_block_max - 1) * kBlockN * params.k_row_stride + (bidh / params.h_h_k_ratio) * params.k_head_stride + : (bidh / params.h_h_k_ratio) * params.k_head_stride; + const index_t row_offset_v = block_table == nullptr + ? binfo.k_offset(params.v_batch_stride, params.v_row_stride, bidb_cache) + + (n_block_max - 1) * kBlockN * params.v_row_stride + (bidh / params.h_h_k_ratio) * params.v_head_stride + : (bidh / params.h_h_k_ratio) * params.v_head_stride; + + Tensor gQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + row_offset_q), + Shape, Int>{}, + make_stride(params.q_row_stride, _1{})); + Tensor gK = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_ptr) + row_offset_k), + Shape, Int>{}, + make_stride(params.k_row_stride, _1{})); + // if (threadIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { printf("k_ptr = %p, row_offset_k = %d, gK_ptr = %p\n", params.k_ptr, row_offset_k, gK.data()); } + Tensor gV = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_ptr) + row_offset_v), + Shape, Int>{}, + make_stride(params.v_row_stride, _1{})); + + Tensor sQ = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), + typename Kernel_traits::SmemLayoutQ{}); + Tensor sQ_NoSwizzle = make_tensor(sQ.data(), typename Kernel_traits::SmemLayoutQNoSwizzle{}); + Tensor sK = make_tensor(sQ.data(), typename Kernel_traits::SmemLayoutK242{}); + Tensor sK_NoSwizzle = make_tensor(sK.data(), typename Kernel_traits::SmemLayoutKNoswizzle{}); + Tensor sVt = make_tensor(sK.data(), typename Kernel_traits::SmemLayoutVtransposed242{}); + Tensor sVtNoSwizzle = make_tensor(sVt.data(), typename Kernel_traits::SmemLayoutVtransposedNoSwizzle{}); + Tensor sP = make_tensor(sK.data() + size(sK), typename Kernel_traits::SmemLayoutP{}); // sP is also bf16 + Tensor sRowMax = make_tensor(make_smem_ptr(reinterpret_cast(smem_) + size(sP) / 2 + size(sK) / 2), //sK is bf16, yet row_max is fp32 + typename Kernel_traits::SmemLayoutRowMax{}); + + const int swz333_offset = cute::get_swizzle_offset<8, 3, 3, 3>(tidx); + typename Kernel_traits::GmemTiledCopyB128 gmem_tiled_copy_Q; + auto gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(tidx); + Tensor tQgQ = gmem_thr_copy_Q.partition_S(gQ); + tQgQ = make_tensor(tQgQ.data() + swz333_offset, layout(tQgQ)); + Tensor tQsQ = gmem_thr_copy_Q.partition_D(sQ_NoSwizzle); + const int swz242_swap_offset = cute::get_swizzle_offset<8,2,4,2>(tidx) + ((tidx & 63) >= 32 ? ((tidx & 1) == 0 ? 8 : -8): 0); + + typename Kernel_traits::GmemTiledCopyB128 gmem_tiled_copy_KV; + int tidx_load_k = tidx & 0xFF; + auto gmem_thr_copy_KV = gmem_tiled_copy_KV.get_thread_slice(tidx_load_k); + Tensor tKgK = gmem_thr_copy_KV.partition_S(gK); // (KCPY, KCPY_N, KCPY_K) + tKgK = make_tensor(tKgK.data() + swz242_swap_offset, layout(tKgK)); + + Tensor tKsK = gmem_thr_copy_KV.partition_D(sK_NoSwizzle); + typename Kernel_traits::TiledMmaS_16x16x32_4x2 tiled_mma_s; + auto thr_mma_s = tiled_mma_s.get_thread_slice(tidx); + Tensor tSrQ = thr_mma_s.partition_fragment_A(sQ_NoSwizzle); // (MMA,MMA_M,MMA_K) + Tensor tSrK = thr_mma_s.partition_fragment_B(sK_NoSwizzle(_, _, 0)); // (MMA,MMA_N,MMA_K) + typename Kernel_traits::TiledMmaO tiled_mma_o; + auto thr_mma_o = tiled_mma_o.get_thread_slice(tidx); + Tensor tOrVt = thr_mma_o.partition_fragment_B(sVtNoSwizzle(_, _, 0)); // (MMA, MMA_K,MMA_N) + + auto smem_tiled_copy_S = make_tiled_copy_C(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_s); + auto smem_thr_copy_S = smem_tiled_copy_S.get_thread_slice(tidx); + Tensor tPsP = smem_thr_copy_S.partition_D(sP); // ((Atom,AtomNum),PIPE_M,PIPE_N) + + auto smem_tiled_copy_P = make_tiled_copy_A(typename Kernel_traits::UniversalCopyAtomB64{}, tiled_mma_o); + auto smem_thr_copy_P = smem_tiled_copy_P.get_thread_slice(tidx); + Tensor tOsP = smem_thr_copy_P.partition_S(sP); + Tensor tOrP = thr_mma_o.partition_fragment_A(sP); + + Tensor acc_o = partition_fragment_C(tiled_mma_o, Shape, Int>{}); // MMA, MMA_M, MMA_K + + // + // Copy Atom retiling + // + + auto smem_tiled_copy_Q = make_tiled_copy_A(typename Kernel_traits::UniversalCopyAtomB128{}, tiled_mma_s); + auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(tidx); + Tensor tSsQ = smem_thr_copy_Q.partition_S(sQ); + + const int swz242_diff_lds_b128 = (tidx & 7) >= 4 ? ((tidx & 31) < 16 ? 8 : -8) : 0; + auto smem_tiled_copy_K = make_tiled_copy_B(typename Kernel_traits::UniversalCopyAtomB128{}, tiled_mma_s); + auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx); + Tensor tSsK = smem_thr_copy_K.partition_S(sK); + tSsK = make_tensor(tSsK.data() + swz242_diff_lds_b128, layout(tSsK)); + const int swz242_diff_lds_trans_b128 = (tidx & 31) >= 16 ? ((tidx & 3) < 2 ? 8 : -8) : 0; + auto smem_tiled_copy_V = make_tiled_copy_B(typename Kernel_traits::LDSB64Trans4x16Atom{}, tiled_mma_o); + auto smem_thr_copy_V = smem_tiled_copy_V.get_thread_slice(tidx); + auto tOsVt = smem_thr_copy_V.partition_S(sVt); + tOsVt = make_tensor(tOsVt.data() + swz242_diff_lds_trans_b128, layout(tOsVt)); + + + // PREDICATES + + // Construct identity layout for sQ and sK + Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + + // Repeat the partitioning with identity layouts + Tensor tQcQ_noSwizzle = gmem_thr_copy_Q.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tKVcKV_noSwizzle = gmem_thr_copy_KV.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + Tensor tKVcKV = make_tensor(tKVcKV_noSwizzle.data() + make_coord(0, swz242_swap_offset), layout(tKVcKV_noSwizzle)); + Tensor tQcQ = make_tensor(tQcQ_noSwizzle.data() + make_coord(0, swz333_offset), layout(tQcQ_noSwizzle)); + + + // Prologue + + + // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs + flash::copy_b128_bsm_async(tQgQ, tQsQ, tQcQ, params.d, binfo.actual_seqlen_q - m_block * kBlockM); + flash::barrier_gvm<0>(); + cute::copy(smem_tiled_copy_Q, tSsQ, tSrQ); + flash::sync_threads(); + + + + int n_block = n_block_max - 1; + uint32_t Ksmem_read_index = 0; + uint32_t Ksmem_write_index = 0; + uint32_t page_idx[size<1>(tKgK)]; + uint32_t page_offset[size<1>(tKgK)]; + // We don't need to clear the sK smem tiles since we'll mask out the scores anyway. + if (warp_idx < 4) { + flash::copy_page(tKgK, page_idx, page_offset, n_block, block_table, params.page_block_size); + flash::copy_b128_page_bsm_async(gK, tKgK, tKsK(_, _, _, Ksmem_write_index), tKVcKV, params.d, n_block, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, page_idx, page_offset, swz242_swap_offset, binfo.actual_seqlen_k - n_block * kBlockN); + } + Ksmem_write_index ^= 1; + clear(acc_o); + + flash::Softmax(acc_o)> softmax; + + flash::Mask mask(binfo.actual_seqlen_k, binfo.actual_seqlen_q, params.ngroups, binfo.tot_seqlen_k, params.cp_world_size, params.cp_rank); + + for (; n_block >= n_block_min; --n_block) { + Tensor acc_s = partition_fragment_C(tiled_mma_s, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) + + clear(acc_s); + flash::barrier_gvm<0>(); + + + flash::gemm_prefetch_lds( + acc_s, tSrQ, tSrK, tSsQ, tSsK(_, _, _, Ksmem_read_index), tiled_mma_s, smem_tiled_copy_Q, smem_tiled_copy_K, + smem_thr_copy_Q, smem_thr_copy_K + ); + + /* + * NOTE:add this schedbound for separate gemmQK and gemmPV into two scheduling blocks + * bring opportunity to lds prefetching + */ + __builtin_mxc_schedbound_begin(); + + + + mask.template apply_mask( + acc_s, n_block * kBlockN + (tidx / 64) / 4 * 16, m_block * kBlockM + (tidx / 64) % kAtomLayoutMS * 16 + (tidx & 0xf), kAtomLayoutMS * 16 + ); + + Tensor scores_max_prev = make_fragment_like(softmax.row_max); + softmax.template get_row_max(acc_s, scores_max_prev, sRowMax, params.scale_softmax_log2); + if (n_block > n_block_min && warp_idx < 4) { + // Advance gK + flash::copy_page(tKgK, page_idx, page_offset, n_block - 1, block_table, params.page_block_size); + flash::copy_b128_page_bsm_async(gK, tKgK, tKsK(_, _, _, Ksmem_write_index), tKVcKV, params.d, n_block - 1, + block_table, params.k_batch_stride, params.k_row_stride, params.page_block_size, page_idx, page_offset, swz242_swap_offset); + Ksmem_write_index ^= 1; + } + softmax.template softmax_rescale_o_without_row_max(acc_s, acc_o, scores_max_prev, params.scale_softmax_log2); + + CONVERT_TENSOR_TYPE(ElementAccum, Element, acc_s, rP) + + cute::copy(smem_tiled_copy_S, rP, tPsP); + flash::sync_threads(); + + flash::gemm( + acc_o, tOrP, tOrVt, tOsP, tOsVt(_, _, _, Ksmem_read_index), tiled_mma_o, smem_tiled_copy_P, smem_tiled_copy_V, + smem_thr_copy_P, smem_thr_copy_V + ); + Ksmem_read_index ^= 1; + } + // Epilogue + if (NoSplit) { + store_64x32_xcore1500(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax, sRowMax); + }else{ + store_64x32_xcore1500(params, bidb, bidh, m_block, n_split_idx, smem_, acc_o, softmax, sRowMax); + } + +} + +} // namespace flash diff --git a/csrc/flash_run/run_mha_fwd.cpp b/csrc/flash_run/run_mha_fwd.cpp deleted file mode 100644 index 5691d8d..0000000 --- a/csrc/flash_run/run_mha_fwd.cpp +++ /dev/null @@ -1,12 +0,0 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - -#include -#include "run_mha.h" -#include "flash_fwd_dispatch_template.h" - -void run_mha_fwd(mcFlashAttn::Flash_fwd_mla_params ¶ms, cudaStream_t stream, bool force_split_kernel) { - - constexpr int kHeadDim = 576; - mcFlashAttn::run_mha_fwd_splitkv_dispatch(params, stream); -} - diff --git a/csrc/flash_run/run_mha.h b/csrc/flash_run/run_mla.h similarity index 52% rename from csrc/flash_run/run_mha.h rename to csrc/flash_run/run_mla.h index a2673b0..cee8df2 100644 --- a/csrc/flash_run/run_mha.h +++ b/csrc/flash_run/run_mla.h @@ -4,4 +4,5 @@ #include "flash_mla.h" -void run_mha_fwd(mcFlashAttn::Flash_fwd_mla_params ¶ms, cudaStream_t stream, bool force_split_kernel=false); +void run_mla_fwd(mcFlashAttn::Flash_fwd_mla_params ¶ms, cudaStream_t stream); +void run_mla_fwd(SparsePrefillParams ¶ms); diff --git a/csrc/flash_run/run_mla_fwd.cpp b/csrc/flash_run/run_mla_fwd.cpp new file mode 100644 index 0000000..b032ec5 --- /dev/null +++ b/csrc/flash_run/run_mla_fwd.cpp @@ -0,0 +1,24 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#include +#include +#include "run_mla.h" +#include "static_switch.h" +#include "flash_fwd_dispatch_template.h" + +void run_mla_fwd(mcFlashAttn::Flash_fwd_mla_params ¶ms, cudaStream_t stream) { + + constexpr int kHeadDim = 576; + ARCH_SWITCH(params.arch, kArch, [&] { + mcFlashAttn::run_mla_fwd_splitkv_dispatch(params, stream); + }); +} + + +void run_mla_fwd(SparsePrefillParams ¶ms) { + + constexpr int kHeadDim = 576; + ARCH_SWITCH(params.arch, kArch, [&] { + mcFlashAttn::run_flash_mla_sparse_prefill_dispatch(params, params.stream); + }); +} diff --git a/csrc/instantiations/flash_fwd_mla_metadata.cu b/csrc/instantiations/flash_fwd_mla_metadata.cu new file mode 100644 index 0000000..b7cef55 --- /dev/null +++ b/csrc/instantiations/flash_fwd_mla_metadata.cu @@ -0,0 +1,97 @@ +#pragma once + +#include "flash_mla.h" +#include "flash_dense_mla_decode_kernel.h" +#include "static_switch.h" + +__global__ void __launch_bounds__(32, 1, 1) +get_mla_metadata_kernel(const GetDecodingMetadataParams params) { + int *seqlens_k_ptr = params.seqlens_k_ptr; + int *tile_scheduler_metadata_ptr = params.tile_scheduler_metadata_ptr; + int *num_splits_ptr = params.num_splits_ptr; + int batch_size = params.batch_size; + int block_size_n = params.block_size_n; + int fixed_overhead_num_blocks = params.fixed_overhead_num_blocks; + int num_sm_parts = params.num_sm_parts; + + extern __shared__ int shared_mem[]; + int* num_blocks_shared = shared_mem; // [batch_size] + int* num_splits_shared = shared_mem + batch_size; // [batch_size+1] + int* seqlens_k_shared = shared_mem + batch_size*2+1; // [batch_size] + int* first_block_idx_shared = shared_mem + batch_size*3+1; // [batch_size] + int* last_block_idx_shared = shared_mem + batch_size*4+1; // [batch_size] + + int total_num_blocks = 0; + for (int i = threadIdx.x; i < batch_size; i += 32) { + int cur_s_k = params.topk == -1 ? __ldg(seqlens_k_ptr + i) : params.topk; + seqlens_k_shared[i] = cur_s_k; + int first_token_idx = 0; + int last_token_idx = max(cur_s_k-1, 0); + int cur_first_block_idx = first_token_idx / block_size_n; + int cur_last_block_idx = last_token_idx / block_size_n; + // NOTE Should attend to tokens [first_token_idx, last_token_idx], i.e. blocks [cur_first_block_idx, cur_last_block_idx] + // NOTE Before clamping, first_token_idx <= last_token_idx always holds, so after clamping, first_token_idx <= last_token_idx still holds. + // NOTE if seqlens_k is 0, then first_token_idx == last_token_idx == cur_first_block_idx == cur_last_block_idx == 0. So the sequence will have 1 block. We will correct this later in this kernel. + int num_blocks = cur_last_block_idx - cur_first_block_idx + 1; + total_num_blocks += num_blocks + fixed_overhead_num_blocks; + num_blocks_shared[i] = num_blocks; + first_block_idx_shared[i] = cur_first_block_idx; + last_block_idx_shared[i] = cur_last_block_idx; + } + for (int offset = 16; offset >= 1; offset /= 2) { + total_num_blocks += __shfl_xor_sync(uint32_t(-1), total_num_blocks, offset); + } + __syncwarp(); + + if (threadIdx.x == 0) { + int payload = mctlass::ceil_div(total_num_blocks, num_sm_parts) + fixed_overhead_num_blocks; + + int now_idx = 0, now_block = 0, now_n_split_idx = 0, cum_num_splits = 0; + num_splits_shared[0] = 0; + for (int i = 0; i < num_sm_parts; ++i) { + int tile_scheduler_metadata0[4], tile_scheduler_metadata1; + tile_scheduler_metadata0[0] = now_idx; + tile_scheduler_metadata0[1] = now_block + first_block_idx_shared[now_idx]; + tile_scheduler_metadata1 = now_n_split_idx; + int remain_payload = payload; + while (now_idx < batch_size) { + int num_blocks = num_blocks_shared[now_idx]; + int now_remain_blocks = num_blocks - now_block; + if (remain_payload >= now_remain_blocks + fixed_overhead_num_blocks) { + cum_num_splits += now_n_split_idx + 1; + num_splits_shared[now_idx + 1] = cum_num_splits; + remain_payload -= now_remain_blocks + fixed_overhead_num_blocks; + ++now_idx; + now_block = 0; + now_n_split_idx = 0; + } else { + if (remain_payload - fixed_overhead_num_blocks > 0) { + now_block += remain_payload - fixed_overhead_num_blocks; + ++now_n_split_idx; + remain_payload = 0; + } + break; + } + } + tile_scheduler_metadata0[2] = now_block > 0 ? now_idx : now_idx - 1; + tile_scheduler_metadata0[3] = now_block > 0 ? + now_block + first_block_idx_shared[now_idx] : (seqlens_k_shared[now_idx-1] == 0 ? + 0 : last_block_idx_shared[now_idx-1] + 1); + *reinterpret_cast(tile_scheduler_metadata_ptr + i * TileSchedulerMetaDataSize) = *reinterpret_cast(tile_scheduler_metadata0); + tile_scheduler_metadata_ptr[i * TileSchedulerMetaDataSize + 4] = tile_scheduler_metadata1; + } + FLASH_DEVICE_ASSERT(now_idx == batch_size && now_block == 0 && now_n_split_idx == 0); + } + __syncwarp(); + + for (int i = threadIdx.x; i <= batch_size; i += 32) { + num_splits_ptr[i] = num_splits_shared[i]; + } +} + +void run_get_mla_metadata_kernel(GetDecodingMetadataParams ¶ms, cudaStream_t stream) { + int smem_size = sizeof(int) * (params.batch_size*5+1); + CUDA_CHECK(cudaFuncSetAttribute(get_mla_metadata_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + get_mla_metadata_kernel<<<1, 32, smem_size, stream>>>(params); + CUDA_KERNEL_LAUNCH_CHECK(); +} \ No newline at end of file diff --git a/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_bf16_True_True_1stage_xcore1000.cu b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_bf16_True_True_1stage_xcore1000.cu new file mode 100644 index 0000000..581de3a --- /dev/null +++ b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_bf16_True_True_1stage_xcore1000.cu @@ -0,0 +1,18 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#include "flash_mla.h" +#include "flash_fwd_run_template.h" +#include + +template void run_flash_splitkv_fwd_mla_template< + 576, + 16, + 16, + 4, + true, + true, + mctlass::bfloat16_t, + true, + 512, + 1 + >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); diff --git a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_bf16_True_True_split_sm80.cu b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_bf16_True_True_split_xcore1000.cu similarity index 76% rename from csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_bf16_True_True_split_sm80.cu rename to csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_bf16_True_True_split_xcore1000.cu index 28da785..20f3f82 100644 --- a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_bf16_True_True_split_sm80.cu +++ b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_bf16_True_True_split_xcore1000.cu @@ -1,17 +1,17 @@ // Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) #include "flash_mla.h" -#include "flash_run_fwd_template_impl.h" +#include "flash_fwd_run_template.h" #include -template void run_flash_splitkv_fwd_template< +template void run_flash_splitkv_fwd_mla_template< 576, 16, 16, 4, true, true, - cutlass::bfloat16_t, + mctlass::bfloat16_t, true, 512, 2 diff --git a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_fp16_True_True_sm80.cu b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_fp16_True_True_1stage_xcore1000.cu similarity index 69% rename from csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_fp16_True_True_sm80.cu rename to csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_fp16_True_True_1stage_xcore1000.cu index 03236fc..f9e62cc 100644 --- a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_fp16_True_True_sm80.cu +++ b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_fp16_True_True_1stage_xcore1000.cu @@ -1,18 +1,18 @@ // Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) #include "flash_mla.h" -#include "flash_run_fwd_template_impl.h" +#include "flash_fwd_run_template.h" #include -template void run_flash_splitkv_fwd_template< +template void run_flash_splitkv_fwd_mla_template< 576, 16, 16, 4, true, true, - cutlass::half_t, - false, + mctlass::half_t, + true, 512, - 2 + 1 >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); diff --git a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_fp16_True_True_split_sm80.cu b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_fp16_True_True_split_xcore1000.cu similarity index 77% rename from csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_fp16_True_True_split_sm80.cu rename to csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_fp16_True_True_split_xcore1000.cu index 5cf5c65..e061636 100644 --- a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_fp16_True_True_split_sm80.cu +++ b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_fp16_True_True_split_xcore1000.cu @@ -1,17 +1,17 @@ // Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) #include "flash_mla.h" -#include "flash_run_fwd_template_impl.h" +#include "flash_fwd_run_template.h" #include -template void run_flash_splitkv_fwd_template< +template void run_flash_splitkv_fwd_mla_template< 576, 16, 16, 4, true, true, - cutlass::half_t, + mctlass::half_t, true, 512, 2 diff --git a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_bf16_True_True_split_sm80.cu b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m32n16_bf16_True_True_split_xcore1000.cu similarity index 76% rename from csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_bf16_True_True_split_sm80.cu rename to csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m32n16_bf16_True_True_split_xcore1000.cu index 835e094..abeeb13 100644 --- a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_bf16_True_True_split_sm80.cu +++ b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m32n16_bf16_True_True_split_xcore1000.cu @@ -1,17 +1,17 @@ // Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) #include "flash_mla.h" -#include "flash_run_fwd_template_impl.h" +#include "flash_fwd_run_template.h" #include -template void run_flash_splitkv_fwd_template< +template void run_flash_splitkv_fwd_mla_template< 576, 32, 16, 4, true, true, - cutlass::bfloat16_t, + mctlass::bfloat16_t, true, 512, 2 diff --git a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_fp16_True_True_split_sm80.cu b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m32n16_fp16_True_True_split_xcore1000.cu similarity index 77% rename from csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_fp16_True_True_split_sm80.cu rename to csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m32n16_fp16_True_True_split_xcore1000.cu index 6a68272..2c9ed6e 100644 --- a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_fp16_True_True_split_sm80.cu +++ b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m32n16_fp16_True_True_split_xcore1000.cu @@ -1,17 +1,17 @@ // Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) #include "flash_mla.h" -#include "flash_run_fwd_template_impl.h" +#include "flash_fwd_run_template.h" #include -template void run_flash_splitkv_fwd_template< +template void run_flash_splitkv_fwd_mla_template< 576, 32, 16, 4, true, true, - cutlass::half_t, + mctlass::half_t, true, 512, 2 diff --git a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_bf16_True_True_split_sm80.cu b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu similarity index 76% rename from csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_bf16_True_True_split_sm80.cu rename to csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu index 052518c..8210f0f 100644 --- a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_bf16_True_True_split_sm80.cu +++ b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu @@ -1,17 +1,17 @@ // Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) #include "flash_mla.h" -#include "flash_run_fwd_template_impl.h" +#include "flash_fwd_run_template.h" #include -template void run_flash_splitkv_fwd_template< +template void run_flash_splitkv_fwd_mla_template< 576, 64, 16, 8, true, true, - cutlass::bfloat16_t, + mctlass::bfloat16_t, true, 512, 2 diff --git a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_fp16_True_True_split_sm80.cu b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1000.cu similarity index 77% rename from csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_fp16_True_True_split_sm80.cu rename to csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1000.cu index e4f33b5..c68b3c4 100644 --- a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_fp16_True_True_split_sm80.cu +++ b/csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1000.cu @@ -1,17 +1,17 @@ // Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) #include "flash_mla.h" -#include "flash_run_fwd_template_impl.h" +#include "flash_fwd_run_template.h" #include -template void run_flash_splitkv_fwd_template< +template void run_flash_splitkv_fwd_mla_template< 576, 64, 16, 8, true, true, - cutlass::half_t, + mctlass::half_t, true, 512, 2 diff --git a/csrc/instantiations/xcore1000/flash_fwd_prefill_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu b/csrc/instantiations/xcore1000/flash_fwd_prefill_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu new file mode 100644 index 0000000..a23c9e4 --- /dev/null +++ b/csrc/instantiations/xcore1000/flash_fwd_prefill_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu @@ -0,0 +1,18 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#include "flash_mla.h" +#include "flash_fwd_run_template.h" +#include + +template void run_flash_mla_sparse_prefill_template< + 576, + 64, + 16, + 8, + true, + true, + mctlass::bfloat16_t, + true, + 512, + 2 + >(SparsePrefillParams ¶ms, cudaStream_t stream); diff --git a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_fp16_True_True_sm80.cu b/csrc/instantiations/xcore1000/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu similarity index 71% rename from csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_fp16_True_True_sm80.cu rename to csrc/instantiations/xcore1000/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu index 1e9fc7a..519daae 100644 --- a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_fp16_True_True_sm80.cu +++ b/csrc/instantiations/xcore1000/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu @@ -1,18 +1,18 @@ // Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) #include "flash_mla.h" -#include "flash_run_fwd_template_impl.h" +#include "flash_fwd_run_template.h" #include -template void run_flash_splitkv_fwd_template< +template void run_flash_splitkv_fwd_sparse_mla_template< 576, 64, 16, 8, true, true, - cutlass::half_t, - false, + mctlass::bfloat16_t, + true, 512, 2 - >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); + >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); \ No newline at end of file diff --git a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_bf16_True_True_sm80.cu b/csrc/instantiations/xcore1000/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1000.cu similarity index 72% rename from csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_bf16_True_True_sm80.cu rename to csrc/instantiations/xcore1000/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1000.cu index b6ae743..6cdbb6e 100644 --- a/csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_bf16_True_True_sm80.cu +++ b/csrc/instantiations/xcore1000/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1000.cu @@ -1,18 +1,18 @@ // Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) #include "flash_mla.h" -#include "flash_run_fwd_template_impl.h" +#include "flash_fwd_run_template.h" #include -template void run_flash_splitkv_fwd_template< +template void run_flash_splitkv_fwd_sparse_mla_template< 576, 64, 16, 8, true, true, - cutlass::bfloat16_t, - false, + mctlass::half_t, + true, 512, 2 - >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); + >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); \ No newline at end of file diff --git a/csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu b/csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu new file mode 100644 index 0000000..234c933 --- /dev/null +++ b/csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu @@ -0,0 +1,19 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#include "flash_mla.h" +#include "flash_fwd_run_template.h" +#include + +template void run_flash_splitkv_fwd_mla_template< + 576, + 64, + 16, + 8, + true, + true, + mctlass::bfloat16_t, + true, + 512, + 2, + Arch::xcore1500 + >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); diff --git a/csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1500.cu b/csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1500.cu new file mode 100644 index 0000000..21b1eca --- /dev/null +++ b/csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1500.cu @@ -0,0 +1,19 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#include "flash_mla.h" +#include "flash_fwd_run_template.h" +#include + +template void run_flash_splitkv_fwd_mla_template< + 576, + 64, + 16, + 8, + true, + true, + mctlass::half_t, + true, + 512, + 2, + Arch::xcore1500 + >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); diff --git a/csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n32_bf16_True_True_split_xcore1500.cu b/csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n32_bf16_True_True_split_xcore1500.cu new file mode 100644 index 0000000..cba61b6 --- /dev/null +++ b/csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n32_bf16_True_True_split_xcore1500.cu @@ -0,0 +1,19 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#include "flash_mla.h" +#include "flash_fwd_run_template.h" +#include + +template void run_flash_splitkv_fwd_mla_template< + 576, + 64, + 32, + 8, + true, + true, + mctlass::bfloat16_t, + true, + 512, + 2, + Arch::xcore1500 + >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); diff --git a/csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n32_fp16_True_True_split_xcore1500.cu b/csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n32_fp16_True_True_split_xcore1500.cu new file mode 100644 index 0000000..da7492b --- /dev/null +++ b/csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n32_fp16_True_True_split_xcore1500.cu @@ -0,0 +1,19 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#include "flash_mla.h" +#include "flash_fwd_run_template.h" +#include + +template void run_flash_splitkv_fwd_mla_template< + 576, + 64, + 32, + 8, + true, + true, + mctlass::half_t, + true, + 512, + 2, + Arch::xcore1500 + >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); diff --git a/csrc/instantiations/xcore1500/flash_fwd_prefill_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu b/csrc/instantiations/xcore1500/flash_fwd_prefill_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu new file mode 100644 index 0000000..39d3b0a --- /dev/null +++ b/csrc/instantiations/xcore1500/flash_fwd_prefill_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu @@ -0,0 +1,19 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#include "flash_mla.h" +#include "flash_fwd_run_template.h" +#include + +template void run_flash_mla_sparse_prefill_template< + 576, + 64, + 16, + 8, + true, + true, + mctlass::bfloat16_t, + true, + 512, + 2, + Arch::xcore1500 + >(SparsePrefillParams ¶ms, cudaStream_t stream); diff --git a/csrc/instantiations/xcore1500/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu b/csrc/instantiations/xcore1500/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu new file mode 100644 index 0000000..cb17ccf --- /dev/null +++ b/csrc/instantiations/xcore1500/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu @@ -0,0 +1,19 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#include "flash_mla.h" +#include "flash_fwd_run_template.h" +#include + +template void run_flash_splitkv_fwd_sparse_mla_template< + 576, + 64, + 16, + 8, + true, + true, + mctlass::bfloat16_t, + true, + 512, + 2, + Arch::xcore1500 + >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); \ No newline at end of file diff --git a/csrc/instantiations/xcore1500/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1500.cu b/csrc/instantiations/xcore1500/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1500.cu new file mode 100644 index 0000000..899c331 --- /dev/null +++ b/csrc/instantiations/xcore1500/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1500.cu @@ -0,0 +1,19 @@ +// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) + +#include "flash_mla.h" +#include "flash_fwd_run_template.h" +#include + +template void run_flash_splitkv_fwd_sparse_mla_template< + 576, + 64, + 16, + 8, + true, + true, + mctlass::half_t, + true, + 512, + 2, + Arch::xcore1500 + >(Flash_fwd_mla_params ¶ms, cudaStream_t stream); \ No newline at end of file diff --git a/csrc/mctlass/.gitignore b/csrc/mctlass/.gitignore new file mode 100644 index 0000000..ef9dc32 --- /dev/null +++ b/csrc/mctlass/.gitignore @@ -0,0 +1,42 @@ +# Compiled Object files +*.slo +*.lo +*.o +*.obj + +# Precompiled Headers +*.gch +*.pch + +# Compiled Dynamic libraries +*.so +*.dylib +*.dll + +# Fortran module files +*.mod + +# Compiled Static libraries +*.lai +*.la +*.a +*.lib + +# Executables +*.exe +*.out +*.app + +# vim tags +tags +.tags +.*.swp + +# Visual Studio Code +.vscode + +# install.sh build dir +build + +# PyCache files +__pycache__ diff --git a/csrc/mctlass/include/cute/algorithm/axpby.hpp b/csrc/mctlass/include/cute/algorithm/axpby.hpp new file mode 100644 index 0000000..a613417 --- /dev/null +++ b/csrc/mctlass/include/cute/algorithm/axpby.hpp @@ -0,0 +1,79 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +namespace cute +{ + +// +// Accept mutable temporaries +// +template +CUTE_HOST_DEVICE +void +axpby(Alpha const& alpha, + Tensor const& x, + Beta const& beta, + Tensor && y) +{ + return axpby(alpha, x, beta, y); +} + +// +// AXPBY +// +template +CUTE_HOST_DEVICE +void +axpby(Alpha const& alpha, + Tensor const& x, + Beta const& beta, + Tensor & y) +{ + auto isBetaZero = (beta == Int<0>{}); + + CUTE_UNROLL + for (int i = 0; i < size(x); ++i) { + y(i) = (isBetaZero ? alpha * x(i) : alpha * x(i) + beta * y(i)); + } +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/algorithm/clear.hpp b/csrc/mctlass/include/cute/algorithm/clear.hpp new file mode 100644 index 0000000..ce7b510 --- /dev/null +++ b/csrc/mctlass/include/cute/algorithm/clear.hpp @@ -0,0 +1,66 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +#include + +namespace cute +{ + +// +// Accept mutable temporaries +// +template +CUTE_HOST_DEVICE +void +clear(Tensor&& tensor) +{ + return clear(tensor); +} + +// +// Set elements to zero +// +template +CUTE_HOST_DEVICE +void +clear(Tensor& tensor) +{ + using T = typename Tensor::value_type; + + fill(tensor, T{}); +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/algorithm/copy.hpp b/csrc/mctlass/include/cute/algorithm/copy.hpp new file mode 100644 index 0000000..a3f08e9 --- /dev/null +++ b/csrc/mctlass/include/cute/algorithm/copy.hpp @@ -0,0 +1,523 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include + +#include + +namespace cute +{ + +// +// Accept mutable temporaries +// + +template +CUTE_HOST_DEVICE +void +copy_if(PrdTensor const& pred, + Tensor const& src, + Tensor && dst) +{ + return copy_if(pred, src, dst); +} + +template +CUTE_HOST_DEVICE +void +copy_if(Copy_Atom const& copy_atom, + PrdTensor const& pred, + Tensor const& src, + Tensor && dst) +{ + return copy_if(copy_atom, pred, src, dst); +} + +template +CUTE_HOST_DEVICE +void +copy_vec(Tensor const& src, + Tensor && dst) +{ + return copy_vec(src, dst); +} + +template +CUTE_HOST_DEVICE +void +copy(Tensor const& src, + Tensor && dst) +{ + return copy(src, dst); +} + +template +CUTE_HOST_DEVICE +void +copy(Copy_Atom const& copy_atom, + Tensor const& src, + Tensor && dst) +{ + return copy(copy_atom, src, dst); +} + +// +// copy_if -- Predicated Copy +// + +template +CUTE_HOST_DEVICE +void +copy_if(PrdTensor const& pred, + Tensor const& src, + Tensor & dst) +{ + auto copy_op = select_elementwise_copy(src, dst); + + CUTE_UNROLL + for (int i = 0; i < size(src); ++i) { + if (pred(i)) { + copy_op.copy(src(i), dst(i)); + } + } +} + +// +// copy_if -- Predicated CopyAtom +// + +template +CUTE_HOST_DEVICE +void +copy_if(Copy_Atom const& copy_atom, + PredTensor const& pred, // (Rest...) + Tensor const& src, // (V,Rest...) + Tensor & dst) // (V,Rest...) +{ + static_assert(SrcLayout::rank == DstLayout::rank, "CopyAtom rank-mismatch."); + if constexpr (SrcLayout::rank == 1) { // Dispatch the copy + copy_atom.call(src, dst); + } else { // Loop over all but the first mode + constexpr int R = SrcLayout::rank; + auto src_v = group_modes<1,R>(src); + auto dst_v = group_modes<1,R>(dst); + CUTE_UNROLL + for (int i = 0; i < size<1>(src_v); ++i) { + if (pred(i)) { + copy_atom.call(src_v(_,i), dst_v(_,i)); + } + } + } +} + +// +// copy_vec -- attempt vectorized copy with VecType +// + +template +CUTE_HOST_DEVICE +void +copy_vec(Tensor const& src, + Tensor & dst) +{ + using SrcType = typename SrcEngine::value_type; + using DstType = typename DstEngine::value_type; + if constexpr (sizeof(SrcType) == sizeof(DstType) && sizeof(VecType) > sizeof(DstType)) + { + /* @pre is_aligned(src.data()) && + * is_aligned(dst.data()) + */ + auto src_v = recast(src); + auto dst_v = recast(dst); + +#if 0 + if (thread0()) { + print("copy_vec -- vectorizing copy from %3db to %3db\n", int(8*sizeof(SrcType)), int(8*sizeof(VecType))); + print(" "); print(layout(src)); print(" => "); print(layout(src_v)); print("\n"); + print(" "); print(layout(dst)); print(" => "); print(layout(dst_v)); print("\n"); + } +#endif + + return copy_if(TrivialPredTensor{}, src_v, dst_v); + } else { +#if 0 + if (thread0()) { + print("copy_vec -- not vectorizing, copy with %3db and %3db\n", int(8*sizeof(SrcType)), int(8*sizeof(DstType))); + print(" "); print(layout(src)); print("\n"); + print(" "); print(layout(dst)); print("\n"); + } +#endif + + return copy_if(TrivialPredTensor{}, src, dst); + } +} + +// +// copy -- auto-vectorizing copy +// + +template +CUTE_HOST_DEVICE +void +copy(Tensor const& src, + Tensor & dst) +{ + constexpr int N = decltype(max_common_vector(src, dst))::value; + +#if 0 + if (thread0()) { + print("copy -- found a max_common_vector of %d\n", N); + print(" "); print(src.data()); print(" o "); print(layout(src)); print("\n"); + print(" "); print(dst.data()); print(" o "); print(layout(dst)); print("\n"); + } +#endif + + if constexpr (N <= 1) { + return copy_if(TrivialPredTensor{}, src, dst); + } else { + constexpr int vec_bits = N * sizeof_bits::value; + using VecType = uint_bit_t; + return copy_vec(src, dst); + } +} + +// +// copy -- CopyAtom +// + +template +CUTE_HOST_DEVICE +void +copy(Copy_Atom const& copy_atom, + Tensor const& src, + Tensor & dst) +{ + return copy_if(copy_atom, TrivialPredTensor{}, src, dst); +} + +template +CUTE_HOST_DEVICE +void +copy(Copy_Atom const&, + Tensor const& src, + Tensor & dst) +{ + return copy(src, dst); +} + +#if defined(__MERGE_LDS_B32) + +CUTE_DEVICE +void reg_trans(uint32_t &a) { + + /* ************************************************************ + ** tmp_0[n]=a[(n&0x3c)+shfl[n%4]] + ** 0x0b1 means shfl[0]=1, shfl[1]=0, shfl[2]=3, shfl[3]=2 + ** and tmp_0[n]/a[n] means the value of tmp_0/a while lane_id=n + * ************************************************************/ + auto tmp_0 = __builtin_mxc_mov_raw_shfl(a, 0x0b1, 0xf, 0xf, false); + auto tmp_1 = __builtin_mxc_byte_perm(a, tmp_0, 0x07060302); + a = __builtin_mxc_byte_perm(tmp_0, a, 0x05040100); + + if (__lane_id() & 0x1) { + a = tmp_1; + } + +} + +CUTE_DEVICE +void reg_trans(uint32_t &a, uint32_t &b) { + + reg_trans(a); + reg_trans(b); + +} + +template < + class SrcEngine, class SrcLayout, + class DstEngine, class DstLayout> +CUTE_DEVICE +void copy_trans(Tensor const && src, + Tensor && dst, + const uint32_t &src_stride, + const uint32_t &dst_stride, + const uint32_t *cpy_offset) +{ + + auto dst_ptr = reinterpret_cast(dst.data()); + auto src_addr = reinterpret_cast(src.data().ptr_); + src_addr = src_addr - cpy_offset[8]; + + /* ************************************************* + ** The address attribute of src_addr has benn destoried, + ** So we need to use __attribute__((address_space (3))) + * *************************************************/ + uint32_t __attribute__((address_space(3))) *src_ptr[8]; + CUTE_UNROLL + for (uint32_t i = 0; i < 4; ++i) { + src_ptr[2 * i] = (uint32_t __attribute__((address_space(3))) *)(src_addr) + cpy_offset[2 * i]; + src_ptr[2 * i + 1] = (uint32_t __attribute__((address_space(3))) *)(src_addr) + cpy_offset[2 * i + 1]; + } + + CUTE_UNROLL + for (uint32_t i = 0; i < size(dst) / 4; ++i) { + dst_ptr[i] = src_ptr[i][0]; + } + + CUTE_UNROLL + for (uint32_t i = 0; i < 8; ++i) { + src_ptr[i] = src_ptr[i] + src_stride; + } + dst_ptr = dst_ptr + dst_stride; + + CUTE_UNROLL + for (uint32_t i = 0; i < size(dst) / 2 - size(dst) / 4; ++i) { + dst_ptr[i] = src_ptr[i][0]; + } + +} + + +#elif defined(__MERGE_LDS_B64) + +CUTE_DEVICE +void reg_trans(uint32_t &a, uint32_t &b) { + + const int laneId = __lane_id(); + + /* ************************************************************ + ** tmp_0[n]=a[(n&0x3c)+shfl[n%4]] + ** 0x0b1 means shfl[0]=1, shfl[1]=0, shfl[2]=3, shfl[3]=2 + ** and tmp_0[n]/a[n] means the value of tmp_0/a while lane_id=n + * ************************************************************/ + auto tmp_0 = __builtin_mxc_mov_raw_shfl(a, 0x0b1, 0xf, 0xf, false); + auto tmp_1 = __builtin_mxc_byte_perm(a, tmp_0, 0x07060302); + a = __builtin_mxc_byte_perm(tmp_0, a, 0x05040100); + + auto tmp_2 = __builtin_mxc_mov_raw_shfl(b, 0x0b1, 0xf, 0xf, false); + auto tmp_3 = __builtin_mxc_byte_perm(b, tmp_2, 0x07060302); + b = __builtin_mxc_byte_perm(tmp_2, b, 0x05040100); + + if (laneId & 0x1) { + a = tmp_1; + b = tmp_3; + } + + /* ************************************************************ + ** tmp_0[n]=a[(n&0x3c)+shfl[n%4]] + ** 0x04e means shfl[0]=2, shfl[1]=3, shfl[2]=0, shfl[3]=1 + ** and tmp_0[n]/a[n] means the value of tmp_0/a while lane_id=n + * ************************************************************/ + tmp_0 = __builtin_mxc_mov_raw_shfl(a, 0x04e, 0xf, 0xf, false); + tmp_1 = __builtin_mxc_mov_raw_shfl(b, 0x04e, 0xf, 0xf, false); + + if ((laneId & 0x3) >> 1) { + a = tmp_1; + } + else { + b = tmp_0; + } +} + +template < + class SrcEngine, class SrcLayout, + class DstEngine, class DstLayout> +CUTE_DEVICE +void +copy_trans( + Tensor const&& src, + Tensor && dst, + const int src_stride, + const int dst_stride, + const uint32_t *cpy_offset) +{ + + auto dst_ptr = reinterpret_cast(dst.data()); + auto src_addr = reinterpret_cast(src.data().ptr_); + src_addr = src_addr - cpy_offset[4]; + + /* ************************************************ + ** The address attribute of src_addr has benn destoried + ** So we need to use __attribute__((address_space (3))) + * *************************************************/ + uint32_t __attribute__((address_space (3))) *src_ptr[4]; + CUTE_UNROLL + for (int i = 0; i < 4; ++i) { + src_ptr[i] = (uint32_t __attribute__((address_space(3))) *)(src_addr) + cpy_offset[i]; + } + CUTE_UNROLL + for (int i = 0; i < size(dst) / 8; ++i) { + dst_ptr[2 * i] = src_ptr[i][0]; + dst_ptr[2 * i + 1] = src_ptr[i][1]; + } + + CUTE_UNROLL + for (int i = 0; i < 4; ++i) { + src_ptr[i] = src_ptr[i] + src_stride; + } + + dst_ptr = dst_ptr + dst_stride; + CUTE_UNROLL + for (int i = 0; i < size(dst) / 4 - size(dst) / 8; ++i) { + dst_ptr[2 * i] = src_ptr[i][0]; + dst_ptr[2 * i + 1] = src_ptr[i][1]; + } + +} + +#endif + +#if defined(__MERGE_LDS_B32) || defined(__MERGE_LDS_B64) +template +CUTE_DEVICE +void tensor_trans(Tensor && dst, + const uint32_t stride) { + + auto dst_ptr = reinterpret_cast(dst.data()); + CUTE_UNROLL + for (uint32_t i = 0; i < size(dst) / 8; ++i) { + reg_trans(dst_ptr[2 * i], dst_ptr[2 * i + 1]); + } + + dst_ptr = dst_ptr + stride; + CUTE_UNROLL + for (uint32_t i = 0; i < size(dst) / 4 - size(dst) / 8; ++i) { + reg_trans(dst_ptr[2 * i], dst_ptr[2 * i + 1]); + } + +} +#endif + +template +CUTE_HOST_DEVICE +void +copy_global_to_reg( + Tensor const&& src, + uint32_t *dst) +{ + + typedef __NATIVE_VECTOR__(4, int) VecType; + auto src_ptr = (VecType *)(src.data().ptr_); + auto dst_ptr = (VecType *)(dst); + dst_ptr[0] = __builtin_mxc_load_global_async128(src_ptr); + +} + +template +CUTE_HOST_DEVICE +void +copy_reg_to_share( + uint32_t *src_ptr, + Tensor && dst) +{ + + auto dst_ptr = reinterpret_cast(dst.data().ptr_); + dst_ptr[0] = src_ptr[0]; + dst_ptr[1] = src_ptr[1]; + dst_ptr[2] = src_ptr[2]; + dst_ptr[3] = src_ptr[3]; + +} + +////////////////////////////////////////// +// Special Auto-Vectorizing Overloads +////////////////////////////////////////// + +#if defined(CUTE_COPY_ATOM_TMA_SM90_ENABLED) +template +CUTE_HOST_DEVICE +void +copy(Copy_Atom, CA_Args...> const& atom, + Tensor const& src, + Tensor & dst) +{ + using SrcType = typename SrcEngine::value_type; + using DstType = typename DstEngine::value_type; + static_assert(sizeof_bits::value == sizeof_bits::value); + static_assert((is_gmem::value && is_smem::value) || + (is_smem::value && is_gmem::value), + "Bulk Copy only supports gmem -> smem or smem -> gmem movement."); + // Do BulkCopy dispatch + using BULK_COPY_OP = conditional_t::value, + SM90_BULK_COPY_G2S, + SM90_BULK_COPY_S2G>; + + constexpr int N = decltype(max_common_vector(src, dst))::value; + + // Construct a new concrete Atom of the vector size + using N_BITS = Int::value>; + using COPY_ATOM = Copy_Atom, SrcType>; + auto bulk_atom = apply(atom.opargs_, [&](auto const&... args) { return COPY_ATOM{args...}; }); + + // Tile the src and dst to the Atom + auto tiler = right_inverse(dst.layout()).compose(Int{}); + +#if 0 + if (thread0()) { + print("copy -- found a max_common_vector of %d\n", N); + print(" "); print(src.data()); print(" o "); print(layout(src)); print("\n"); + print(" "); print(dst.data()); print(" o "); print(layout(dst)); print("\n"); + } +#endif + + return copy(bulk_atom, logical_divide(src, tiler), logical_divide(dst, tiler)); +} +#endif // #if defined(CUTE_COPY_ATOM_TMA_SM90_ENABLED) + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/algorithm/fill.hpp b/csrc/mctlass/include/cute/algorithm/fill.hpp new file mode 100644 index 0000000..bc0c4ad --- /dev/null +++ b/csrc/mctlass/include/cute/algorithm/fill.hpp @@ -0,0 +1,87 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include + +namespace cute +{ + +// +// Accept mutable temporaries +// +template +CUTE_HOST_DEVICE +void +fill(Tensor&& tensor, T const& value) +{ + return fill(tensor, value); +} + +namespace detail +{ + +// Prefer fill(tensor.data(), value), if possible +template +CUTE_HOST_DEVICE +auto +fill(Tensor& tensor, T const& value, prefer<1>) + -> decltype(fill(tensor.data(), value)) +{ + fill(tensor.data(), value); +} + +// Default implementation +template +CUTE_HOST_DEVICE +void +fill(Tensor& tensor, T const& value, prefer<0>) +{ + CUTE_UNROLL + for (int i = 0; i < size(tensor); ++i) { + tensor(i) = value; + } +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE +void +fill(Tensor& tensor, T const& value) +{ + return detail::fill(tensor, value, prefer<1>{}); +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/algorithm/functional.hpp b/csrc/mctlass/include/cute/algorithm/functional.hpp new file mode 100644 index 0000000..ea17ecb --- /dev/null +++ b/csrc/mctlass/include/cute/algorithm/functional.hpp @@ -0,0 +1,198 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +/** C++14 extensions */ + +namespace cute { + +/**************/ +/** Identity **/ +/**************/ + +struct identity { + template + CUTE_HOST_DEVICE constexpr + decltype(auto) operator()(T&& arg) const { + return std::forward(arg); + } +}; + +template +struct constant_fn { + template + CUTE_HOST_DEVICE constexpr + decltype(auto) operator()(T&&...) const { + return r_; + } + R r_; +}; + +/***********/ +/** Unary **/ +/***********/ + +#define CUTE_LEFT_UNARY_OP(NAME,OP) \ + struct NAME { \ + template \ + CUTE_HOST_DEVICE constexpr \ + decltype(auto) operator()(T&& arg) const { \ + return OP std::forward(arg); \ + } \ + } +#define CUTE_RIGHT_UNARY_OP(NAME,OP) \ + struct NAME { \ + template \ + CUTE_HOST_DEVICE constexpr \ + decltype(auto) operator()(T&& arg) const { \ + return std::forward(arg) OP ; \ + } \ + } +#define CUTE_NAMED_UNARY_OP(NAME,OP) \ + struct NAME { \ + template \ + CUTE_HOST_DEVICE constexpr \ + decltype(auto) operator()(T&& arg) const { \ + return OP (std::forward(arg)); \ + } \ + } + +CUTE_LEFT_UNARY_OP(unary_plus, +); +CUTE_LEFT_UNARY_OP(negate, -); +CUTE_LEFT_UNARY_OP(bit_not, ~); +CUTE_LEFT_UNARY_OP(logical_not, !); +CUTE_LEFT_UNARY_OP(dereference, *); +CUTE_LEFT_UNARY_OP(address_of, &); +CUTE_LEFT_UNARY_OP(pre_increment, ++); +CUTE_LEFT_UNARY_OP(pre_decrement, --); + +CUTE_RIGHT_UNARY_OP(post_increment, ++); +CUTE_RIGHT_UNARY_OP(post_decrement, --); + +CUTE_NAMED_UNARY_OP(abs_fn, abs); +CUTE_NAMED_UNARY_OP(conjugate, cute::conj); + +#undef CUTE_LEFT_UNARY_OP +#undef CUTE_RIGHT_UNARY_OP +#undef CUTE_NAMED_UNARY_OP + +/************/ +/** Binary **/ +/************/ + +#define CUTE_BINARY_OP(NAME,OP) \ + struct NAME { \ + template \ + CUTE_HOST_DEVICE constexpr \ + decltype(auto) operator()(T&& lhs, U&& rhs) const { \ + return std::forward(lhs) OP std::forward(rhs); \ + } \ + } +#define CUTE_NAMED_BINARY_OP(NAME,OP) \ + struct NAME { \ + template \ + CUTE_HOST_DEVICE constexpr \ + decltype(auto) operator()(T&& lhs, U&& rhs) const { \ + return OP (std::forward(lhs), std::forward(rhs)); \ + } \ + } + + +CUTE_BINARY_OP(plus, +); +CUTE_BINARY_OP(minus, -); +CUTE_BINARY_OP(multiplies, *); +CUTE_BINARY_OP(divides, /); +CUTE_BINARY_OP(modulus, %); + +CUTE_BINARY_OP(plus_assign, +=); +CUTE_BINARY_OP(minus_assign, -=); +CUTE_BINARY_OP(multiplies_assign, *=); +CUTE_BINARY_OP(divides_assign, /=); +CUTE_BINARY_OP(modulus_assign, %=); + +CUTE_BINARY_OP(bit_and, &); +CUTE_BINARY_OP(bit_or, |); +CUTE_BINARY_OP(bit_xor, ^); +CUTE_BINARY_OP(left_shift, <<); +CUTE_BINARY_OP(right_shift, >>); + +CUTE_BINARY_OP(bit_and_assign, &=); +CUTE_BINARY_OP(bit_or_assign, |=); +CUTE_BINARY_OP(bit_xor_assign, ^=); +CUTE_BINARY_OP(left_shift_assign, <<=); +CUTE_BINARY_OP(right_shift_assign, >>=); + +CUTE_BINARY_OP(logical_and, &&); +CUTE_BINARY_OP(logical_or, ||); + +CUTE_BINARY_OP(equal_to, ==); +CUTE_BINARY_OP(not_equal_to, !=); +CUTE_BINARY_OP(greater, >); +CUTE_BINARY_OP(less, <); +CUTE_BINARY_OP(greater_equal, >=); +CUTE_BINARY_OP(less_equal, <=); + +CUTE_NAMED_BINARY_OP(max_fn, cute::max); +CUTE_NAMED_BINARY_OP(min_fn, cute::min); + +#undef CUTE_BINARY_OP +#undef CUTE_NAMED_BINARY_OP + +/**********/ +/** Meta **/ +/**********/ + +template +struct bound_fn { + + template + CUTE_HOST_DEVICE constexpr + decltype(auto) + operator()(T&& arg) { + return fn_(arg_, std::forward(arg)); + } + + Fn fn_; + Arg arg_; +}; + +template +CUTE_HOST_DEVICE constexpr +auto +bind(Fn const& fn, Arg const& arg) { + return bound_fn{fn, arg}; +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/algorithm/gemm.hpp b/csrc/mctlass/include/cute/algorithm/gemm.hpp new file mode 100644 index 0000000..329a1fe --- /dev/null +++ b/csrc/mctlass/include/cute/algorithm/gemm.hpp @@ -0,0 +1,744 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include + +#include + +#include + +/** The gemm algorithm takes four (or three) tensors and computes + * D += A * B + C + * It dispatches based on the number of modes each tensor has: + * + * 1. `(V) x (V) => (V)`. + * The element-wise product of vectors. Dispatches to FMA or MMA. + * 2. `(M) x (N) => (M,N)`. + * The outer product of vectors. Dispatches to [3] with new mode K=(1). + * 3. `(M,K) x (N,K) => (M,N)`. + * The product of matrices. Dispatches to [5] with MMA vector-mode V. + * 4. `(V,M) x (V,N) => (V,M,N)`. + * The batched outer product of vectors. Accounts for register reuse and dispatches to [1] for each (m,n). + * 5. `(V,M,K) x (V,N,K) => (V,M,N)`. + * The batched product of matrices. Dispatches to [4] for each (k). + */ + +namespace cute +{ + +// +// Three arguments to four +// + +template +CUTE_HOST_DEVICE +void +gemm(Tensor const& A, + Tensor const& B, + Tensor & C) +{ + return gemm(C, A, B, C); +} + +template +CUTE_HOST_DEVICE +void +gemm(MMA_Atom const& mma, + Tensor const& A, + Tensor const& B, + Tensor & C) +{ + return gemm(mma, C, A, B, C); +} + +// +// Accept mutable temporaries +// + +template +CUTE_HOST_DEVICE +void +gemm(Tensor const& A, + Tensor const& B, + Tensor && C) +{ + return gemm(C, A, B, C); +} + +template +CUTE_HOST_DEVICE +void +gemm(Tensor && D, + Tensor const& A, + Tensor const& B, + Tensor const& C) +{ + return gemm(D, A, B, C); +} + +template +CUTE_HOST_DEVICE +void +gemm(MMA_Atom const& mma, + Tensor const& A, + Tensor const& B, + Tensor && C) +{ + return gemm(mma, C, A, B, C); +} + +template +CUTE_HOST_DEVICE +void +gemm(MMA_Atom const& mma, + Tensor && D, + Tensor const& A, + Tensor const& B, + Tensor const& C) +{ + return gemm(mma, D, A, B, C); +} + +// +// Default MMA is UniversalFMA +// + +template +CUTE_HOST_DEVICE +void +gemm(Tensor & D, + Tensor const& A, + Tensor const& B, + Tensor const& C) +{ + using MMA = MMA_Atom::value_type, + typename Tensor::value_type, + typename Tensor::value_type, + typename Tensor::value_type>>; + + return gemm(MMA{}, D, A, B, C); +} + +// +// Thread-Local Register-Memory GEMMs +// + +// Dispatch [1]: (V) x (V) => (V) +template ::value && + ALayout::rank == 1 && is_rmem::value && + BLayout::rank == 1 && is_rmem::value && + CLayout::rank == 1 && is_rmem::value)> +CUTE_HOST_DEVICE +void +gemm(MMA_Atom const& mma, + Tensor & D, // (V) Logical data + Tensor const& A, // (V) Logical data + Tensor const& B, // (V) Logical data + Tensor const& C) // (V) Logical data +{ + // No static assertions on (V), MMA checks compatibility + mma.call(D, A, B, C); +} + +// Dispatch [2]: (M) x (N) => (M,N) +template ::value && + ALayout::rank == 1 && is_rmem::value && + BLayout::rank == 1 && is_rmem::value && + CLayout::rank == 2 && is_rmem::value)> +CUTE_HOST_DEVICE +void +gemm(MMA_Atom const& mma, + Tensor & D, // (M,N) Logical data + Tensor const& A, // (M) Logical data + Tensor const& B, // (N) Logical data + Tensor const& C) // (M,N) Logical data +{ + CUTE_STATIC_ASSERT_V(size<0>(A) == size<0>(C)); // AM == CM + CUTE_STATIC_ASSERT_V(size<0>(B) == size<1>(C)); // BN == CN + CUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D)); + + gemm(mma, + D, // (M,N) + make_tensor(A.data(), append<2>(A.layout())), // (M,1) + make_tensor(B.data(), append<2>(B.layout())), // (N,1) + C); // (M,N) +} + +// Dispatch [3]: (M,K) x (N,K) => (M,N) +template ::value && + ALayout::rank == 2 && is_rmem::value && + BLayout::rank == 2 && is_rmem::value && + CLayout::rank == 2 && is_rmem::value)> +CUTE_HOST_DEVICE +void +gemm(MMA_Atom const& mma, + Tensor & D, // (M,N) Logical data + Tensor const& A, // (M,K) Logical data + Tensor const& B, // (N,K) Logical data + Tensor const& C) // (M,N) Logical data +{ + CUTE_STATIC_ASSERT_V(size<0>(A) == size<0>(C)); // AM == CM + CUTE_STATIC_ASSERT_V(size<0>(B) == size<1>(C)); // BN == CN + CUTE_STATIC_ASSERT_V(size<1>(A) == size<1>(B)); // AK == BK + CUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D)); + + // Assert this is a 1-value MMA + CUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom::LayoutC_TV{}) == Int<1>{}); + CUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom::LayoutA_TV{}) == Int<1>{}); + CUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom::LayoutB_TV{}) == Int<1>{}); + + gemm(mma, + make_tensor(D.data(), prepend<3>(D.layout())), // (1,M,N) + make_tensor(A.data(), prepend<3>(A.layout())), // (1,M,K) + make_tensor(B.data(), prepend<3>(B.layout())), // (1,N,K) + make_tensor(C.data(), prepend<3>(C.layout()))); // (1,M,N) +} + +// Dispatch [4]: (V,M) x (V,N) => (V,M,N) +template ::value && + ALayout::rank == 2 && is_rmem::value && + BLayout::rank == 2 && is_rmem::value && + CLayout::rank == 3 && is_rmem::value)> +CUTE_HOST_DEVICE +void +gemm(MMA_Atom const& mma, + Tensor & D, // (V,M,N) Logical data + Tensor const& A, // (V,M) Logical data + Tensor const& B, // (V,N) Logical data + Tensor const& C) // (V,M,N) Logical data +{ + CUTE_STATIC_ASSERT_V(size<1>(A) == size<1>(C)); // AM == CM + CUTE_STATIC_ASSERT_V(size<1>(B) == size<2>(C)); // BN == CN + CUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D) && size<2>(C) == size<2>(D)); + + auto M = size<1>(A); + auto N = size<1>(B); + // REGISTER .reuse OPTIMIZATIONS + // 64-bit traversal specialization -- serpentine path + if constexpr (decltype(size<0>(A))::value * sizeof(typename TA::value_type) == 8 && + decltype(size<0>(B))::value * sizeof(typename TB::value_type) == 8) + { +#if 1 // NOTE: Row- vs Col- major could depend on the C-matrix order... (which we can test) + // Row-major serpentine iteration + CUTE_UNROLL + for (int m = 0; m < M; ++m) { + CUTE_UNROLL + for (int n = 0; n < N; ++n) { + int ns = (m & 1) ? N-1-n : n; // Serpentine coordinate + gemm(mma, D(_,m,ns), A(_,m), B(_,ns), C(_,m,ns)); + } + } +#else + // Col-major serpentine iteration + CUTE_UNROLL + for (int n = 0; n < N; ++n) { + CUTE_UNROLL + for (int m = 0; m < M; ++m) { + int ms = (n & 1) ? M-1-m : m; // Serpentine coordinate + gemm(mma, D(_,ms,n), A(_,ms), B(_,n), C(_,ms,n)); + } + } +#endif + } else + // 32-bit traversal specialization -- kinked serpentine path + if constexpr (decltype(size<0>(A))::value * sizeof(typename TA::value_type) == 4 && + decltype(size<0>(B))::value * sizeof(typename TB::value_type) == 4) + { +#if 1 // NOTE: Row- vs Col- major could depend on the C-matrix order... (which we can test) + // Row-major kinked serpentine iteration + CUTE_UNROLL + for (int m = 0; m < M; m += 2) { + CUTE_UNROLL + for (int n = 0; n < N; ++n) { + int ns = (m & 2) ? N-1-n : n; + gemm(mma, D(_,m+0,ns), A(_,m+0), B(_,ns), C(_,m+0,ns)); + + if (m+1 < M) { + gemm(mma, D(_,m+1,ns), A(_,m+1), B(_,ns), C(_,m+1,ns)); + } + } + } +#else + // Col-major kinked serpentine iteration + CUTE_UNROLL + for (int n = 0; n < N; n += 2) { + CUTE_UNROLL + for (int m = 0; m < M; ++m) { + // Kinked serpentine traversal for maximum register reuse + int ms = (n & 2) ? M-1-m : m; + gemm(mma, D(_,ms,n+0), A(_,ms), B(_,n+0), C(_,ms,n+0)); + + if (n+1 < N) { + gemm(mma, D(_,ms,n+1), A(_,ms), B(_,n+1), C(_,ms,n+1)); + } + } + } +#endif + } else + // 64-bit + 32-bit traversal order -- keep A (64-bit) in the outer loop and serpentine B + if constexpr (decltype(size<0>(A))::value * sizeof(typename TA::value_type) == 8 && + decltype(size<0>(B))::value * sizeof(typename TB::value_type) == 4) { + // Row-major serpentine iteration + CUTE_UNROLL + for (int m = 0; m < M; ++m) { + CUTE_UNROLL + for (int n = 0; n < N; ++n) { + int ns = (m & 1) ? N-1-n : n; // Serpentine coordinate + gemm(mma, D(_,m,ns), A(_,m), B(_,ns), C(_,m,ns)); + } + } + } else + // 32-bit + 64-bit traversal order -- keep B (64-bit) in the outer loop and serpentine A + if constexpr (decltype(size<0>(A))::value * sizeof(typename TA::value_type) == 4 && + decltype(size<0>(B))::value * sizeof(typename TB::value_type) == 8) { + // Col-major serpentine iteration + CUTE_UNROLL + for (int n = 0; n < N; ++n) { + CUTE_UNROLL + for (int m = 0; m < M; ++m) { + int ms = (n & 1) ? M-1-m : m; // Serpentine coordinate + gemm(mma, D(_,ms,n), A(_,ms), B(_,n), C(_,ms,n)); + } + } + } else + // Fallback to serpentine loop + { + // Col-major serpentine iteration + CUTE_UNROLL + for (int n = 0; n < N; ++n) { + CUTE_UNROLL + for (int m = 0; m < M; ++m) { + int ms = (n & 1) ? M-1-m : m; // Serpentine coordinate + gemm(mma, D(_,ms,n), A(_,ms), B(_,n), C(_,ms,n)); + } + } + } +} + +// Dispatch [5]: (V,M,K) x (V,N,K) => (V,M,N) +template ::value && + ALayout::rank == 3 && is_rmem::value && + BLayout::rank == 3 && is_rmem::value && + CLayout::rank == 3 && is_rmem::value)> +CUTE_HOST_DEVICE +void +gemm(MMA_Atom const& mma, + Tensor & D, // (V,M,N) Logical data + Tensor const& A, // (V,M,K) Logical data + Tensor const& B, // (V,N,K) Logical data + Tensor const& C) // (V,M,N) Logical data +{ + CUTE_STATIC_ASSERT_V(size<1>(A) == size<1>(C)); // AM == CM + CUTE_STATIC_ASSERT_V(size<1>(B) == size<2>(C)); // BN == CN + CUTE_STATIC_ASSERT_V(size<2>(A) == size<2>(B)); // AK == BK + CUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D) && size<2>(C) == size<2>(D)); + + auto K = size<2>(A); + + CUTE_UNROLL + for (int k = 0; k < K; ++k) { + gemm(mma, D, A(_,_,k), B(_,_,k), C); + } +} + +// +// Thread-Local Shared-Memory GEMMs +// + +// Dispatch [1]: (V) x (V) => (V) +// Dispatch [2]: (M) x (N) => (M,N) +// Dispatch [3]: (M,K) x (N,K) => (M,N) +// Dispatch [4]: (V,M) x (V,N) => (V,M,N) +// Dispatch [5]: (V,M,K) x (V,N,K) => (V,M,N) +// Dispatch [3]: (M,K) x (N,K) => (M,N) +template ::value && + ALayout::rank == 2 && is_smem::value && + BLayout::rank == 2 && is_smem::value && + CLayout::rank == 2 && is_rmem::value)> +CUTE_HOST_DEVICE +void +gemm(MMA_Atom const& mma, + Tensor & D, // (M,N) Logical data + Tensor const& A, // (M,K) Logical data + Tensor const& B, // (N,K) Logical data + Tensor const& C) // (M,N) Logical data +{ + CUTE_STATIC_ASSERT_V(size<0>(A) == size<0>(C)); // AM == CM + CUTE_STATIC_ASSERT_V(size<0>(B) == size<1>(C)); // BN == CN + CUTE_STATIC_ASSERT_V(size<1>(A) == size<1>(B)); // AK == BK + CUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D)); + + // Assert this is a 1-value MMA + CUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom::LayoutC_TV{}) == Int<1>{}); + CUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom::LayoutA_TV{}) == Int<1>{}); + CUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom::LayoutB_TV{}) == Int<1>{}); + + gemm(mma, + make_tensor(D.data(), prepend<3>(D.layout())), // (1,M,N) + make_tensor(A.data(), prepend<3>(A.layout())), // (1,M,K) + make_tensor(B.data(), prepend<3>(B.layout())), // (1,N,K) + make_tensor(C.data(), prepend<3>(C.layout()))); // (1,M,N) +} + +// Dispatch [5]: (V,M,K) x (V,N,K) => (V,M,N) +template ::value && + ALayout::rank == 3 && is_smem::value && + BLayout::rank == 3 && is_smem::value && + CLayout::rank == 3 && is_rmem::value)> +CUTE_HOST_DEVICE +void +gemm(MMA_Atom const& mma, + Tensor & D, // (V,M,N) Logical data + Tensor const& A, // (V,M,K) Logical data + Tensor const& B, // (V,N,K) Logical data + Tensor const& C) // (V,M,N) Logical data +{ + CUTE_STATIC_ASSERT_V(size<1>(A) == size<1>(C)); // AM == CM + CUTE_STATIC_ASSERT_V(size<1>(B) == size<2>(C)); // BN == CN + CUTE_STATIC_ASSERT_V(size<2>(A) == size<2>(B)); // AK == BK + CUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D) && size<2>(C) == size<2>(D)); + + auto rA = MMA_Atom::make_fragment_A(A); + auto rB = MMA_Atom::make_fragment_B(B); + + auto K = size<2>(A); + + CUTE_UNROLL + for (int k = 0; k < K; ++k) + { + copy(A(_,_,k), rA(_,_,k)); + copy(B(_,_,k), rB(_,_,k)); + // Thread-level register gemm for k + gemm(mma, D, rA(_,_,k), rB(_,_,k), C); + } +} + +// +// Collective Shared-Memory GEMMs +// + +template ::value && + BLayout::rank == 2 && is_smem::value && + CLayout::rank == 2 && is_smem::value)> +CUTE_HOST_DEVICE +void +gemm(ThrMMA const& thr_mma, + Alpha const& alpha, + Tensor sA, + Tensor sB, + Beta const& beta, + Tensor sC, + ALoadTransformOp const& sA_load_op /* transforms A values before used in GEMM */, + BLoadTransformOp const& sB_load_op /* transforms B values before used in GEMM */) +{ + CUTE_STATIC_ASSERT_V(size<0>(sA) == size<0>(sC)); // AM == CM + CUTE_STATIC_ASSERT_V(size<0>(sB) == size<1>(sC)); // BN == CN + CUTE_STATIC_ASSERT_V(size<1>(sA) == size<1>(sB)); // AK == BK + + using TypeA = typename TA::value_type; + using TypeB = typename TB::value_type; + using TypeC = typename TC::value_type; + + static_assert(is_same_v>, TypeA>, + "ALoadTransformOp functor must accept and return value of type TA::value_type"); + static_assert(is_same_v>, TypeB>, + "BLoadTransformOp functor must accept and return value of type TB::value_type"); + + // Original, static size of the problem + auto M = size<0>(sC); + auto N = size<1>(sC); + auto K = size<1>(sA); + + // Block size of the compute tile + auto BLK_M = tile_size<0>(thr_mma); + auto BLK_N = tile_size<1>(thr_mma); + auto BLK_K = tile_size<2>(thr_mma); + + // Compute the "residues" + auto m_residue = M - BLK_M * (ceil_div(M, BLK_M) - Int<1>{}); // (0,BLK_M] + auto n_residue = N - BLK_N * (ceil_div(N, BLK_N) - Int<1>{}); // (0,BLK_N] + auto k_residue = K - BLK_K * (ceil_div(K, BLK_K) ); // (-BLK_K,0] + + // Shift the origin so k_residue is zeroth tile + sA.data() = &sA(0,k_residue); + sB.data() = &sB(0,k_residue); + +#if 0 + if (thread0()) { + printf("%d in BLK_M (%d)\n", int(m_residue), int(BLK_M)); + printf("%d in BLK_N (%d)\n", int(n_residue), int(BLK_N)); + printf("%d in BLK_K (%d)\n", int(k_residue), int(BLK_K)); + } +#endif + + // + // MMA Partitioning + // + + // Round the layout extents up to BLK_X + Tensor rounded_sA = sA.compose(make_shape(ceil_div(M, BLK_M) * BLK_M, ceil_div(K, BLK_K) * BLK_K)); + Tensor rounded_sB = sB.compose(make_shape(ceil_div(N, BLK_N) * BLK_N, ceil_div(K, BLK_K) * BLK_K)); + Tensor rounded_sC = sC.compose(make_shape(ceil_div(M, BLK_M) * BLK_M, ceil_div(N, BLK_N) * BLK_N)); + +#if 0 + if (thread0()) { + print(rounded_sA.layout()); print("\n"); + print(rounded_sB.layout()); print("\n"); + print(rounded_sC.layout()); print("\n"); + } +#endif + + // Partition the sA and sB tiles across the threads for the MMA + Tensor tCsA = thr_mma.partition_A(rounded_sA); // (MMA,MMA_M,MMA_K) + Tensor tCsB = thr_mma.partition_B(rounded_sB); // (MMA,MMA_N,MMA_K) + Tensor tCsC = thr_mma.partition_C(rounded_sC); // (MMA,MMA_M,MMA_N) + // Create register tensors for the MMA to operate on + Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K) + Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K) + Tensor tCrC = thr_mma.make_fragment_C(tCsC); // (MMA,MMA_M,MMA_N) + +#if 0 + if (thread0()) { + print(tCsA.layout()); print("\n"); + print(tCsB.layout()); print("\n"); + print(tCsC.layout()); print("\n"); + print(tCrA.layout()); print("\n"); + print(tCrB.layout()); print("\n"); + print(tCrC.layout()); print("\n"); + } +#endif + + // + // PREDICATION + // + + // Allocate the preds for only the MMA-mode of tCsA and tCsB + Tensor tCpA = make_tensor(size<0>(tCsA)); + Tensor tCpB = make_tensor(size<0>(tCsB)); + + // Create coordinate tensors on a single compute block for predication + Tensor cA = make_identity_tensor(make_shape(BLK_M, BLK_K)); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cB = make_identity_tensor(make_shape(BLK_N, BLK_K)); // (BLK_M,BLK_K) -> (blk_n,blk_k) + + // Repeat partitioning with thr_mma + Tensor tCcA = thr_mma.partition_A(cA); // (MMA,1,1) -> (blk_m,blk_k) + Tensor tCcB = thr_mma.partition_B(cB); // (MMA,1,1) -> (blk_n,blk_k) + + // Populate the m and n predicates + CUTE_UNROLL + for (int i = 0; i < size(tCpA); ++i) { + tCpA(i) = elem_less(get<0>(tCcA(i)), m_residue); + } + CUTE_UNROLL + for (int i = 0; i < size(tCpB); ++i) { + tCpB(i) = elem_less(get<0>(tCcB(i)), n_residue); + } + +#if 0 + printf("Thr %d: A(%d,%d):%d B(%d,%d):%d\n", + threadIdx.x, + int(get<0>(tCcA(0))), int(get<1>(tCcA(0))), int(tCpA(0)), + int(get<0>(tCcB(0))), int(get<1>(tCcB(0))), int(tCpB(0))); +#endif + + // + // PREFETCH k_block = 0 (with k-predication) + // + + CUTE_UNROLL + for (int i = 0; i < size<0>(tCsA); ++i) { // Copy MMA_I + if (k_residue == 0 || get<1>(tCcA(i)) >= -k_residue) { // k_block = 0, predicated on k + CUTE_UNROLL + for (int m = 0; m < size<1>(tCsA); ++m) { // Copy MMA_M, predicated on m + tCrA(i,m,0) = (m_residue == BLK_M || m < size<1>(tCsA)-1 || tCpA(i)) ? sA_load_op(tCsA(i,m,0)) : TypeA{}; + } + } + } + + CUTE_UNROLL + for (int i = 0; i < size<0>(tCsB); ++i) { // Copy MMA_I + if (k_residue == 0 || get<1>(tCcB(i)) >= -k_residue) { // k_block = 0, predicated on k + CUTE_UNROLL + for (int n = 0; n < size<1>(tCsB); ++n) { // Copy MMA_N, predicated on n + tCrB(i,n,0) = (n_residue == BLK_N || n < size<1>(tCsB)-1 || tCpB(i)) ? sB_load_op(tCsB(i,n,0)) : TypeB{}; + } + } + } + // + // MAINLOOP + // + + // Clear accumulators + clear(tCrC); + + constexpr int K_BLOCK_MAX = size<2>(tCrA); + + CUTE_UNROLL + for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) + { + // static-if load the next k_block. No k-predication required on these loads. + if (k_block < K_BLOCK_MAX-1) + { + // Load the next k_block + int k_next = k_block + 1; + + CUTE_UNROLL + for (int m = 0; m < size<1>(tCsA); ++m) { // Copy MMA_M + CUTE_UNROLL + for (int i = 0; i < size<0>(tCsA); ++i) { // Copy_if MMA_I predicated on m + tCrA(i,m,k_next) = (m_residue == BLK_M || m < size<1>(tCsA)-1 || tCpA(i)) ? sA_load_op(tCsA(i,m,k_next)) : TypeA{}; + } + } + + CUTE_UNROLL + for (int n = 0; n < size<1>(tCsB); ++n) { // Copy MMA_N + CUTE_UNROLL + for (int i = 0; i < size<0>(tCsB); ++i) { // Copy MMA_I predicated on n + tCrB(i,n,k_next) = (n_residue == BLK_N || n < size<1>(tCsB)-1 || tCpB(i)) ? sB_load_op(tCsB(i,n,k_next)) : TypeB{}; + } + } + } + + // GEMM on k_block in registers + gemm(thr_mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC); + } + + // + // Epilogue + // + + Tensor cC = make_identity_tensor(make_shape(BLK_M, BLK_N)); // (BLK_M,BLK_N) -> (blk_m,blk_n) + Tensor tCcC = thr_mma.partition_C(cC); // (MMA, 1, 1) -> (blk_m,blk_n) + + const bool isBetaZero = (beta == Beta{}); + + // Custom axpby_if for now + CUTE_UNROLL + for (int m = 0; m < size<1>(tCsC); ++m) + { + CUTE_UNROLL + for (int n = 0; n < size<2>(tCsC); ++n) + { + CUTE_UNROLL + for (int i = 0; i < size<0>(tCsC); ++i) + { + if ((m_residue == BLK_M || m < size<1>(tCrC)-1 || get<0>(tCcC(i)) < m_residue) && + (n_residue == BLK_N || n < size<2>(tCrC)-1 || get<1>(tCcC(i)) < n_residue)) + { + tCsC(i,m,n) = isBetaZero ? alpha * tCrC(i,m,n) : alpha * tCrC(i,m,n) + beta * tCsC(i,m,n); + } + } + } + } +} + +template ::value && + BLayout::rank == 2 && is_smem::value && + CLayout::rank == 2 && is_smem::value)> +CUTE_HOST_DEVICE +void +gemm(ThrMMA const& thr_mma, + Alpha const& alpha, + Tensor sA, + Tensor sB, + Beta const& beta, + Tensor sC) +{ + gemm(thr_mma, alpha, sA, sB, beta, sC, identity() /* sA_load_op */, identity() /* sB_load_op */); +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/algorithm/prefer.hpp b/csrc/mctlass/include/cute/algorithm/prefer.hpp new file mode 100644 index 0000000..804896c --- /dev/null +++ b/csrc/mctlass/include/cute/algorithm/prefer.hpp @@ -0,0 +1,46 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +namespace cute +{ + +// Infinite types that inherit from each other +template +struct prefer : prefer {}; + +template <> +struct prefer<0> {}; + +// Can be used to preferencially overload implementations +// Higher N in prefer have higher priority. + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/algorithm/tensor_algorithms.hpp b/csrc/mctlass/include/cute/algorithm/tensor_algorithms.hpp new file mode 100644 index 0000000..5fac8f9 --- /dev/null +++ b/csrc/mctlass/include/cute/algorithm/tensor_algorithms.hpp @@ -0,0 +1,123 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/** Common algorithms on (hierarchical) tensors */ + +#pragma once + +#include + +#include + +namespace cute +{ + +// +// for_each +// + +template +CUTE_HOST_DEVICE constexpr +void +for_each(Tensor const& tensor, UnaryOp&& op) +{ + CUTE_UNROLL + for (int i = 0; i < size(tensor); ++i) { + static_cast(op)(tensor(i)); + } +} + +template +CUTE_HOST_DEVICE constexpr +void +for_each(Tensor& tensor, UnaryOp&& op) +{ + CUTE_UNROLL + for (int i = 0; i < size(tensor); ++i) { + static_cast(op)(tensor(i)); + } +} + +// Accept mutable temporaries +template +CUTE_HOST_DEVICE constexpr +void +for_each(Tensor&& tensor, UnaryOp&& op) +{ + return for_each(tensor, static_cast(op)); +} + +// +// transform +// + +// Similar to std::transform but does not return number of elements affected +template +CUTE_HOST_DEVICE constexpr +void +transform(Tensor& tensor, UnaryOp&& op) +{ + CUTE_UNROLL + for (int i = 0; i < size(tensor); ++i) { + tensor(i) = static_cast(op)(tensor(i)); + } +} + +// Accept mutable temporaries +template +CUTE_HOST_DEVICE constexpr +void +transform(Tensor&& tensor, UnaryOp&& op) +{ + return transform(tensor, std::forward(op)); +} + +// Similar to std::transform transforms one tensors and assigns it to another +template +CUTE_HOST_DEVICE constexpr +void +transform(Tensor& tensor_in, Tensor& tensor_out, UnaryOp&& op) +{ + CUTE_UNROLL + for (int i = 0; i < size(tensor_in); ++i) { + tensor_out(i) = static_cast(op)(tensor_in(i)); + } +} + +// Accept mutable temporaries +template +CUTE_HOST_DEVICE constexpr +void +transform(Tensor&& tensor_in, Tensor&& tensor_out, UnaryOp&& op) +{ + return transform(tensor_in, tensor_out, std::forward(op)); +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/algorithm/tuple_algorithms.hpp b/csrc/mctlass/include/cute/algorithm/tuple_algorithms.hpp new file mode 100644 index 0000000..ed338cc --- /dev/null +++ b/csrc/mctlass/include/cute/algorithm/tuple_algorithms.hpp @@ -0,0 +1,875 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include +#include +#include +#include + +/** Common algorithms on (hierarchical) tuples */ +/** Style choice: + * Forward params [using static_cast(.)] for const/non-const/ref/non-ref args + * but don't bother forwarding functions as ref-qualified member fns are extremely rare + */ + +namespace cute +{ + +// +// Apply (Unpack) +// (t, f) => f(t_0,t_1,...,t_n) +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +apply(T&& t, F&& f, seq) +{ + return f(get(static_cast(t))...); +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +apply(T&& t, F&& f) +{ + return detail::apply(static_cast(t), f, tuple_seq{}); +} + +// +// Transform Apply +// (t, f, g) => g(f(t_0),f(t_1),...) +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +tapply(T&& t, F&& f, G&& g, seq) +{ + return g(f(get(static_cast(t)))...); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tapply(T0&& t0, T1&& t1, F&& f, G&& g, seq) +{ + return g(f(get(static_cast(t0)), + get(static_cast(t1)))...); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tapply(T0&& t0, T1&& t1, T2&& t2, F&& f, G&& g, seq) +{ + return g(f(get(static_cast(t0)), + get(static_cast(t1)), + get(static_cast(t2)))...); +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +transform_apply(T&& t, F&& f, G&& g) +{ + return detail::tapply(static_cast(t), f, g, tuple_seq{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +transform_apply(T0&& t0, T1&& t1, F&& f, G&& g) +{ + return detail::tapply(static_cast(t0), static_cast(t1), f, g, tuple_seq{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +transform_apply(T0&& t0, T1&& t1, T2&& t2, F&& f, G&& g) +{ + return detail::tapply(static_cast(t0), static_cast(t1), static_cast(t2), f, g, tuple_seq{}); +} + +// +// For Each +// (t, f) => f(t_0),f(t_1),...,f(t_n) +// + +template +CUTE_HOST_DEVICE constexpr +void +for_each(T&& t, F&& f) +{ + detail::apply(t, [&](auto&&... a) { (f(static_cast(a)), ...); }, tuple_seq{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +for_each_leaf(T&& t, F&& f) +{ + if constexpr (is_tuple>::value) { + return detail::apply(static_cast(t), [&](auto&&... a){ return (for_each_leaf(static_cast(a), f), ...); }, tuple_seq{}); + } else { + return f(static_cast(t)); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Transform +// (t, f) => (f(t_0),f(t_1),...,f(t_n)) +// + +template +CUTE_HOST_DEVICE constexpr +auto +transform(T const& t, F&& f) +{ + return detail::tapply(t, f, [](auto const&... a){ return cute::make_tuple(a...); }, tuple_seq{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +transform(T0 const& t0, T1 const& t1, F&& f) +{ + static_assert(tuple_size::value == tuple_size::value, "Mismatched tuple_size"); + return detail::tapply(t0, t1, f, [](auto const&... a){ return cute::make_tuple(a...); }, tuple_seq{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +transform(T0 const& t0, T1 const& t1, T2 const& t2, F&& f) +{ + static_assert(tuple_size::value == tuple_size::value, "Mismatched tuple_size"); + static_assert(tuple_size::value == tuple_size::value, "Mismatched tuple_size"); + return detail::tapply(t0, t1, t2, f, [](auto const&... a){ return cute::make_tuple(a...); }, tuple_seq{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +transform_leaf(T const& t, F&& f) +{ + if constexpr (is_tuple::value) { + return transform(t, [&](auto const& a) { return transform_leaf(a, f); }); + } else { + return f(t); + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +transform_leaf(T0 const& t0, T1 const& t1, F&& f) +{ + if constexpr (is_tuple::value) { + return transform(t0, t1, [&](auto const& a, auto const& b) { return transform_leaf(a, b, f); }); + } else { + return f(t0, t1); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// find and find_if +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +find_if(T const& t, F&& f, seq<>) +{ + return cute::integral_constant::value>{}; +} + +template +CUTE_HOST_DEVICE constexpr +auto +find_if(T const& t, F&& f, seq) +{ + if constexpr (decltype(f(get(t)))::value) { + return cute::integral_constant{}; + } else { + return find_if(t, f, seq{}); + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +find_if(T const& t, F&& f) +{ + if constexpr (is_tuple::value) { + return detail::find_if(t, f, tuple_seq{}); + } else { + return cute::integral_constant{}; + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +find(T const& t, X const& x) +{ + return find_if(t, [&](auto const& v) { return v == x; }); // This should always return a static true/false +} + +template +CUTE_HOST_DEVICE constexpr +auto +none_of(T const& t, F&& f) +{ + if constexpr (is_tuple::value) { + return cute::integral_constant::value>{}; + } else { + return not f(t); + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +all_of(T const& t, F&& f) +{ + if constexpr (is_tuple::value) { + auto not_f = [&](auto const& a) { return not f(a); }; + return cute::integral_constant::value>{}; + } else { + return f(t); + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +any_of(T const& t, F&& f) +{ + return not none_of(t, f); +} + +// +// Filter +// (t, f) => +// + +template +CUTE_HOST_DEVICE constexpr +auto +filter_tuple(T const& t, F&& f) +{ + return transform_apply(t, f, [](auto const&... a) { return cute::tuple_cat(a...); }); +} + +template +CUTE_HOST_DEVICE constexpr +auto +filter_tuple(T0 const& t0, T1 const& t1, F&& f) +{ + return transform_apply(t0, t1, f, [](auto const&... a) { return cute::tuple_cat(a...); }); +} + +// +// Fold (Reduce, Accumulate) +// (t, v, f) => f(...f(f(v,t_0),t_1),...,t_n) +// + +namespace detail { + +// This impl compiles much faster than cute::apply and variadic args +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +fold(T&& t, V&& v, F&& f, seq<>) +{ + return static_cast(v); +} + +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +fold(T&& t, V&& v, F&& f, seq) +{ + if constexpr (sizeof...(Is) == 0) { + return f(static_cast(v), get(static_cast(t))); + } else { + return fold(static_cast(t), + f(static_cast(v), get(static_cast(t))), + f, + seq{}); + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +fold(T&& t, V&& v, F&& f) +{ + if constexpr (is_tuple>::value) { + return detail::fold(static_cast(t), + static_cast(v), + f, + tuple_seq{}); + } else { + return f(static_cast(v), static_cast(t)); + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +fold_first(T&& t, F&& f) +{ + if constexpr (is_tuple>::value) { + return detail::fold(static_cast(t), + get<0>(static_cast(t)), + f, + make_range<1,tuple_size>::value>{}); + } else { + return static_cast(t); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// front, back, take, unwrap +// + +// Get the first non-tuple element in a hierarchical tuple +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +front(T&& t) +{ + if constexpr (is_tuple>::value) { + return front(get<0>(static_cast(t))); + } else { + return static_cast(t); + } + + CUTE_GCC_UNREACHABLE; +} + +// Get the last non-tuple element in a hierarchical tuple +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +back(T&& t) +{ + if constexpr (is_tuple>::value) { + constexpr int N = tuple_size>::value; + return back(get(static_cast(t))); + } else { + return static_cast(t); + } + + CUTE_GCC_UNREACHABLE; +} + +// Takes the elements in the range [B,E) +template +CUTE_HOST_DEVICE constexpr +auto +take(T const& t) +{ + return detail::apply(t, [](auto const&... a) { return cute::make_tuple(a...); }, make_range{}); +} + +// Unwrap rank-1 tuples until we're left with a rank>1 tuple or a non-tuple +template +CUTE_HOST_DEVICE constexpr +auto +unwrap(T const& t) +{ + if constexpr (is_tuple::value) { + if constexpr (tuple_size::value == 1) { + return unwrap(get<0>(t)); + } else { + return t; + } + } else { + return t; + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Flatten a hierarchical tuple to a tuple of depth one. +// + +template +CUTE_HOST_DEVICE constexpr +auto +flatten_to_tuple(T const& t) +{ + if constexpr (is_tuple::value) { + return filter_tuple(t, [](auto const& a) { return flatten_to_tuple(a); }); + } else { + return cute::make_tuple(t); + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +flatten(T const& t) +{ + if constexpr (is_tuple::value) { + return filter_tuple(t, [](auto const& a) { return flatten_to_tuple(a); }); + } else { + return t; + } + + CUTE_GCC_UNREACHABLE; +} + +// +// insert and remove and replace +// + +namespace detail { + +// Shortcut around cute::tuple_cat for common insert/remove/repeat cases +template +CUTE_HOST_DEVICE constexpr +auto +construct(T const& t, X const& x, seq, seq, seq) +{ + return cute::make_tuple(get(t)..., (void(J),x)..., get(t)...); +} + +} // end namespace detail + +// Insert x into the Nth position of the tuple +template +CUTE_HOST_DEVICE constexpr +auto +insert(T const& t, X const& x) +{ + return detail::construct(t, x, make_seq{}, seq<0>{}, make_range::value>{}); +} + +// Remove the Nth element of the tuple +template +CUTE_HOST_DEVICE constexpr +auto +remove(T const& t) +{ + return detail::construct(t, 0, make_seq{}, seq<>{}, make_range::value>{}); +} + +// Replace the Nth element of the tuple with x +template +CUTE_HOST_DEVICE constexpr +auto +replace(T const& t, X const& x) +{ + return detail::construct(t, x, make_seq{}, seq<0>{}, make_range::value>{}); +} + +// Replace the first element of the tuple with x +template +CUTE_HOST_DEVICE constexpr +auto +replace_front(T const& t, X const& x) +{ + if constexpr (is_tuple::value) { + return detail::construct(t, x, seq<>{}, seq<0>{}, make_range<1,tuple_size::value>{}); + } else { + return x; + } + + CUTE_GCC_UNREACHABLE; +} + +// Replace the last element of the tuple with x +template +CUTE_HOST_DEVICE constexpr +auto +replace_back(T const& t, X const& x) +{ + if constexpr (is_tuple::value) { + return detail::construct(t, x, make_seq::value-1>{}, seq<0>{}, seq<>{}); + } else { + return x; + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Make a tuple of Xs of tuple_size N +// + +template +CUTE_HOST_DEVICE constexpr +auto +repeat(X const& x) +{ + return detail::construct(0, x, seq<>{}, make_seq{}, seq<>{}); +} + +// +// Make a tuple of Xs the same profile as tuple +// + +template +CUTE_HOST_DEVICE constexpr +auto +repeat_like(T const& t, X const& x) +{ + if constexpr (is_tuple::value) { + return transform(t, [&](auto const& a) { return repeat_like(a,x); }); + } else { + return x; + } + + CUTE_GCC_UNREACHABLE; +} + +// Group the elements [B,E) of a T into a single element +// e.g. group<2,4>(T<_1,_2,_3,_4,_5,_6>{}) +// => T<_1,_2,T<_3,_4>,_5,_6>{} +template +CUTE_HOST_DEVICE constexpr +auto +group(T const& t) +{ + return detail::construct(t, take(t), make_seq{}, seq<0>{}, make_range::value>{}); +} + +// +// Extend a T to rank N by appending/prepending an element +// + +template +CUTE_HOST_DEVICE constexpr +auto +append(T const& a, X const& x) +{ + if constexpr (is_tuple::value) { + if constexpr (N == tuple_size::value) { + return a; + } else { + static_assert(N > tuple_size::value); + return detail::construct(a, x, make_seq::value>{}, make_seq::value>{}, seq<>{}); + } + } else { + if constexpr (N == 1) { + return a; + } else { + return detail::construct(cute::make_tuple(a), x, seq<0>{}, make_seq{}, seq<>{}); + } + } + + CUTE_GCC_UNREACHABLE; +} +template +CUTE_HOST_DEVICE constexpr +auto +append(T const& a, X const& x) +{ + if constexpr (is_tuple::value) { + return detail::construct(a, x, make_seq::value>{}, seq<0>{}, seq<>{}); + } else { + return cute::make_tuple(a, x); + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +prepend(T const& a, X const& x) +{ + if constexpr (is_tuple::value) { + if constexpr (N == tuple_size::value) { + return a; + } else { + static_assert(N > tuple_size::value); + return detail::construct(a, x, seq<>{}, make_seq::value>{}, make_seq::value>{}); + } + } else { + if constexpr (N == 1) { + return a; + } else { + static_assert(N > 1); + return detail::construct(cute::make_tuple(a), x, seq<>{}, make_seq{}, seq<0>{}); + } + } + + CUTE_GCC_UNREACHABLE; +} +template +CUTE_HOST_DEVICE constexpr +auto +prepend(T const& a, X const& x) +{ + if constexpr (is_tuple::value) { + return detail::construct(a, x, seq<>{}, seq<0>{}, make_seq::value>{}); + } else { + return cute::make_tuple(x, a); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Inclusive scan (prefix sum) +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +iscan(T const& t, V const& v, F&& f, seq) +{ + // Apply the function to v and the element at I + auto v_next = f(v, get(t)); + // Replace I with v_next + auto t_next = replace(t, v_next); + +#if 0 + std::cout << "ISCAN i" << I << std::endl; + std::cout << " t " << t << std::endl; + std::cout << " i " << v << std::endl; + std::cout << " f(i,t) " << v_next << std::endl; + std::cout << " t_n " << t_next << std::endl; +#endif + + if constexpr (sizeof...(Is) == 0) { + return t_next; + } else { + return iscan(t_next, v_next, f, seq{}); + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +iscan(T const& t, V const& v, F&& f) +{ + return detail::iscan(t, v, f, tuple_seq{}); +} + +// +// Exclusive scan (prefix sum) +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +escan(T const& t, V const& v, F&& f, seq) +{ + if constexpr (sizeof...(Is) == 0) { + // Replace I with v + return replace(t, v); + } else { + // Apply the function to v and the element at I + auto v_next = f(v, get(t)); + // Replace I with v + auto t_next = replace(t, v); + +#if 0 + std::cout << "ESCAN i" << I << std::endl; + std::cout << " t " << t << std::endl; + std::cout << " i " << v << std::endl; + std::cout << " f(i,t) " << v_next << std::endl; + std::cout << " t_n " << t_next << std::endl; +#endif + + // Recurse + return escan(t_next, v_next, f, seq{}); + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +escan(T const& t, V const& v, F&& f) +{ + return detail::escan(t, v, f, tuple_seq{}); +} + +// +// Zip (Transpose) +// + +// Take ((a,b,c,...),(x,y,z,...),...) rank-R0 x rank-R1 input +// to produce ((a,x,...),(b,y,...),(c,z,...),...) rank-R1 x rank-R0 output + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +zip_(Ts const&... ts) +{ + return cute::make_tuple(get(ts)...); +} + +template +CUTE_HOST_DEVICE constexpr +auto +zip(T const& t, seq, seq) +{ + static_assert(conjunction>::value == tuple_size>::value>...>::value, "Mismatched Ranks"); + return cute::make_tuple(zip_(get(t)...)...); +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +zip(T const& t) +{ + if constexpr (is_tuple::value) { + if constexpr (is_tuple>::value) { + return detail::zip(t, tuple_seq{}, tuple_seq>{}); + } else { + return cute::make_tuple(t); + } + } else { + return t; + } + + CUTE_GCC_UNREACHABLE; +} + +// Convenient to pass them in separately +template +CUTE_HOST_DEVICE constexpr +auto +zip(T0 const& t0, T1 const& t1, Ts const&... ts) +{ + return zip(cute::make_tuple(t0, t1, ts...)); +} + +// +// zip2_by -- A guided zip for rank-2 tuples +// Take a tuple like ((A,a),((B,b),(C,c)),d) +// and produce a tuple ((A,(B,C)),(a,(b,c),d)) +// where the rank-2 modes are selected by the terminals of the guide (X,(X,X)) +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +zip2_by(T const& t, TG const& guide, seq, seq) +{ + // zip2_by produces the modes like ((A,a),(B,b),...) + auto split = cute::make_tuple(zip2_by(get(t), get(guide))...); + + // Rearrange and append missing modes from t to make ((A,B,...),(a,b,...,x,y)) + return cute::make_tuple(cute::make_tuple(get<0>(get(split))...), + cute::make_tuple(get<1>(get(split))..., get(t)...)); +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +zip2_by(T const& t, TG const& guide) +{ + if constexpr (is_tuple::value) { + constexpr int TR = tuple_size::value; + constexpr int GR = tuple_size::value; + static_assert(TR >= GR, "Mismatched ranks"); + return detail::zip2_by(t, guide, + make_range< 0, GR>{}, + make_range{}); + } else { + static_assert(tuple_size::value == 2, "Mismatched ranks"); + return t; + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/arch/cluster_sm90.hpp b/csrc/mctlass/include/cute/arch/cluster_sm90.hpp new file mode 100644 index 0000000..a31871d --- /dev/null +++ b/csrc/mctlass/include/cute/arch/cluster_sm90.hpp @@ -0,0 +1,243 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +// Config +// #if (defined(__MACA_ARCH__) && (__MACA_ARCH__ >= 900) && \ +// ((__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8)))) +// # define CUTE_ARCH_CLUSTER_SM90_ENABLED +// #endif + +// #if (defined(__MACA_ARCH__) && (__MACA_ARCH__ >= 900) && (__CUDACC_VER_MAJOR__ >= 12)) +// # define CUTE_ARCH_ELECT_ONE_SM90_ENABLED +// #endif + +namespace cute { + +CUTE_DEVICE void cluster_arrive_relaxed() +{ +#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED) + asm volatile("barrier.cluster.arrive.relaxed.aligned;\n" : : ); +#else + CUTE_RUNTIME_ASSERT("CUTE_ARCH_CLUSTER_SM90_ENABLED is not defined"); +#endif +} + +CUTE_DEVICE void cluster_arrive() +{ +#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED) + asm volatile("barrier.cluster.arrive.aligned;\n" : : ); +#else + CUTE_RUNTIME_ASSERT("CUTE_ARCH_CLUSTER_SM90_ENABLED is not defined"); +#endif +} + +CUTE_DEVICE void cluster_wait() +{ +#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED) + asm volatile("barrier.cluster.wait.aligned;\n" : : ); +#else + CUTE_RUNTIME_ASSERT("CUTE_ARCH_CLUSTER_SM90_ENABLED is not defined"); +#endif +} + +CUTE_DEVICE void cluster_sync() +{ +#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED) + cluster_arrive(); + cluster_wait(); +#else + CUTE_RUNTIME_ASSERT("CUTE_ARCH_CLUSTER_SM90_ENABLED is not defined"); +#endif +} + +// Returns the dim3 grid size in terms of number of clusters. +CUTE_DEVICE dim3 cluster_grid_dims() +{ +#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED) + uint32_t x, y, z; + asm volatile("mov.u32 %0, %nclusterid.x;\n" : "=r"(x) : ); + asm volatile("mov.u32 %0, %nclusterid.y;\n" : "=r"(y) : ); + asm volatile("mov.u32 %0, %nclusterid.z;\n" : "=r"(z) : ); + return {x, y, z}; +#elif defined(__MACA_ARCH__) + // MSVC requires protecting use of gridDim with __MACA_ARCH__. + return gridDim; +#elif defined(_MSC_VER) + CUTE_RUNTIME_ASSERT("cluster_grid_dims() can only be called on device"); +#else + return {0, 0, 0}; +#endif +} + +// Returns the dim3 cluster rank in the grid. +CUTE_DEVICE dim3 cluster_id_in_grid() +{ +#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED) + uint32_t x, y, z; + asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(x) : ); + asm volatile("mov.u32 %0, %clusterid.y;\n" : "=r"(y) : ); + asm volatile("mov.u32 %0, %clusterid.z;\n" : "=r"(z) : ); + return {x, y, z}; +#elif defined(__MACA_ARCH__) + // MSVC requires protecting use of blockIdx with __MACA_ARCH__. + return blockIdx; +#elif defined(_MSC_VER) + CUTE_RUNTIME_ASSERT("cluster_id_in_grid() can only be called on device"); +#else + return {0, 0, 0}; +#endif +} + +// Returns the relative dim3 block rank local to the cluster. +CUTE_DEVICE dim3 block_id_in_cluster() +{ +#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED) + uint32_t x, y, z; + asm volatile("mov.u32 %0, %cluster_ctaid.x;\n" : "=r"(x) : ); + asm volatile("mov.u32 %0, %cluster_ctaid.y;\n" : "=r"(y) : ); + asm volatile("mov.u32 %0, %cluster_ctaid.z;\n" : "=r"(z) : ); + return {x, y, z}; +#else + return {0,0,0}; +#endif +} + +// Returns the dim3 cluster shape. +CUTE_DEVICE dim3 cluster_shape() +{ +#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED) + uint32_t x, y, z; + asm volatile("mov.u32 %0, %cluster_nctaid.x;\n" : "=r"(x) : ); + asm volatile("mov.u32 %0, %cluster_nctaid.y;\n" : "=r"(y) : ); + asm volatile("mov.u32 %0, %cluster_nctaid.z;\n" : "=r"(z) : ); + return {x, y, z}; +#else + return {1,1,1}; +#endif +} + +// Get 1D ctaid in a cluster. +MCTLASS_DEVICE uint32_t block_rank_in_cluster() +{ +#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED) + uint32_t rank; + asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(rank) :); + return rank; +#else + return 0; +#endif +} + +// Set the destination block-ID in cluster for a given SMEM Address +MCTLASS_DEVICE uint32_t set_block_rank(uint32_t smemAddr, uint32_t rank) +{ +#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED) + uint32_t result; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;\n" + : "=r"(result) + : "r"(smemAddr), "r"(rank)); + return result; +#else + return smemAddr; +#endif +} + +// Elect one thread in the warp. The elected thread gets its predicate set to true, all others obtain false. +CUTE_HOST_DEVICE uint32_t elect_one_sync() +{ +#if defined(CUTE_ARCH_ELECT_ONE_SM90_ENABLED) + uint32_t pred = 0; + uint32_t laneid = 0; + asm volatile( + "{\n" + ".reg .b32 %rx;\n" + ".reg .pred %px;\n" + " elect.sync %rx|%px, %2;\n" + "@%px mov.s32 %1, 1;\n" + " mov.s32 %0, %rx;\n" + "}\n" + : "+r"(laneid), "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +#elif defined(__MACA_ARCH__) + return (threadIdx.x % 64) == 0; +#else + return true; +#endif +} + +struct ElectOneLaneIdReturnType { + uint32_t is_leader; + uint32_t leader_lane_id; +}; + +CUTE_HOST_DEVICE +ElectOneLaneIdReturnType +elect_one_leader_sync() +{ +#if defined(CUTE_ARCH_ELECT_ONE_SM90_ENABLED) + uint32_t pred = 0; + uint32_t laneid = 0; + asm volatile( + "{\n" + ".reg .b32 %rx;\n" + ".reg .pred %px;\n" + " elect.sync %rx|%px, %2;\n" + "@%px mov.s32 %1, 1;\n" + " mov.s32 %0, %rx;\n" + "}\n" + : "+r"(laneid), "+r"(pred) + : "r"(0xFFFFFFFF)); + return {pred, laneid}; +#elif defined(__MACA_ARCH__) + return {(threadIdx.x % 64) == 0, 0}; +#else + return {true, 0}; +#endif +} + +// Store value to remote shared memory in the cluster +CUTE_DEVICE +void +store_shared_remote(uint32_t value, uint32_t smem_addr, uint32_t mbarrier_addr, uint32_t dst_cta_rank) +{ +#if defined(CUTE_ARCH_CLUSTER_SM90_ENABLED) + uint32_t dsmem_addr = set_block_rank(smem_addr, dst_cta_rank); + uint32_t remote_barrier_addr = set_block_rank(mbarrier_addr, dst_cta_rank); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.u32 [%0], %1, [%2];" + : : "r"(dsmem_addr), "r"(value), "r"(remote_barrier_addr)); +#endif +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/arch/copy.hpp b/csrc/mctlass/include/cute/arch/copy.hpp new file mode 100644 index 0000000..aa7bb33 --- /dev/null +++ b/csrc/mctlass/include/cute/arch/copy.hpp @@ -0,0 +1,71 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include + +namespace cute +{ + +// +// Direct Copy for any type +// + +template +struct UniversalCopy +{ + using SRegisters = S[1]; + using DRegisters = D[1]; + + CUTE_HOST_DEVICE static constexpr void + copy(S const& src, + D & dst) + { + dst = src; + } +}; + +// +// Placeholder for the copy algorithm's default, auto-vectorizing behavior +// + +struct DefaultCopy +{ + using SRegisters = uint128_t[1]; + using DRegisters = uint128_t[1]; +}; + +using AutoVectorizingCopy = DefaultCopy; + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/arch/copy_sm75.hpp b/csrc/mctlass/include/cute/arch/copy_sm75.hpp new file mode 100644 index 0000000..73a12f1 --- /dev/null +++ b/csrc/mctlass/include/cute/arch/copy_sm75.hpp @@ -0,0 +1,322 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +// Config +#if defined(__clang__) && defined(__MACA__) + // ldmatrix PTX instructions added in Clang 14: https://reviews.llvm.org/D107046 + // ... but will not work until Clang 15: + // * https://reviews.llvm.org/D121666 + // * https://reviews.llvm.org/D126846 + #define CUTE_ARCH_CLANG_SUPPORTS_LDSM_SM75 (__clang_major__ >= 15) +#endif + +#if defined(__MXCC__) || defined(__MACACC_RTC__) + // ldmatrix PTX instruction added in CUDA 10.2+ + #define CUTE_ARCH_NVCC_SUPPORTS_LDSM_SM75 ((__CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 2) || __CUDACC_VER_MAJOR__ >= 11) +#endif + +#if ! defined(CUTE_ARCH_LDSM_SM75_SUPPORTED) + #define CUTE_ARCH_LDSM_SM75_SUPPORTED (CUTE_ARCH_NVCC_SUPPORTS_LDSM_SM75 || CUTE_ARCH_CLANG_SUPPORTS_LDSM_SM75) +#endif + +#if ! defined(CUTE_ARCH_LDSM_SM75_ENABLED) + #define CUTE_ARCH_LDSM_SM75_ENABLED (CUTE_ARCH_LDSM_SM75_SUPPORTED) +#endif + +#if 0 + #define CUTE_ARCH_LDSM_SM75_ACTIVATED 1 +#endif + +namespace cute +{ + +struct SM75_U32x1_LDSM_N +{ + using SRegisters = uint128_t[1]; + using DRegisters = uint32_t[1]; + + CUTE_HOST_DEVICE static void + copy(uint128_t const& smem_src, + uint32_t& dst) + { +#if defined(CUTE_ARCH_LDSM_SM75_ACTIVATED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_src); + asm volatile ("ldmatrix.sync.aligned.x1.m8n8.shared.b16 {%0}, [%1];\n" + : "=r"(dst) + : "r"(smem_int_ptr)); +#else + CUTE_RUNTIME_ASSERT("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); +#endif + } +}; + +struct SM75_U32x2_LDSM_N +{ + using SRegisters = uint128_t[1]; + using DRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + copy(uint128_t const& smem_src, + uint32_t& dst0, uint32_t& dst1) + { +#if defined(CUTE_ARCH_LDSM_SM75_ACTIVATED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_src); + asm volatile ("ldmatrix.sync.aligned.x2.m8n8.shared.b16 {%0, %1}, [%2];\n" + : "=r"(dst0), "=r"(dst1) + : "r"(smem_int_ptr)); +#else + CUTE_RUNTIME_ASSERT("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); +#endif + } +}; + +struct SM75_U32x4_LDSM_N +{ + using SRegisters = uint128_t[1]; + using DRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + copy(uint128_t const& smem_src, + uint32_t& dst0, uint32_t& dst1, uint32_t& dst2, uint32_t& dst3) + { +#if defined(CUTE_ARCH_LDSM_SM75_ACTIVATED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_src); + asm volatile ("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(dst0), "=r"(dst1), "=r"(dst2), "=r"(dst3) + : "r"(smem_int_ptr)); + +#elif defined(__MACA_ARCH__) + const int lane_id = __lane_id(); + if (lane_id >= 32) return; + uint64_t sm_ptr = reinterpret_cast(&smem_src); + uint64_t row_ptr[32]; + for (int i = 0; i < 32; ++i) { + row_ptr[i] = __shfl_down_sync(ULONG_MAX, sm_ptr, i - lane_id); + } + + const int row_id = lane_id / 4; + const int col_offset = lane_id % 4; + dst0 = *(reinterpret_cast(row_ptr[0 + row_id]) + col_offset); + dst1 = *(reinterpret_cast(row_ptr[8 + row_id]) + col_offset); + dst2 = *(reinterpret_cast(row_ptr[16 + row_id]) + col_offset); + dst3 = *(reinterpret_cast(row_ptr[24 + row_id]) + col_offset); +#else + CUTE_RUNTIME_ASSERT("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); +#endif + } +}; + +struct SM75_U32x4_LDSM_N_B +{ + using SRegisters = uint128_t[1]; + using DRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + copy(uint128_t const& smem_src, + uint32_t& dst0, uint32_t& dst1, uint32_t& dst2, uint32_t& dst3) + { +#if defined(CUTE_ARCH_LDSM_SM75_ACTIVATED) + + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_src); + asm volatile ("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(dst0), "=r"(dst1), "=r"(dst2), "=r"(dst3) + : "r"(smem_int_ptr)); +#elif defined(__MACA_ARCH__) + const int lane_id = __lane_id(); + if (lane_id >= 32) return; + uint64_t sm_ptr = reinterpret_cast(&smem_src); + uint64_t row_ptr[32]; + for (int i = 0; i < 32; ++i) { + row_ptr[i] = __shfl_down_sync(ULONG_MAX, sm_ptr, i - lane_id); + } + + const int row_id = lane_id / 4; + const int col_offset = lane_id % 4; + dst0 = *(reinterpret_cast(row_ptr[0 + row_id]) + col_offset); + dst1 = *(reinterpret_cast(row_ptr[8 + row_id]) + col_offset); + dst2 = *(reinterpret_cast(row_ptr[16 + row_id]) + col_offset); + dst3 = *(reinterpret_cast(row_ptr[24 + row_id]) + col_offset); +#else + CUTE_RUNTIME_ASSERT("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); +#endif + } +}; + +struct SM75_U16x2_LDSM_T +{ + using SRegisters = uint128_t[1]; + using DRegisters = uint32_t[1]; + + CUTE_HOST_DEVICE static void + copy(uint128_t const& smem_src, + uint32_t& dst) + { +#if defined(CUTE_ARCH_LDSM_SM75_ACTIVATED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_src); + asm volatile ("ldmatrix.sync.aligned.x1.trans.m8n8.shared.b16 {%0}, [%1];\n" + : "=r"(dst) + : "r"(smem_int_ptr)); +#else + CUTE_RUNTIME_ASSERT("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); +#endif + } +}; + +struct SM75_U16x4_LDSM_T +{ + using SRegisters = uint128_t[1]; + using DRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + copy(uint128_t const& smem_src, + uint32_t& dst0, uint32_t& dst1) + { +#if defined(CUTE_ARCH_LDSM_SM75_ACTIVATED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_src); + asm volatile ("ldmatrix.sync.aligned.x2.trans.m8n8.shared.b16 {%0, %1}, [%2];\n" + : "=r"(dst0), "=r"(dst1) + : "r"(smem_int_ptr)); +#else + CUTE_RUNTIME_ASSERT("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); +#endif + } +}; + +struct SM75_U16x8_LDSM_T +{ + using SRegisters = uint128_t[1]; + using DRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + copy(uint128_t const& smem_src, + uint32_t& dst0, uint32_t& dst1, uint32_t& dst2, uint32_t& dst3) + { +#if defined(CUTE_ARCH_LDSM_SM75_ACTIVATED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_src); + asm volatile ("ldmatrix.sync.aligned.x4.trans.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(dst0), "=r"(dst1), "=r"(dst2), "=r"(dst3) + : "r"(smem_int_ptr)); +#elif defined(__MACA_ARCH__) + const int lane_id = __lane_id(); + if (lane_id >= 32) return; + uint64_t sm_ptr = reinterpret_cast(&smem_src); + uint64_t row_ptr[32]; + for (int i = 0; i < 32; ++i) { + row_ptr[i] = __shfl_down_sync(ULONG_MAX, sm_ptr, i - lane_id); + } + + const int row_offset = lane_id % 4 * 2; + const int col_offset = lane_id / 4; + auto low_b16_addr = reinterpret_cast(row_ptr[0 + row_offset]) + col_offset; + auto high_b16_addr = reinterpret_cast(row_ptr[0 + row_offset + 1]) + col_offset; + auto dst_b16 = reinterpret_cast(&dst0); + *dst_b16 = *low_b16_addr; + *(dst_b16 + 1) = *high_b16_addr; + + low_b16_addr = reinterpret_cast(row_ptr[8 + row_offset]) + col_offset; + high_b16_addr = reinterpret_cast(row_ptr[8 + row_offset + 1]) + col_offset; + dst_b16 = reinterpret_cast(&dst1); + *dst_b16 = *low_b16_addr; + *(dst_b16 + 1) = *high_b16_addr; + + low_b16_addr = reinterpret_cast(row_ptr[16 + row_offset]) + col_offset; + high_b16_addr = reinterpret_cast(row_ptr[16 + row_offset + 1]) + col_offset; + dst_b16 = reinterpret_cast(&dst2); + *dst_b16 = *low_b16_addr; + *(dst_b16 + 1) = *high_b16_addr; + + low_b16_addr = reinterpret_cast(row_ptr[24 + row_offset]) + col_offset; + high_b16_addr = reinterpret_cast(row_ptr[24 + row_offset + 1]) + col_offset; + dst_b16 = reinterpret_cast(&dst3); + *dst_b16 = *low_b16_addr; + *(dst_b16 + 1) = *high_b16_addr; +#else + CUTE_RUNTIME_ASSERT("Trying to use ldmatrix without CUTE_ARCH_LDSM_SM75_ACTIVATED."); +#endif + } +}; + +// +// Legacy LDSM interfaces that aren't very useful +// + +template +CUTE_HOST_DEVICE +void +copy_ldsm(uint128_t const* const smem_ptr, + T* rmem_ptr) +{ + uint32_t* reg_ptr = reinterpret_cast(rmem_ptr); + + // if constexpr + if (sizeof(T) == 4) { + SM75_U32x1_LDSM_N::copy(smem_ptr[0], reg_ptr[0]); + } + else if (sizeof(T) == 8) { + SM75_U32x2_LDSM_N::copy(smem_ptr[0], reg_ptr[0], reg_ptr[1]); + } + else if (sizeof(T) == 16) { + SM75_U32x4_LDSM_N::copy(smem_ptr[0], reg_ptr[0], reg_ptr[1], reg_ptr[2], reg_ptr[3]); + } + else { + static_assert(sizeof(T) == 4 || sizeof(T) == 8 || sizeof(T) == 16, "sizeof(T) is not supported"); + } +} + +template +CUTE_HOST_DEVICE +void +copy_ldsm_trans(uint128_t const* const smem_ptr, + T* rmem_ptr) +{ + uint32_t* reg_ptr = reinterpret_cast(rmem_ptr); + + // if constexpr + if (sizeof(T) == 4) { + SM75_U16x2_LDSM_T::copy(smem_ptr[0], reg_ptr[0]); + } + else if (sizeof(T) == 8) { + SM75_U16x4_LDSM_T::copy(smem_ptr[0], reg_ptr[0], reg_ptr[1]); + } + else if (sizeof(T) == 16) { + SM75_U16x8_LDSM_T::copy(smem_ptr[0], reg_ptr[0], reg_ptr[1], reg_ptr[2], reg_ptr[3]); + } + else { + static_assert(sizeof(T) == 4 || sizeof(T) == 8 || sizeof(T) == 16, "sizeof(T) is not supported"); + } +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/arch/copy_sm80.hpp b/csrc/mctlass/include/cute/arch/copy_sm80.hpp new file mode 100644 index 0000000..3c19eb8 --- /dev/null +++ b/csrc/mctlass/include/cute/arch/copy_sm80.hpp @@ -0,0 +1,201 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +#if defined(__MACA_ARCH__) && 0 +# define CUTE_ARCH_CP_ASYNC_SM80_ENABLED +#endif + +#if defined(__MACA_ARCH__) && (__MACA_ARCH__ == 1500) +# define MACA_ARCH_LDS_TRANS_ENABLED +#endif + +namespace cute +{ + +extern __device__ void +_vmov_cp_async(__attribute__((address_space (3)))int8_t *, __attribute__((address_space (1)))int8_t *, int32_t, int32_t) + __asm("llvm.mxc.cp.async.global.to.shared"); + +/// Copy via cp.async with caching at all levels +template +struct SM80_CP_ASYNC_CACHEALWAYS +{ + using SRegisters = TS[1]; + using DRegisters = TD[1]; + + static_assert(sizeof(TS) == sizeof(TD), "cp.async requires sizeof(src_value_type) == sizeof(dst_value_type)"); + static_assert(sizeof(TS) == 4 || sizeof(TS) == 8 || sizeof(TS) == 16, "cp.async sizeof(TS) is not supported"); + + CUTE_HOST_DEVICE static void + copy(TS const& gmem_src, + TD & smem_dst) + { +#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED) + TS const* gmem_ptr = &gmem_src; + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_dst); + asm volatile("cp.async.ca.shared.global [%0], [%1], %2;\n" + :: "r"(smem_int_ptr), + "l"(gmem_ptr), + "n"(sizeof(TS))); +#elif defined(__MACA_ARCH__) + TS const *gmem_ptr = &gmem_src; + TD *smem_ptr = &smem_dst; + *static_cast(smem_ptr) = *static_cast(gmem_ptr); + +#else + CUTE_RUNTIME_ASSERT("Support for cp.async instructions has not been enabled"); +#endif + } +}; + +/// Copy via cp.async with caching at global level +template +struct SM80_CP_ASYNC_CACHEGLOBAL +{ + using SRegisters = TS[1]; + using DRegisters = TD[1]; + + static_assert(sizeof(TS) == sizeof(TD), "cp.async requires sizeof(src_value_type) == sizeof(dst_value_type)"); + static_assert(sizeof(TS) == 4 || sizeof(TS) == 8 || sizeof(TS) == 16, "cp.async sizeof(TS) is not supported"); + + CUTE_HOST_DEVICE static void + copy(TS const& gmem_src, + TD & smem_dst) + { +#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED) + TS const* gmem_ptr = &gmem_src; + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_dst); + asm volatile("cp.async.cg.shared.global [%0], [%1], %2;\n" + :: "r"(smem_int_ptr), + "l"(gmem_ptr), + "n"(sizeof(TS))); +#elif defined(__MACA_ARCH__) + TS const *gmem_ptr = &gmem_src; + TD *smem_ptr = &smem_dst; + *static_cast(smem_ptr) = *static_cast(gmem_ptr); +#else + CUTE_RUNTIME_ASSERT("Support for cp.async instructions has not been enabled"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Copy via cp.async with caching at global level +template +struct MACA_CP_ASYNC_CACHEGLOBAL +{ + using SRegisters = TS[1]; + using DRegisters = TD[1]; + + static_assert(sizeof(TS) == sizeof(TD), "cp.async requires sizeof(src_value_type) == sizeof(dst_value_type)"); + static_assert(sizeof(TS) == 4 || sizeof(TS) == 8 || sizeof(TS) == 16, "cp.async sizeof(TS) is not supported"); + + CUTE_HOST_DEVICE static void + copy(TS const& gmem_src, + TD & smem_dst) + { +#if defined(__MACA_ARCH__) + TS const *gmem_ptr = &gmem_src; + TD *smem_ptr = &smem_dst; + _vmov_cp_async((__attribute__((address_space (3)))int8_t *)(smem_ptr), + (__attribute__((address_space (1)))int8_t *)(gmem_ptr), 0, sizeof(TS)); + //__builtin_mxc_arrive(64); + //__builtin_mxc_barrier(); +#else + CUTE_RUNTIME_ASSERT("Support for cp.async instructions has not been enabled"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Establishes an ordering w.r.t previously issued cp.async instructions. Does not block. +CUTE_HOST_DEVICE +void +cp_async_fence() +{ +#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED) + asm volatile("cp.async.commit_group;\n" ::); +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Blocks until all but N previous cp.async.commit_group operations have committed. +template +CUTE_HOST_DEVICE +void +cp_async_wait() +{ +#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED) + if constexpr (N == 0) { + asm volatile("cp.async.wait_all;\n" ::); + } else { + asm volatile("cp.async.wait_group %0;\n" :: "n"(N)); + } +#endif +} + +template +CUTE_HOST_DEVICE +void +cp_async_wait(Int) +{ + return cp_async_wait(); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +struct MACA_LDS_TRANS_4X16 +{ + using SRegisters = uint64_t[1]; + using DRegisters = uint64_t[1]; + + CUTE_HOST_DEVICE static void + copy(uint64_t& smem_src, + uint64_t& dst) + { +#if defined(MACA_ARCH_LDS_TRANS_ENABLED) + int64_t *smem_src_ptr = reinterpret_cast(&smem_src); + dst = __builtin_mxc_load_shared_trans_4x16_i64(smem_src_ptr); +#else + CUTE_RUNTIME_ASSERT("Trying to use lds_b64_trans_4x16 without MACA_ARCH_LDS_TRANS_ENABLED."); +#endif + } +}; + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/arch/copy_sm90.hpp b/csrc/mctlass/include/cute/arch/copy_sm90.hpp new file mode 100644 index 0000000..a642694 --- /dev/null +++ b/csrc/mctlass/include/cute/arch/copy_sm90.hpp @@ -0,0 +1,225 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +// Config +// #if (defined(__MACA_ARCH__) && (__MACA_ARCH__ >= 900) && (__CUDACC_VER_MAJOR__ >= 12)) +// # define CUTE_ARCH_STSM_SM90_ENABLED +// # define CUTE_ARCH_TMA_SM90_ENABLED +// #endif + +namespace cute +{ + +struct SM90_U32x1_STSM_N +{ + using SRegisters = uint32_t[1]; + using DRegisters = uint128_t[1]; + + CUTE_HOST_DEVICE static void + copy(uint32_t const& src, + uint128_t & smem_dst) + { +#if defined(CUTE_ARCH_STSM_SM90_ENABLED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_dst); + asm volatile ("stmatrix.sync.aligned.x1.m8n8.shared.b16 [%0], {%1};\n" + :: "r"(smem_int_ptr), + "r"(src)); +#else + CUTE_RUNTIME_ASSERT("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); +#endif + } +}; + +struct SM90_U32x2_STSM_N +{ + using SRegisters = uint32_t[2]; + using DRegisters = uint128_t[1]; + + CUTE_HOST_DEVICE static void + copy(uint32_t const& src0, uint32_t const& src1, + uint128_t& smem_dst) + { +#if defined(CUTE_ARCH_STSM_SM90_ENABLED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_dst); + asm volatile ("stmatrix.sync.aligned.x2.m8n8.shared.b16 [%0], {%1, %2};\n" + :: "r"(smem_int_ptr), + "r"(src0), "r"(src1)); +#else + CUTE_RUNTIME_ASSERT("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); +#endif + } +}; + +struct SM90_U32x4_STSM_N +{ + using SRegisters = uint32_t[4]; + using DRegisters = uint128_t[1]; + + CUTE_HOST_DEVICE static void + copy(uint32_t const& src0, uint32_t const& src1, uint32_t const& src2, uint32_t const& src3, + uint128_t& smem_dst) + { +#if defined(CUTE_ARCH_STSM_SM90_ENABLED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_dst); + asm volatile ("stmatrix.sync.aligned.x4.m8n8.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(smem_int_ptr), + "r"(src0), "r"(src1), "r"(src2), "r"(src3)); +#else + CUTE_RUNTIME_ASSERT("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); +#endif + } +}; + +struct SM90_U16x2_STSM_T +{ + using SRegisters = uint32_t[1]; + using DRegisters = uint128_t[1]; + + CUTE_HOST_DEVICE static void + copy(uint32_t const& src, + uint128_t& smem_dst) + { +#if defined(CUTE_ARCH_STSM_SM90_ENABLED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_dst); + asm volatile ("stmatrix.sync.aligned.x1.trans.m8n8.shared.b16 [%0], {%1};\n" + :: "r"(smem_int_ptr), + "r"(src)); +#else + CUTE_RUNTIME_ASSERT("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); +#endif + } +}; + +struct SM90_U16x4_STSM_T +{ + using SRegisters = uint32_t[2]; + using DRegisters = uint128_t[1]; + + CUTE_HOST_DEVICE static void + copy(uint32_t const& src0, uint32_t const& src1, + uint128_t& smem_dst) + { +#if defined(CUTE_ARCH_STSM_SM90_ENABLED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_dst); + asm volatile ("stmatrix.sync.aligned.x2.trans.m8n8.shared.b16 [%0], {%1, %2};\n" + :: "r"(smem_int_ptr), + "r"(src0), "r"(src1)); +#else + CUTE_RUNTIME_ASSERT("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); +#endif + } +}; + +struct SM90_U16x8_STSM_T +{ + using SRegisters = uint32_t[4]; + using DRegisters = uint128_t[1]; + + CUTE_HOST_DEVICE static void + copy(uint32_t const& src0, uint32_t const& src1, uint32_t const& src2, uint32_t const& src3, + uint128_t& smem_dst) + { +#if defined(CUTE_ARCH_STSM_SM90_ENABLED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_dst); + asm volatile ("stmatrix.sync.aligned.x4.trans.m8n8.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(smem_int_ptr), + "r"(src0), "r"(src1), "r"(src2), "r"(src3)); +#else + CUTE_RUNTIME_ASSERT("Trying to use stmatrix without CUTE_ARCH_STSM_SM90_ENABLED."); +#endif + } +}; + +// +// Legacy STSM interfaces that aren't very useful +// + +template +CUTE_HOST_DEVICE +void +copy_stsm(T const* const rmem_ptr, + uint128_t* const smem_ptr) +{ + uint32_t const* reg_ptr = reinterpret_cast(rmem_ptr); + + // if constexpr + if (sizeof(T) == 4) { + SM90_U32x1_STSM_N::copy(reg_ptr[0], smem_ptr[0]); + } + else if (sizeof(T) == 8) { + SM90_U32x2_STSM_N::copy(reg_ptr[0], reg_ptr[1], smem_ptr[0]); + } + else if (sizeof(T) == 16) { + SM90_U32x4_STSM_N::copy(reg_ptr[0], reg_ptr[1], reg_ptr[2], reg_ptr[3], smem_ptr[0]); + } + else { + static_assert(sizeof(T) == 4 || sizeof(T) == 8 || sizeof(T) == 16, "sizeof(T) is not supported"); + } +} + +template +CUTE_HOST_DEVICE +void +copy_stsm_trans(T const* const rmem_ptr, + uint128_t* const smem_ptr) +{ + uint32_t const* reg_ptr = reinterpret_cast(rmem_ptr); + + // if constexpr + if (sizeof(T) == 4) { + SM90_U16x2_STSM_T::copy(reg_ptr[0], smem_ptr[0]); + } + else if (sizeof(T) == 8) { + SM90_U16x4_STSM_T::copy(reg_ptr[0], reg_ptr[1], smem_ptr[0]); + } + else if (sizeof(T) == 16) { + SM90_U16x8_STSM_T::copy(reg_ptr[0], reg_ptr[1], reg_ptr[2], reg_ptr[3], smem_ptr[0]); + } + else { + static_assert(sizeof(T) == 4 || sizeof(T) == 8 || sizeof(T) == 16, "sizeof(T) is not supported"); + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // end namespace cute + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#include +#include + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/cute/arch/copy_sm90_desc.hpp b/csrc/mctlass/include/cute/arch/copy_sm90_desc.hpp new file mode 100644 index 0000000..81b6630 --- /dev/null +++ b/csrc/mctlass/include/cute/arch/copy_sm90_desc.hpp @@ -0,0 +1,201 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#if !defined(__MACACC_RTC__) +#include +#include +#endif + +#include + +#include +#include + +#include +#include +#include // to_Format<[u]intX> +#include // to_Format + +namespace cute +{ + +////////////////////////////////////////////////////////////////////////////////////////////////////// +/// Barriers are 64-bit of user-managed information used in broadly two types syncronization patterns +/// 1) arrive/wait on threads (usage: cp.async and warp-specialized kernels) +/// 2) transaction-based (usage: TMA transaction where a CTA issues one transaction) +////////////////////////////////////////////////////////////////////////////////////////////////////// + +// Initialize barrier present in shared memory +CUTE_HOST_DEVICE +void +initialize_barrier(uint64_t& smem_barrier, // 64 bits user-manged barrier in smem + int thread_count = 1) // Thread count expected to arrive/wait on this barrier +{ +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_barrier); + asm volatile ("mbarrier.init.shared.b64 [%0], %1;\n" + :: "r"(smem_int_ptr), + "r"(thread_count)); +#endif +} + +// Set the number of bytes transfered per transaction +CUTE_HOST_DEVICE +void +set_barrier_transaction_bytes(uint64_t& smem_barrier, // 64 bits user-manged barrier in smem + uint32_t bytes) // Number of bytes transfered by per TMA transaction +{ +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_barrier); + asm volatile ("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;\n" + :: "r"(smem_int_ptr), + "r"(bytes)); +#endif +} + +// Barrier wait +CUTE_HOST_DEVICE +void +wait_barrier(uint64_t& smem_barrier, // 64 bits user-manged barrier in smem + int phase_bit) // Current phase bit the barrier waiting to flip +{ +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_barrier); + asm volatile( + "{\n" + ".reg .pred P1;\n" + "LAB_WAIT:\n" + "mbarrier.try_wait.parity.shared.b64 P1, [%0], %1;\n" + "@P1 bra.uni DONE;\n" + "bra.uni LAB_WAIT;\n" + "DONE:\n" + "}\n" + :: "r"(smem_int_ptr), + "r"(phase_bit)); + +#endif +} + +// Barrier arrive +CUTE_HOST_DEVICE +void +arrive_barrier(uint64_t& smem_barrier) // 64 bits user-manged barrier in smem +{ +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_barrier); + asm volatile( + "{\n" + ".reg .b64 state; \n" + "mbarrier.arrive.shared.b64 state, [%0];\n" + "}\n" + :: "r"(smem_int_ptr)); +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// TMA Descriptor and utilities +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace TMA { + +enum class SmemSwizzleBits : uint8_t { + DISABLE = 0, + B32 = 1, + B64 = 2, + B128 = 3, +}; + +#if !defined(__MACACC_RTC__) +// #if (__CUDACC_VER_MAJOR__ >= 12) +#if 0 + +template +inline CUtensorMapDataType to_CUtensorMapDataType() { + if constexpr (is_same::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT8; } else + if constexpr (is_same::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT8; } else + if constexpr (is_same::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT16; } else + if constexpr (is_same::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT32; } else + if constexpr (is_same::value) { return CU_TENSOR_MAP_DATA_TYPE_UINT64; } else + if constexpr (is_same::value) { return CU_TENSOR_MAP_DATA_TYPE_INT32; } else + if constexpr (is_same::value) { return CU_TENSOR_MAP_DATA_TYPE_INT64; } else + if constexpr (is_same::value) { return CU_TENSOR_MAP_DATA_TYPE_FLOAT16; } else + if constexpr (is_same::value) { return CU_TENSOR_MAP_DATA_TYPE_FLOAT32; } else + if constexpr (is_same::value) { return CU_TENSOR_MAP_DATA_TYPE_FLOAT64; } else + if constexpr (is_same::value) { return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; } else + if constexpr (is_same::value) { return CU_TENSOR_MAP_DATA_TYPE_TFLOAT32; } else + { static_assert(sizeof(T) < 0, "Unknown TMA Format!"); } +} + +inline CUtensorMapSwizzle to_CUtensorMapSwizzle(SmemSwizzleBits const& t) { + switch (t) { + default: assert(false && "Unknown SmemSwizzleBits!"); + case SmemSwizzleBits::DISABLE: return CU_TENSOR_MAP_SWIZZLE_NONE; + case SmemSwizzleBits::B32: return CU_TENSOR_MAP_SWIZZLE_32B; + case SmemSwizzleBits::B64: return CU_TENSOR_MAP_SWIZZLE_64B; + case SmemSwizzleBits::B128: return CU_TENSOR_MAP_SWIZZLE_128B; + } +} + +#endif // (__CUDACC_VER_MAJOR__ >= 12) +#endif // !defined(__MACACC_RTC__) +} // end namespace TMA + +// #if (__CUDACC_VER_MAJOR__ >= 12) && !defined(__MACACC_RTC__) +#if 0 +using TmaDescriptor = CUtensorMap; +#else +using TmaDescriptor = struct { char bytes[128]; }; +#endif +//////////////////////////////////////////////////////////////////////////////////////////////////// +/// Initiates a TensorMap Prefetch +//////////////////////////////////////////////////////////////////////////////////////////////////// + +CUTE_HOST_DEVICE +void +prefetch_tma_descriptor(TmaDescriptor const* desc_ptr) +{ +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + // Prefetch TMA Descriptor using generic addressing (i.e. no specific state space: const or param) + asm volatile ( + "prefetch.tensormap [%0];" + : + : "l"(gmem_int_desc) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use TMA Descriptor Prefetch without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif +} + +/////////////////////////////////////////////////////////////////////////////// + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/arch/copy_sm90_tma.hpp b/csrc/mctlass/include/cute/arch/copy_sm90_tma.hpp new file mode 100644 index 0000000..412754c --- /dev/null +++ b/csrc/mctlass/include/cute/arch/copy_sm90_tma.hpp @@ -0,0 +1,861 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include + +namespace cute +{ + +//////////////////////////////////////////////////////////////////////////////////////////////////// +/// TMA_LOAD : Initiates a TMA copy from global memory to shared memory +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct SM90_TMA_LOAD_1D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& crd0) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3}], [%2];" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "r"(crd0) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_2D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4}], [%2];" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "r"(crd0), "r"(crd1) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_3D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5}], [%2];" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "r"(crd0), "r"(crd1), "r"(crd2) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_4D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6}], [%2];" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_5D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3, int32_t const& crd4) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2];" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "r"(crd4) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& crd0) + { + return SM90_TMA_LOAD_1D::copy(desc_ptr, smem_mbar, smem_ptr, crd0); + } + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1) + { + return SM90_TMA_LOAD_2D::copy(desc_ptr, smem_mbar, smem_ptr, crd0, crd1); + } + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2) + { + return SM90_TMA_LOAD_3D::copy(desc_ptr, smem_mbar, smem_ptr, crd0, crd1, crd2); + } + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3) + { + return SM90_TMA_LOAD_4D::copy(desc_ptr, smem_mbar, smem_ptr, crd0, crd1, crd2, crd3); + } + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3, int32_t const& crd4) + { + return SM90_TMA_LOAD_5D::copy(desc_ptr, smem_mbar, smem_ptr, crd0, crd1, crd2, crd3, crd4); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +/// TMA_LOAD im2col: Initiates a TMA copy, in im2col mode, from global memory to shared memory +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct SM90_TMA_LOAD_IM2COL_3D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_n, + uint16_t const& offset_w) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + // Copy from global to shared::cluster. + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5}], [%2], {%6};" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "r"(coord_c), "r"(coord_w), "r"(coord_n), + "h"(offset_w) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_IM2COL_4D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_n, + uint16_t const& offset_w, + uint16_t const& offset_h) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + // Copy from global to shared::cluster. + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6}], [%2], {%7, %8};" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_n), + "h"(offset_w), "h"(offset_h) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_IM2COL_5D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_d, int32_t const& coord_n, + uint16_t const& offset_w, + uint16_t const& offset_h, + uint16_t const& offset_d) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + // Copy from global to shared::cluster. + asm volatile ( + "cp.async.bulk.tensor.5d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], {%8, %9, %10};" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_d), "r"(coord_n), + "h"(offset_w), "h"(offset_h), "h"(offset_d) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_IM2COL +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_n, + uint16_t const& offset_w) + { + return SM90_TMA_LOAD_IM2COL_3D::copy(desc_ptr, smem_mbar, smem_ptr, + coord_c, coord_w, coord_n, + offset_w); + } + + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_n, + uint16_t const& offset_w, + uint16_t const& offset_h) + { + return SM90_TMA_LOAD_IM2COL_4D::copy(desc_ptr, smem_mbar, smem_ptr, + coord_c, coord_w, coord_h, coord_n, + offset_w, offset_h); + } + + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, + int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_d, int32_t const& coord_n, + uint16_t const& offset_w, + uint16_t const& offset_h, + uint16_t const& offset_d) + { + return SM90_TMA_LOAD_IM2COL_5D::copy(desc_ptr, smem_mbar, smem_ptr, + coord_c, coord_w, coord_h, coord_d, coord_n, + offset_w, offset_h, offset_d); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +/// TMA_LOAD_MULTICAST: Initiates a TMA copy from global memory to shared memory +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct SM90_TMA_LOAD_MULTICAST_1D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask, + void const* const smem_ptr, + int32_t const& crd0) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4}], [%2], %3;" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "h"(multicast_mask), + "r"(crd0) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_MULTICAST_2D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5}], [%2], %3;" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "h"(multicast_mask), + "r"(crd0), "r"(crd1) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_MULTICAST_3D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6}], [%2], %3;" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "h"(multicast_mask), + "r"(crd0), "r"(crd1), "r"(crd2) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_MULTICAST_4D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6, %7}], [%2], %3;" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "h"(multicast_mask), + "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_MULTICAST_5D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3, int32_t const& crd4) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%4, %5, %6, %7, %8}], [%2], %3;" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "h"(multicast_mask), + "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "r"(crd4) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_MULTICAST +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask, + void const* const smem_ptr, + int32_t const& crd0) + { + return SM90_TMA_LOAD_MULTICAST_1D::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0); + } + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1) + { + return SM90_TMA_LOAD_MULTICAST_2D::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0, crd1); + } + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2) + { + return SM90_TMA_LOAD_MULTICAST_3D::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0, crd1, crd2); + } + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3) + { + return SM90_TMA_LOAD_MULTICAST_4D::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0, crd1, crd2, crd3); + } + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, uint16_t multicast_mask, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3, int32_t const& crd4) + { + return SM90_TMA_LOAD_MULTICAST_5D::copy(desc_ptr, smem_mbar, multicast_mask, smem_ptr, crd0, crd1, crd2, crd3, crd4); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +/// TMA_LOAD_MULTICAST im2col: Initiates a TMA copy, in im2col mode, from global memory to shared memory +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct SM90_TMA_LOAD_IM2COL_MULTICAST_3D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + uint16_t const& multicast_mask, + void const* const smem_ptr, + int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_n, + uint16_t const& offset_w) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + // Copy from global to shared::cluster. + asm volatile ( + "cp.async.bulk.tensor.3d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes.multicast::cluster" + " [%0], [%1, {%3, %4, %5}], [%2], {%6}, %7;" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "r"(coord_c), "r"(coord_w), "r"(coord_n), + "h"(offset_w), + "h"(multicast_mask) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_IM2COL_MULTICAST_4D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + uint16_t const& multicast_mask, + void const* const smem_ptr, + int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_n, + uint16_t const& offset_w, + uint16_t const& offset_h) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + // Copy from global to shared::cluster. + asm volatile ( + "cp.async.bulk.tensor.4d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6}], [%2], {%7, %8}, %9;" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_n), + "h"(offset_w), "h"(offset_h), + "h"(multicast_mask) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_IM2COL_MULTICAST_5D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + uint16_t const& multicast_mask, + void const* const smem_ptr, + int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_d, int32_t const& coord_n, + uint16_t const& offset_w, + uint16_t const& offset_h, + uint16_t const& offset_d) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + // Copy from global to shared::cluster. + asm volatile ( + "cp.async.bulk.tensor.5d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes" + " [%0], [%1, {%3, %4, %5, %6, %7}], [%2], {%8, %9, %10}, %11;" + : + : "r"(smem_int_ptr), "l"(gmem_int_desc), "r"(smem_int_mbar), + "r"(coord_c), "r"(coord_w), "r"(coord_h), "r"(coord_d), "r"(coord_n), + "h"(offset_w), "h"(offset_h), "h"(offset_d), + "h"(multicast_mask) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_LOAD_IM2COL_MULTICAST +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + uint16_t const& multicast_mask, + void const* const smem_ptr, + int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_n, + uint16_t const& offset_w) + { + return SM90_TMA_LOAD_IM2COL_MULTICAST_3D::copy(desc_ptr, smem_mbar, + multicast_mask, smem_ptr, + coord_c, coord_w, coord_n, + offset_w); + } + + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + uint16_t const& multicast_mask, + void const* const smem_ptr, + int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_n, + uint16_t const& offset_w, + uint16_t const& offset_h) + { + return SM90_TMA_LOAD_IM2COL_MULTICAST_4D::copy(desc_ptr, smem_mbar, + multicast_mask, smem_ptr, + coord_c, coord_w, coord_h, coord_n, + offset_w, offset_h); + } + + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, uint64_t& smem_mbar, + uint16_t const& multicast_mask, + void const* const smem_ptr, + int32_t const& coord_c, int32_t const& coord_w, int32_t const& coord_h, int32_t const& coord_d, int32_t const& coord_n, + uint16_t const& offset_w, + uint16_t const& offset_h, + uint16_t const& offset_d) + { + return SM90_TMA_LOAD_IM2COL_MULTICAST_5D::copy(desc_ptr, smem_mbar, + multicast_mask, smem_ptr, + coord_c, coord_w, coord_h, coord_d, coord_n, + offset_w, offset_h, offset_d); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +/// TMA_STORE : Initiates a TMA copy from shared memory to global memory +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct SM90_TMA_STORE_1D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, + void const* const smem_ptr, + int32_t const& crd0) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.1d.global.shared::cta.bulk_group [%0, {%2}], [%1];" + : + : "l"(gmem_int_desc), "r"(smem_int_ptr), + "r"(crd0) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_STORE_2D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.2d.global.shared::cta.bulk_group [%0, {%2, %3}], [%1];" + : + : "l"(gmem_int_desc), "r"(smem_int_ptr), + "r"(crd0), "r"(crd1) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_STORE_3D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.3d.global.shared::cta.bulk_group [%0, {%2, %3, %4}], [%1];" + : + : "l"(gmem_int_desc), "r"(smem_int_ptr), + "r"(crd0), "r"(crd1), "r"(crd2) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_STORE_4D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.4d.global.shared::cta.bulk_group [%0, {%2, %3, %4, %5}], [%1];" + : + : "l"(gmem_int_desc), "r"(smem_int_ptr), + "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_STORE_5D +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3, int32_t const& crd4) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint64_t gmem_int_desc = reinterpret_cast(desc_ptr); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile ( + "cp.async.bulk.tensor.5d.global.shared::cta.bulk_group [%0, {%2, %3, %4, %5, %6}], [%1];" + : + : "l"(gmem_int_desc), "r"(smem_int_ptr), + "r"(crd0), "r"(crd1), "r"(crd2), "r"(crd3), "r"(crd4) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_TMA_STORE +{ + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, + void const* const smem_ptr, + int32_t const& crd0) + { + return SM90_TMA_STORE_1D::copy(desc_ptr, smem_ptr, crd0); + } + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1) + { + return SM90_TMA_STORE_2D::copy(desc_ptr, smem_ptr, crd0, crd1); + } + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2) + { + return SM90_TMA_STORE_3D::copy(desc_ptr, smem_ptr, crd0, crd1, crd2); + } + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3) + { + return SM90_TMA_STORE_4D::copy(desc_ptr, smem_ptr, crd0, crd1, crd2, crd3); + } + CUTE_HOST_DEVICE static void + copy(void const* const desc_ptr, + void const* const smem_ptr, + int32_t const& crd0, int32_t const& crd1, int32_t const& crd2, int32_t const& crd3, int32_t const& crd4) + { + return SM90_TMA_STORE_5D::copy(desc_ptr, smem_ptr, crd0, crd1, crd2, crd3, crd4); + } +}; + +// Indicate arrival of warp issuing TMA_STORE +CUTE_HOST_DEVICE static void +tma_store_arrive() { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + asm volatile("cp.async.bulk.commit_group;"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif +} + +// Wait on prior N (Count) TMA_STORE instructions to complete +template +CUTE_HOST_DEVICE static void +tma_store_wait() { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + asm volatile( + "cp.async.bulk.wait_group.read %0;" + : + : "n"(Count) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use tma without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +/// BULK_COPY : Copy a bulk of memory between shared memory and global memory +//////////////////////////////////////////////////////////////////////////////////////////////////// + +struct SM90_BULK_COPY_G2S +{ + CUTE_HOST_DEVICE static void + copy(void const* const gmem_ptr, uint64_t& smem_mbar, + void const* const smem_ptr, int32_t load_bytes) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint32_t smem_int_mbar = cast_smem_ptr_to_uint(&smem_mbar); + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];\n" + : + : "r"(smem_int_ptr), "l"(gmem_ptr), "r"(load_bytes), "r"(smem_int_mbar) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use BULK_COPY without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_BULK_COPY_S2G +{ + CUTE_HOST_DEVICE static void + copy(void const* const smem_ptr, + void const* const gmem_ptr, int32_t store_bytes) + { +#if defined(CUTE_ARCH_TMA_SM90_ENABLED) + uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr); + asm volatile("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;\n" + : + : "l"(gmem_ptr), "r"(smem_int_ptr), "r"(store_bytes) + : "memory"); +#else + CUTE_RUNTIME_ASSERT("Trying to use BULK_COPY without CUTE_ARCH_TMA_SM90_ENABLED."); +#endif + } +}; + +struct SM90_BULK_COPY_AUTO {}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/arch/mma.hpp b/csrc/mctlass/include/cute/arch/mma.hpp new file mode 100644 index 0000000..1c1058f --- /dev/null +++ b/csrc/mctlass/include/cute/arch/mma.hpp @@ -0,0 +1,64 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +namespace cute +{ + +// +// Direct FMA for any type +// + +template +struct UniversalFMA +{ + using DRegisters = D[1]; + using ARegisters = A[1]; + using BRegisters = B[1]; + using CRegisters = C[1]; + + CUTE_HOST_DEVICE static constexpr void + fma(D & d, + A const& a, + B const& b, + C const& c) + { + // Forward to an ADL/cute free function for these types + using cute::fma; + fma(d, a, b, c); + } +}; + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/arch/mma_sm75.hpp b/csrc/mctlass/include/cute/arch/mma_sm75.hpp new file mode 100644 index 0000000..4e3cb72 --- /dev/null +++ b/csrc/mctlass/include/cute/arch/mma_sm75.hpp @@ -0,0 +1,120 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +// Config +// #if ((__CUDACC_VER_MAJOR__ > 10) || (__CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 2)) +// # define CUTE_ARCH_MMA_SM75_SUPPORTED +// # if (defined(__MACA_ARCH__) && (__MACA_ARCH__ >= 750)) +// # define CUTE_ARCH_MMA_SM75_ENABLED +// # endif +// #endif + +namespace cute +{ + +// +// SM75 MMA 1688 F16F16F32 +// + +struct SM75_16x8x8_F32F16F16F32_TN +{ + using DRegisters = float[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = float[4]; + + // Register asm fma + CUTE_HOST_DEVICE static void + fma(float & d0, float & d1, float & d2, float & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + float const& c0, float const& c1, float const& c2, float const& c3) + { +#if defined(CUTE_ARCH_MMA_SM75_ENABLED) + asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32" + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "f"(c0), "f"(c1), "f"(c2), "f"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM75_16x8x8_F32F16F16F32_TN without CUTE_ARCH_MMA_SM75_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// +// SM75 MMA 8816 S8S8S32 +// + +struct SM75_8x8x16_S32S8S8S32_TN +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + // Register asm fma + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM75_ENABLED) + asm volatile("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32" + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM75_8x8x16_S32S8S8S32_TN without CUTE_ARCH_MMA_SM75_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/arch/mma_sm80.hpp b/csrc/mctlass/include/cute/arch/mma_sm80.hpp new file mode 100644 index 0000000..04bbbea --- /dev/null +++ b/csrc/mctlass/include/cute/arch/mma_sm80.hpp @@ -0,0 +1,2714 @@ + /************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include +#include +#include + +// Config +#if defined(__MACA_ARCH__) +// # define CUTE_ARCH_MMA_SM80_ENABLED +#endif + +namespace cute { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x8 TN +struct SM80_16x8x8_F16F16F16F16_TN +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3}," + "{%4}," + "{%5, %6};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x8_F16F16F16F16_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x16 TN +struct SM80_16x8x16_F16F16F16F16_TN +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x16_F16F16F16F16_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x8 TN +struct SM80_16x8x8_F32F16F16F32_TN +{ + using DRegisters = float[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = float[4]; + + CUTE_HOST_DEVICE static void + fma(float & d0, float & d1, float & d2, float & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + float const & c0, float const & c1, float const & c2, float const & c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "f"(c0), "f"(c1), "f"(c2), "f"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x8_F32F16F16F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x16x16 TN +struct MACA_16x16x16_F32F16F16F32 +{ + using DRegisters = float[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[2]; + using CRegisters = float[4]; + using VectorType = __NATIVE_VECTOR__(2, uint32_t); + + CUTE_HOST_DEVICE static void + fma(float & d0, float & d1, float & d2, float & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& b0, uint32_t const& b1, + float const & c0, float const & c1, float const & c2, float const & c3) + { + + VectorType a = {a0, a1}; + VectorType b = {b0, b1}; + + auto result = __builtin_mxc_mma_16x16x16f16(b, a, {c0, c1, c2, c3}); + d0 = result[0]; + d1 = result[1]; + d2 = result[2]; + d3 = result[3]; + + } +}; + +// MMA 16x16x32 TN +struct MACA_16x16x32_F32F16F16F32 +{ + using DRegisters = float[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[4]; + using CRegisters = float[4]; + using VectorType = __NATIVE_VECTOR__(2, uint32_t); + + CUTE_HOST_DEVICE static void + fma(float & d0, float & d1, float & d2, float & d3, + uint32_t const& a0, uint32_t const& a1, // A, k0 + uint32_t const& a2, uint32_t const& a3, // A, k1 + uint32_t const& b0, uint32_t const& b1, // B, k0 + uint32_t const& b2, uint32_t const& b3, // B, k1 + float const & c0, float const & c1, float const & c2, float const & c3) + { + + VectorType a = {a0, a1}; + VectorType b = {b0, b1}; + + // first mma , A(k0) * B(k0) + auto acc = __builtin_mxc_mma_16x16x16f16(b, a, {c0, c1, c2, c3}); + + // second mma , A(k1) * B(k1), accum + a = {a2, a3}; + b = {b2, b3}; + acc = __builtin_mxc_mma_16x16x16f16(b, a, {acc[0], acc[1], acc[2], acc[3]}); + + d0 = acc[0]; + d1 = acc[1]; + d2 = acc[2]; + d3 = acc[3]; + } +}; + +// MMA 16x64x16 TN +// use for lds4x4 + perm4x4 +struct MACA_16x64x16_F32F16F16F32 +{ + using DRegisters = float[16]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[8]; + using CRegisters = float[16]; + using VectorType = __NATIVE_VECTOR__(2, uint32_t); + + CUTE_HOST_DEVICE static void + fma(float & d0, float & d1, float & d2, float & d3, + float & d4, float & d5, float & d6, float & d7, + float & d8, float & d9, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + uint32_t const& a0, uint32_t const& a1, // A, k0 + uint32_t const& b0, uint32_t const& b1, // B, m0 + uint32_t const& b2, uint32_t const& b3, // B, m1 + uint32_t const& b4, uint32_t const& b5, // B, m2 + uint32_t const& b6, uint32_t const& b7, // B, m3 + float const & c0, float const & c1, float const & c2, float const & c3, + float const & c4, float const & c5, float const & c6, float const & c7, + float const & c8, float const & c9, float const & c10, float const & c11, + float const & c12, float const & c13, float const & c14, float const & c15) + { + + VectorType a = {a0, a1}; + VectorType b = {b0, b1}; + + auto acc0 = __builtin_mxc_mma_16x16x16f16(b, a, {c0, c1, c2, c3}); + d0 = acc0[0]; + d1 = acc0[1]; + d2 = acc0[2]; + d3 = acc0[3]; + + b = {b2, b3}; + auto acc1 = __builtin_mxc_mma_16x16x16f16(b, a, {c4, c5, c6, c7}); + d4 = acc1[0]; + d5 = acc1[1]; + d6 = acc1[2]; + d7 = acc1[3]; + + b = {b4, b5}; + auto acc2 = __builtin_mxc_mma_16x16x16f16(b, a, {c8, c9, c10, c11}); + d8 = acc2[0]; + d9 = acc2[1]; + d10 = acc2[2]; + d11 = acc2[3]; + + b = {b6, b7}; + auto acc3 = __builtin_mxc_mma_16x16x16f16(b, a, {c12, c13, c14, c15}); + d12 = acc3[0]; + d13 = acc3[1]; + d14 = acc3[2]; + d15 = acc3[3]; + + } +}; + +// MMA 16x16x16 TN +struct MACA_16x16x16_I32I8I8I32 { + using DRegisters = int32_t[4]; + using ARegisters = int32_t[1]; + using BRegisters = int32_t[1]; + using CRegisters = int32_t[4]; + + CUTE_HOST_DEVICE static void fma( + int32_t &d0, int32_t &d1, int32_t &d2, int32_t &d3, + int32_t const &a0, + int32_t const &b0, + int32_t const &c0, int32_t const &c1, int32_t const &c2, int32_t const &c3) { + +#if defined(CUTE_MACA_XCORE1000_ENABLED) + auto result = __builtin_mxc_mma_16x16x16i8(b0, a0, {c0, c1, c2, c3}); + d0 = result[0]; + d1 = result[1]; + d2 = result[2]; + d3 = result[3]; +#else + CUTE_RUNTIME_ASSERT("Attempting to use MACA_16x16x16_I32I8I8I32 without CUTE_MACA_XCORE1000_ENABLED"); +#endif + + } +}; + +// MMA 16x16x32 TN +struct MACA_16x16x32_I32I8I8I32 { + using DRegisters = int32_t[4]; + using ARegisters = int32_t[2]; + using BRegisters = int32_t[2]; + using CRegisters = int32_t[4]; + + CUTE_HOST_DEVICE static void fma( + int32_t &d0, int32_t &d1, int32_t &d2, int32_t &d3, + int32_t const &a0, int32_t const &a1, + int32_t const &b0, int32_t const &b1, + int32_t const &c0, int32_t const &c1, int32_t const &c2, int32_t const &c3) { + +#if defined(CUTE_MACA_XCORE1500_ENABLED) + auto result = __builtin_mxc_mma_16x16x32i8({b0, b1}, {a0, a1}, {c0, c1, c2, c3}); + d0 = result[0]; + d1 = result[1]; + d2 = result[2]; + d3 = result[3]; +#else + CUTE_RUNTIME_ASSERT("Attempting to use MACA_16x16x32_I32I8I8I32 without CUTE_MACA_XCORE1500_ENABLED"); +#endif + } +}; +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x16 TN +struct SM80_16x8x16_F32F16F16F32_TN +{ + using DRegisters = float[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = float[4]; + + CUTE_HOST_DEVICE static void + fma(float & d0, float & d1, float & d2, float & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + float const & c0, float const & c1, float const & c2, float const & c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "f"(c0), "f"(c1), "f"(c2), "f"(c3)); +#elif defined(__MACA_ARCH__) + + uint32_t aa[2]; + uint32_t bb[2]; + const int lane_id = __lane_id(); + int delta = (lane_id % 32 % 8 * 4 + lane_id % 32 / 16 * 2) - lane_id; + + // shuffle A + uint32_t tmp0 = __shfl_down_sync(ULONG_MAX, a0, delta); + uint32_t tmp0_1 = __shfl_down_sync(ULONG_MAX, a0, delta + 1); + uint32_t tmp1 = __shfl_down_sync(ULONG_MAX, a1, delta); + uint32_t tmp1_1 = __shfl_down_sync(ULONG_MAX, a1, delta + 1); + uint32_t tmp2 = __shfl_down_sync(ULONG_MAX, a2, delta); + uint32_t tmp2_1 = __shfl_down_sync(ULONG_MAX, a2, delta + 1); + uint32_t tmp3 = __shfl_down_sync(ULONG_MAX, a3, delta); + uint32_t tmp3_1 = __shfl_down_sync(ULONG_MAX, a3, delta + 1); + bool up_down_flag = lane_id % 16 / 8; + if (lane_id < 32) { + aa[0] = up_down_flag ? tmp1 : tmp0; + aa[1] = up_down_flag ? tmp1_1 : tmp0_1; + } else if (lane_id < 64) { + aa[0] = up_down_flag ? tmp3 : tmp2; + aa[1] = up_down_flag ? tmp3_1 : tmp2_1; + } + + // shuffle B + tmp0 = __shfl_down_sync(ULONG_MAX, b0, delta); + tmp0_1 = __shfl_down_sync(ULONG_MAX, b0, delta + 1); + tmp1 = __shfl_down_sync(ULONG_MAX, b1, delta); + tmp1_1 = __shfl_down_sync(ULONG_MAX, b1, delta + 1); + if (lane_id < 32) { + bb[0] = tmp0; + bb[1] = tmp0_1; + } else if (lane_id < 64) { + bb[0] = tmp1; + bb[1] = tmp1_1; + } + __half2 const *ha0 = reinterpret_cast<__half2 const *>(&a0); + __half2 const *ha1 = reinterpret_cast<__half2 const *>(&a1); + __half2 const *ha2 = reinterpret_cast<__half2 const *>(&a2); + __half2 const *ha3 = reinterpret_cast<__half2 const *>(&a3); + __half2 const *hb0 = reinterpret_cast<__half2 const *>(&b0); + __half2 const *hb1 = reinterpret_cast<__half2 const *>(&b1); + __half2 const *haa0 = reinterpret_cast<__half2 const *>(&aa[0]); + __half2 const *haa1 = reinterpret_cast<__half2 const *>(&aa[1]); + __half2 const *hbb0 = reinterpret_cast<__half2 const *>(&bb[0]); + __half2 const *hbb1 = reinterpret_cast<__half2 const *>(&bb[1]); + + // shuffle C + delta = (lane_id % 8 / 2 + lane_id % 32 / 16 * 16) - lane_id; + float tmp_c0 = __shfl_down_sync(ULONG_MAX, c0, delta); + float tmp_c0_4 = __shfl_down_sync(ULONG_MAX, c0, delta + 4); + float tmp_c0_8 = __shfl_down_sync(ULONG_MAX, c0, delta + 8); + float tmp_c0_12 = __shfl_down_sync(ULONG_MAX, c0, delta + 12); + float tmp_c1 = __shfl_down_sync(ULONG_MAX, c1, delta); + float tmp_c1_4 = __shfl_down_sync(ULONG_MAX, c1, delta + 4); + float tmp_c1_8 = __shfl_down_sync(ULONG_MAX, c1, delta + 8); + float tmp_c1_12 = __shfl_down_sync(ULONG_MAX, c1, delta + 12); + float tmp_c2 = __shfl_down_sync(ULONG_MAX, c2, delta); + float tmp_c2_4 = __shfl_down_sync(ULONG_MAX, c2, delta + 4); + float tmp_c2_8 = __shfl_down_sync(ULONG_MAX, c2, delta + 8); + float tmp_c2_12 = __shfl_down_sync(ULONG_MAX, c2, delta + 12); + float tmp_c3 = __shfl_down_sync(ULONG_MAX, c3, delta); + float tmp_c3_4 = __shfl_down_sync(ULONG_MAX, c3, delta + 4); + float tmp_c3_8 = __shfl_down_sync(ULONG_MAX, c3, delta + 8); + float tmp_c3_12 = __shfl_down_sync(ULONG_MAX, c3, delta + 12); + float cc0, cc1, cc2, cc3; + if (lane_id < 32) { + if (lane_id % 2) { + cc0 = tmp_c1; + cc1 = tmp_c1_4; + cc2 = tmp_c1_8; + cc3 = tmp_c1_12; + } else { + cc0 = tmp_c0; + cc1 = tmp_c0_4; + cc2 = tmp_c0_8; + cc3 = tmp_c0_12; + } + } else if (lane_id < 64) { + if (lane_id % 2) { + cc0 = tmp_c3; + cc1 = tmp_c3_4; + cc2 = tmp_c3_8; + cc3 = tmp_c3_12; + } else { + cc0 = tmp_c2; + cc1 = tmp_c2_4; + cc2 = tmp_c2_8; + cc3 = tmp_c2_12; + } + } + + auto result = __builtin_mxc_mma_16x16x16f16({static_cast<__fp16>(float(haa0->x)), static_cast<__fp16>(float(haa0->y)), + static_cast<__fp16>(float(haa1->x)), static_cast<__fp16>(float(haa1->y))}, + {static_cast<__fp16>(float(hbb0->x)), static_cast<__fp16>(float(hbb0->y)), + static_cast<__fp16>(float(hbb1->x)), static_cast<__fp16>(float(hbb1->y))}, + {cc0, cc1, cc2, cc3}); + + delta = (lane_id % 4 * 2 + lane_id / 16 * 16) - lane_id; + float tmp_d0 = __shfl_down_sync(ULONG_MAX, result[0], delta); + float tmp_d1 = __shfl_down_sync(ULONG_MAX, result[1], delta); + float tmp_d2 = __shfl_down_sync(ULONG_MAX, result[2], delta); + float tmp_d3 = __shfl_down_sync(ULONG_MAX, result[3], delta); + float tmp_d0_1 = __shfl_down_sync(ULONG_MAX, result[0], delta + 1); + float tmp_d1_1 = __shfl_down_sync(ULONG_MAX, result[1], delta + 1); + float tmp_d2_1 = __shfl_down_sync(ULONG_MAX, result[2], delta + 1); + float tmp_d3_1 = __shfl_down_sync(ULONG_MAX, result[3], delta + 1); + float tmp_d0_32 = __shfl_down_sync(ULONG_MAX, result[0], delta + 32); + float tmp_d1_32 = __shfl_down_sync(ULONG_MAX, result[1], delta + 32); + float tmp_d2_32 = __shfl_down_sync(ULONG_MAX, result[2], delta + 32); + float tmp_d3_32 = __shfl_down_sync(ULONG_MAX, result[3], delta + 32); + float tmp_d0_32_1 = __shfl_down_sync(ULONG_MAX, result[0], delta + 32 + 1); + float tmp_d1_32_1 = __shfl_down_sync(ULONG_MAX, result[1], delta + 32 + 1); + float tmp_d2_32_1 = __shfl_down_sync(ULONG_MAX, result[2], delta + 32 + 1); + float tmp_d3_32_1 = __shfl_down_sync(ULONG_MAX, result[3], delta + 32 + 1); + + if (lane_id < 32) { + if (lane_id % 16 / 4 == 0) { + d0 = tmp_d0; + d1 = tmp_d0_1; + d2 = tmp_d0_32; + d3 = tmp_d0_32_1; + } else if (lane_id % 16 / 4 == 1) { + d0 = tmp_d1; + d1 = tmp_d1_1; + d2 = tmp_d1_32; + d3 = tmp_d1_32_1; + } else if (lane_id % 16 / 4 == 2) { + d0 = tmp_d2; + d1 = tmp_d2_1; + d2 = tmp_d2_32; + d3 = tmp_d2_32_1; + } else if (lane_id % 16 / 4 == 3) { + d0 = tmp_d3; + d1 = tmp_d3_1; + d2 = tmp_d3_32; + d3 = tmp_d3_32_1; + } + } + +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x16_F32F16F16F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x8 TN +struct SM80_16x8x8_F32BF16BF16F32_TN +{ + using DRegisters = float[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = float[4]; + + CUTE_HOST_DEVICE static void + fma(float & d0, float & d1, float & d2, float & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + float const & c0, float const & c1, float const & c2, float const & c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "f"(c0), "f"(c1), "f"(c2), "f"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x8_F32BF16BF16F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x16x16 TN +struct MACA_16x16x16_F32BF16BF16F32 +{ + using DRegisters = float[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[2]; + using CRegisters = float[4]; + using VectorType = __NATIVE_VECTOR__(2, uint32_t); + + CUTE_HOST_DEVICE static void + fma(float & d0, float & d1, float & d2, float & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& b0, uint32_t const& b1, + float const & c0, float const & c1, float const & c2, float const & c3) + { + + VectorType a = {a0, a1}; + VectorType b = {b0, b1}; + + auto result = __builtin_mxc_mma_16x16x16bf16(b, a, {c0, c1, c2, c3}); + + d0 = result[0]; + d1 = result[1]; + d2 = result[2]; + d3 = result[3]; + } +}; + +// MMA 16x16x32 TN +struct MACA_16x16x32_F32BF16BF16F32 +{ + using DRegisters = float[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[4]; + using CRegisters = float[4]; + using VectorType = __NATIVE_VECTOR__(2, uint32_t); + + CUTE_HOST_DEVICE static void + fma(float & d0, float & d1, float & d2, float & d3, + uint32_t const& a0, uint32_t const& a1, // A, k0 + uint32_t const& a2, uint32_t const& a3, // A, k1 + uint32_t const& b0, uint32_t const& b1, // B, k0 + uint32_t const& b2, uint32_t const& b3, // B, k1 + float const & c0, float const & c1, float const & c2, float const & c3) + { + + VectorType a = {a0, a1}; + VectorType b = {b0, b1}; + + // first mma , A(k0) * B(k0) + auto acc = __builtin_mxc_mma_16x16x16bf16(b, a, {c0, c1, c2, c3}); + + // second mma , A(k1) * B(k1), accum + a = {a2, a3}; + b = {b2, b3}; + acc = __builtin_mxc_mma_16x16x16bf16(b, a, {acc[0], acc[1], acc[2], acc[3]}); + + d0 = acc[0]; + d1 = acc[1]; + d2 = acc[2]; + d3 = acc[3]; + } +}; + +// MMA 16x64x16 TN +// use for lds4x4 + perm4x4 +struct MACA_16x64x16_F32BF16BF16F32 +{ + using DRegisters = float[16]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[8]; + using CRegisters = float[16]; + using VectorType = __NATIVE_VECTOR__(2, uint32_t); + + CUTE_HOST_DEVICE static void + fma(float & d0, float & d1, float & d2, float & d3, + float & d4, float & d5, float & d6, float & d7, + float & d8, float & d9, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + uint32_t const& a0, uint32_t const& a1, // A, k0 + uint32_t const& b0, uint32_t const& b1, // B, m0 + uint32_t const& b2, uint32_t const& b3, // B, m1 + uint32_t const& b4, uint32_t const& b5, // B, m2 + uint32_t const& b6, uint32_t const& b7, // B, m3 + float const & c0, float const & c1, float const & c2, float const & c3, + float const & c4, float const & c5, float const & c6, float const & c7, + float const & c8, float const & c9, float const & c10, float const & c11, + float const & c12, float const & c13, float const & c14, float const & c15) + { + + VectorType a = {a0, a1}; + VectorType b = {b0, b1}; + + auto acc0 = __builtin_mxc_mma_16x16x16bf16(b, a, {c0, c1, c2, c3}); + d0 = acc0[0]; + d1 = acc0[1]; + d2 = acc0[2]; + d3 = acc0[3]; + + b = {b2, b3}; + auto acc1 = __builtin_mxc_mma_16x16x16bf16(b, a, {c4, c5, c6, c7}); + d4 = acc1[0]; + d5 = acc1[1]; + d6 = acc1[2]; + d7 = acc1[3]; + + b = {b4, b5}; + auto acc2 = __builtin_mxc_mma_16x16x16bf16(b, a, {c8, c9, c10, c11}); + d8 = acc2[0]; + d9 = acc2[1]; + d10 = acc2[2]; + d11 = acc2[3]; + + b = {b6, b7}; + auto acc3 = __builtin_mxc_mma_16x16x16bf16(b, a, {c12, c13, c14, c15}); + d12 = acc3[0]; + d13 = acc3[1]; + d14 = acc3[2]; + d15 = acc3[3]; + + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x16 TN +struct SM80_16x8x16_F32BF16BF16F32_TN +{ + using DRegisters = float[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = float[4]; + + CUTE_HOST_DEVICE static void + fma(float & d0, float & d1, float & d2, float & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + float const & c0, float const & c1, float const & c2, float const & c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "f"(c0), "f"(c1), "f"(c2), "f"(c3)); + +#elif defined(__MACA_ARCH__) + + uint32_t aa[2]; + uint32_t bb[2]; + const int lane_id = __lane_id(); + int delta = (lane_id % 32 % 8 * 4 + lane_id % 32 / 16 * 2) - lane_id; + + // shuffle A + uint32_t tmp0 = __shfl_down_sync(ULONG_MAX, a0, delta); + uint32_t tmp0_1 = __shfl_down_sync(ULONG_MAX, a0, delta + 1); + uint32_t tmp1 = __shfl_down_sync(ULONG_MAX, a1, delta); + uint32_t tmp1_1 = __shfl_down_sync(ULONG_MAX, a1, delta + 1); + uint32_t tmp2 = __shfl_down_sync(ULONG_MAX, a2, delta); + uint32_t tmp2_1 = __shfl_down_sync(ULONG_MAX, a2, delta + 1); + uint32_t tmp3 = __shfl_down_sync(ULONG_MAX, a3, delta); + uint32_t tmp3_1 = __shfl_down_sync(ULONG_MAX, a3, delta + 1); + bool up_down_flag = lane_id % 16 / 8; + if (lane_id < 32) { + aa[0] = up_down_flag ? tmp1 : tmp0; + aa[1] = up_down_flag ? tmp1_1 : tmp0_1; + } else if (lane_id < 64) { + aa[0] = up_down_flag ? tmp3 : tmp2; + aa[1] = up_down_flag ? tmp3_1 : tmp2_1; + } + + // shuffle B + tmp0 = __shfl_down_sync(ULONG_MAX, b0, delta); + tmp0_1 = __shfl_down_sync(ULONG_MAX, b0, delta + 1); + tmp1 = __shfl_down_sync(ULONG_MAX, b1, delta); + tmp1_1 = __shfl_down_sync(ULONG_MAX, b1, delta + 1); + if (lane_id < 32) { + bb[0] = tmp0; + bb[1] = tmp0_1; + } else if (lane_id < 64) { + bb[0] = tmp1; + bb[1] = tmp1_1; + } + __maca_bfloat162 const *ha0 = reinterpret_cast<__maca_bfloat162 const *>(&a0); + __maca_bfloat162 const *ha1 = reinterpret_cast<__maca_bfloat162 const *>(&a1); + __maca_bfloat162 const *ha2 = reinterpret_cast<__maca_bfloat162 const *>(&a2); + __maca_bfloat162 const *ha3 = reinterpret_cast<__maca_bfloat162 const *>(&a3); + __maca_bfloat162 const *hb0 = reinterpret_cast<__maca_bfloat162 const *>(&b0); + __maca_bfloat162 const *hb1 = reinterpret_cast<__maca_bfloat162 const *>(&b1); + __maca_bfloat162 const *haa0 = reinterpret_cast<__maca_bfloat162 const *>(&aa[0]); + __maca_bfloat162 const *haa1 = reinterpret_cast<__maca_bfloat162 const *>(&aa[1]); + __maca_bfloat162 const *hbb0 = reinterpret_cast<__maca_bfloat162 const *>(&bb[0]); + __maca_bfloat162 const *hbb1 = reinterpret_cast<__maca_bfloat162 const *>(&bb[1]); + + // shuffle C + delta = (lane_id % 8 / 2 + lane_id % 32 / 16 * 16) - lane_id; + float tmp_c0 = __shfl_down_sync(ULONG_MAX, c0, delta); + float tmp_c0_4 = __shfl_down_sync(ULONG_MAX, c0, delta + 4); + float tmp_c0_8 = __shfl_down_sync(ULONG_MAX, c0, delta + 8); + float tmp_c0_12 = __shfl_down_sync(ULONG_MAX, c0, delta + 12); + float tmp_c1 = __shfl_down_sync(ULONG_MAX, c1, delta); + float tmp_c1_4 = __shfl_down_sync(ULONG_MAX, c1, delta + 4); + float tmp_c1_8 = __shfl_down_sync(ULONG_MAX, c1, delta + 8); + float tmp_c1_12 = __shfl_down_sync(ULONG_MAX, c1, delta + 12); + float tmp_c2 = __shfl_down_sync(ULONG_MAX, c2, delta); + float tmp_c2_4 = __shfl_down_sync(ULONG_MAX, c2, delta + 4); + float tmp_c2_8 = __shfl_down_sync(ULONG_MAX, c2, delta + 8); + float tmp_c2_12 = __shfl_down_sync(ULONG_MAX, c2, delta + 12); + float tmp_c3 = __shfl_down_sync(ULONG_MAX, c3, delta); + float tmp_c3_4 = __shfl_down_sync(ULONG_MAX, c3, delta + 4); + float tmp_c3_8 = __shfl_down_sync(ULONG_MAX, c3, delta + 8); + float tmp_c3_12 = __shfl_down_sync(ULONG_MAX, c3, delta + 12); + float cc0, cc1, cc2, cc3; + if (lane_id < 32) { + if (lane_id % 2) { + cc0 = tmp_c1; + cc1 = tmp_c1_4; + cc2 = tmp_c1_8; + cc3 = tmp_c1_12; + } else { + cc0 = tmp_c0; + cc1 = tmp_c0_4; + cc2 = tmp_c0_8; + cc3 = tmp_c0_12; + } + } else if (lane_id < 64) { + if (lane_id % 2) { + cc0 = tmp_c3; + cc1 = tmp_c3_4; + cc2 = tmp_c3_8; + cc3 = tmp_c3_12; + } else { + cc0 = tmp_c2; + cc1 = tmp_c2_4; + cc2 = tmp_c2_8; + cc3 = tmp_c2_12; + } + } + + auto result = __builtin_mxc_mma_16x16x16bf16({*reinterpret_cast(&(haa0->x)), *reinterpret_cast(&(haa0->y)), + *reinterpret_cast(&(haa1->x)), *reinterpret_cast(&(haa1->y))}, + {*reinterpret_cast(&(hbb0->x)), *reinterpret_cast(&(hbb0->y)), + *reinterpret_cast(&(hbb1->x)), *reinterpret_cast(&(hbb1->y))}, + {cc0, cc1, cc2, cc3}); + + delta = (lane_id % 4 * 2 + lane_id / 16 * 16) - lane_id; + float tmp_d0 = __shfl_down_sync(ULONG_MAX, result[0], delta); + float tmp_d1 = __shfl_down_sync(ULONG_MAX, result[1], delta); + float tmp_d2 = __shfl_down_sync(ULONG_MAX, result[2], delta); + float tmp_d3 = __shfl_down_sync(ULONG_MAX, result[3], delta); + float tmp_d0_1 = __shfl_down_sync(ULONG_MAX, result[0], delta + 1); + float tmp_d1_1 = __shfl_down_sync(ULONG_MAX, result[1], delta + 1); + float tmp_d2_1 = __shfl_down_sync(ULONG_MAX, result[2], delta + 1); + float tmp_d3_1 = __shfl_down_sync(ULONG_MAX, result[3], delta + 1); + float tmp_d0_32 = __shfl_down_sync(ULONG_MAX, result[0], delta + 32); + float tmp_d1_32 = __shfl_down_sync(ULONG_MAX, result[1], delta + 32); + float tmp_d2_32 = __shfl_down_sync(ULONG_MAX, result[2], delta + 32); + float tmp_d3_32 = __shfl_down_sync(ULONG_MAX, result[3], delta + 32); + float tmp_d0_32_1 = __shfl_down_sync(ULONG_MAX, result[0], delta + 32 + 1); + float tmp_d1_32_1 = __shfl_down_sync(ULONG_MAX, result[1], delta + 32 + 1); + float tmp_d2_32_1 = __shfl_down_sync(ULONG_MAX, result[2], delta + 32 + 1); + float tmp_d3_32_1 = __shfl_down_sync(ULONG_MAX, result[3], delta + 32 + 1); + + if (lane_id < 32) { + if (lane_id % 16 / 4 == 0) { + d0 = tmp_d0; + d1 = tmp_d0_1; + d2 = tmp_d0_32; + d3 = tmp_d0_32_1; + } else if (lane_id % 16 / 4 == 1) { + d0 = tmp_d1; + d1 = tmp_d1_1; + d2 = tmp_d1_32; + d3 = tmp_d1_32_1; + } else if (lane_id % 16 / 4 == 2) { + d0 = tmp_d2; + d1 = tmp_d2_1; + d2 = tmp_d2_32; + d3 = tmp_d2_32_1; + } else if (lane_id % 16 / 4 == 3) { + d0 = tmp_d3; + d1 = tmp_d3_1; + d2 = tmp_d3_32; + d3 = tmp_d3_32_1; + } + } + +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x16_F32BF16BF16F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x4 TN +struct SM80_16x8x4_F32TF32TF32F32_TN +{ + using DRegisters = float[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = float[4]; + + CUTE_HOST_DEVICE static void + fma(float & d0, float & d1, float & d2, float & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + float const & c0, float const & c1, float const & c2, float const & c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k4.row.col.f32.tf32.tf32.f32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "f"(c0), "f"(c1), "f"(c2), "f"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x4_F32TF32TF32F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x8 TN +struct SM80_16x8x8_F32TF32TF32F32_TN +{ + using DRegisters = float[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = float[4]; + + CUTE_HOST_DEVICE static void + fma(float & d0, float & d1, float & d2, float & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + float const & c0, float const & c1, float const & c2, float const & c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "f"(c0), "f"(c1), "f"(c2), "f"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x8_F32TF32TF32F32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x4 TN +struct SM80_8x8x4_F64F64F64F64_TN +{ + using DRegisters = double[2]; + using ARegisters = double[1]; + using BRegisters = double[1]; + using CRegisters = double[2]; + + CUTE_HOST_DEVICE static void + fma(double & d0, double & d1, + double const& a0, + double const& b0, + double const& c0, double const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k4.row.col.f64.f64.f64.f64 " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=d"(d0), "=d"(d1) + : "d"(a0), + "d"(b0), + "d"(c0), "d"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x4_F64F64F64F64_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +// MMA 8x8x4 TN with Planar Complex multiplication +struct SM80_8x8x4_C64C64C64C64_TN +{ + using DRegisters = complex[2]; + using ARegisters = complex[1]; + using BRegisters = complex[1]; + using CRegisters = complex[2]; + + CUTE_HOST_DEVICE static void + fma(complex & d0, complex & d1, + complex const& a0, + complex const& b0, + complex const& c0, complex const& c1) + { + // Because thrust::complex does not provide a mutable ref + double& rd0 = reinterpret_cast(d0)[0]; + double& id0 = reinterpret_cast(d0)[1]; + double& rd1 = reinterpret_cast(d1)[0]; + double& id1 = reinterpret_cast(d1)[1]; + + // d.real() = a.real() * b.real() + c.real(); + SM80_8x8x4_F64F64F64F64_TN::fma( + rd0, rd1, + a0.real(), + b0.real(), + c0.real(), c1.real()); + + // d.imag() = a.imag() * b.real() + c.imag(); + SM80_8x8x4_F64F64F64F64_TN::fma( + id0, id1, + a0.imag(), + b0.real(), + c0.imag(), c1.imag()); + + // d.real() = -a.imag() * b.imag() + d.real(); + SM80_8x8x4_F64F64F64F64_TN::fma( + rd0, rd1, + -a0.imag(), + b0.imag(), + d0.real(), d1.real()); + + // d.imag() = a.real() * b.imag() + d.imag(); + SM80_8x8x4_F64F64F64F64_TN::fma( + id0, id1, + a0.real(), + b0.imag(), + d0.imag(), d1.imag()); + } +}; + +// MMA 8x8x4 TN with Gaussian Complex multiplication: +// (a + bi)*(c + di) +// yields +// t0 += a*c +// t1 += b*d +// t2 += (a+b)*(c+d) +// then +// re = t0 - t1 +// im = t2 - t0 - t1 +struct SM80_8x8x4_GC64C64C64GC64_TN +{ + struct GaussComplex { + double t0, t1, t2; + + CUTE_HOST_DEVICE //constexpr + operator complex() const { return complex(t0 - t1, t2 - t0 - t1); } + + CUTE_HOST_DEVICE friend //constexpr + complex operator*(GaussComplex const& a, complex const& b) { return static_cast>(a) * b; } + CUTE_HOST_DEVICE friend //constexpr + complex operator*(complex const& a, GaussComplex const& b) { return b * a; } + + CUTE_HOST_DEVICE friend //constexpr + complex operator+(GaussComplex const& a, complex const& b) { return static_cast>(a) + b; } + CUTE_HOST_DEVICE friend //constexpr + complex operator+(complex const& a, GaussComplex const& b) { return b + a; } + }; + + using DRegisters = GaussComplex[2]; + using ARegisters = complex[1]; + using BRegisters = complex[1]; + using CRegisters = GaussComplex[2]; + + CUTE_HOST_DEVICE static void + fma(GaussComplex & d0, GaussComplex & d1, + complex const& a0, + complex const& b0, + GaussComplex const& c0, GaussComplex const& c1) + { + SM80_8x8x4_F64F64F64F64_TN::fma(d0.t0, d1.t0, + a0.real(), + b0.real(), + c0.t0, c1.t0); + SM80_8x8x4_F64F64F64F64_TN::fma(d0.t1, d1.t1, + a0.imag(), + b0.imag(), + c0.t1, c1.t1); + SM80_8x8x4_F64F64F64F64_TN::fma(d0.t2, d1.t2, + a0.real() + a0.imag(), + b0.real() + b0.imag(), + c0.t2, c1.t2); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x16 TN +struct SM80_8x8x16_S32S8S8S32_TN +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32S8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x16 TN +struct SM80_8x8x16_S32S8S8S32_TN_SATURATE +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32S8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x16 TN +struct SM80_16x8x16_S32S8S8S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x16_S32S8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x16 TN +struct SM80_16x8x16_S32S8S8S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x16_S32S8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32S8S8S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32S8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32S8S8S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32S8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x16 TN +struct SM80_8x8x16_S32S8U8S32_TN +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.s8.u8.s32 " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32S8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x16 TN +struct SM80_8x8x16_S32S8U8S32_TN_SATURATE +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.s8.u8.s32.satfinite " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32S8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x16 TN +struct SM80_16x8x16_S32S8U8S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.s8.u8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x16_S32S8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x16 TN +struct SM80_16x8x16_S32S8U8S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.s8.u8.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x16_S32S8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32S8U8S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.u8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32S8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32S8U8S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.u8.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32S8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x16 TN +struct SM80_8x8x16_S32U8S8S32_TN +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.u8.s8.s32 " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32U8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x16 TN +struct SM80_8x8x16_S32U8S8S32_TN_SATURATE +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.u8.s8.s32.satfinite " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32U8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x16 TN +struct SM80_16x8x16_S32U8S8S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.u8.s8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x16_S32U8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x16 TN +struct SM80_16x8x16_S32U8S8S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.u8.s8.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x16_S32U8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32U8S8S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.u8.s8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32U8S8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32U8S8S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.u8.s8.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32U8S8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x16 TN +struct SM80_8x8x16_S32U8U8S32_TN +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.u8.u8.s32 " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32U8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x16 TN +struct SM80_8x8x16_S32U8U8S32_TN_SATURATE +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k16.row.col.s32.u8.u8.s32.satfinite " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x16_S32U8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x16 TN +struct SM80_16x8x16_S32U8U8S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.u8.u8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x16_S32U8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x16 TN +struct SM80_16x8x16_S32U8U8S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.u8.u8.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x16_S32U8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32U8U8S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.u8.u8.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32U8U8S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32U8U8S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.u8.u8.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32U8U8S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x32 TN +struct SM80_8x8x32_S32S4S4S32_TN +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k32.row.col.s32.s4.s4.s32 " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32S4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x32 TN +struct SM80_8x8x32_S32S4S4S32_TN_SATURATE +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k32.row.col.s32.s4.s4.s32.satfinite " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32S4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32S4S4S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32S4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32S4S4S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s4.s4.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32S4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x64 TN +struct SM80_16x8x64_S32S4S4S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x64_S32S4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x64 TN +struct SM80_16x8x64_S32S4S4S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x64_S32S4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x32 TN +struct SM80_8x8x32_S32S4U4S32_TN +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k32.row.col.s32.s4.u4.s32 " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32S4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x32 TN +struct SM80_8x8x32_S32S4U4S32_TN_SATURATE +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k32.row.col.s32.s4.u4.s32.satfinite " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32S4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32S4U4S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s4.u4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32S4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32S4U4S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s4.u4.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32S4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x64 TN +struct SM80_16x8x64_S32S4U4S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.u4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x64_S32S4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x64 TN +struct SM80_16x8x64_S32S4U4S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.u4.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x64_S32S4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x32 TN +struct SM80_8x8x32_S32U4S4S32_TN +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k32.row.col.s32.u4.s4.s32 " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32U4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x32 TN +struct SM80_8x8x32_S32U4S4S32_TN_SATURATE +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k32.row.col.s32.u4.s4.s32.satfinite " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32U4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32U4S4S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.u4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32U4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32U4S4S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.u4.s4.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32U4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x64 TN +struct SM80_16x8x64_S32U4S4S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.u4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x64_S32U4S4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x64 TN +struct SM80_16x8x64_S32U4S4S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.u4.s4.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x64_S32U4S4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x32 TN +struct SM80_8x8x32_S32U4U4S32_TN +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k32.row.col.s32.u4.u4.s32 " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32U4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x32 TN +struct SM80_8x8x32_S32U4U4S32_TN_SATURATE +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k32.row.col.s32.u4.u4.s32.satfinite " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x32_S32U4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32U4U4S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.u4.u4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32U4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x32 TN +struct SM80_16x8x32_S32U4U4S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.u4.u4.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x32_S32U4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x64 TN +struct SM80_16x8x64_S32U4U4S32_TN +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.u4.u4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x64_S32U4U4S32_TN without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x64 TN +struct SM80_16x8x64_S32U4U4S32_TN_SATURATE +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.u4.u4.s32.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x64_S32U4U4S32_TN_SATURATE without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 8x8x128 TN +struct SM80_8x8x128_S32U1U1S32_TN_XORPOPC +{ + using DRegisters = uint32_t[2]; + using ARegisters = uint32_t[1]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, + uint32_t const& a0, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m8n8k128.row.col.s32.b1.b1.s32.xor.popc " + "{%0, %1}," + "{%2}," + "{%3}," + "{%4, %5};\n" + : "=r"(d0), "=r"(d1) + : "r"(a0), + "r"(b0), + "r"(c0), "r"(c1)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_8x8x128_S32U1U1S32_TN_XORPOPC without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x128 TN +struct SM80_16x8x128_S32U1U1S32_TN_XORPOPC +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[2]; + using BRegisters = uint32_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, + uint32_t const& b0, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k128.row.col.s32.b1.b1.s32.xor.popc " + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), + "r"(b0), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x128_S32U1U1S32_TN_XORPOPC without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x256 TN +struct SM80_16x8x256_S32U1U1S32_TN_XORPOPC +{ + using DRegisters = uint32_t[4]; + using ARegisters = uint32_t[4]; + using BRegisters = uint32_t[2]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint32_t const& b0, uint32_t const& b1, + uint32_t const& c0, uint32_t const& c1, uint32_t const& c2, uint32_t const& c3) + { +#if defined(CUTE_ARCH_MMA_SM80_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k256.row.col.s32.b1.b1.s32.xor.popc " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(d0), "=r"(d1), "=r"(d2), "=r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "r"(b0), "r"(b1), + "r"(c0), "r"(c1), "r"(c2), "r"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM80_16x8x256_S32U1U1S32_TN_XORPOPC without CUTE_ARCH_MMA_SM80_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cute diff --git a/csrc/mctlass/include/cute/arch/mma_sm90.hpp b/csrc/mctlass/include/cute/arch/mma_sm90.hpp new file mode 100644 index 0000000..93c9bd8 --- /dev/null +++ b/csrc/mctlass/include/cute/arch/mma_sm90.hpp @@ -0,0 +1,961 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include + +#include + +// Config +// #if defined(__MACA_ARCH__) && (__MACA_ARCH__ >= 900) +// # define CUTE_ARCH_MMA_SM90_ENABLED +// #endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cute { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x4 TN +struct SM90_16x8x4_F64F64F64F64_TN +{ + using DRegisters = double[4]; + using ARegisters = double[2]; + using BRegisters = double[1]; + using CRegisters = double[4]; + + CUTE_HOST_DEVICE static void + fma(double & d0, double & d1, double & d2, double & d3, + double const& a0, double const& a1, + double const& b0, + double const& c0, double const& c1, double const& c2, double const& c3) + { +#if defined(CUTE_ARCH_MMA_SM90_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k4.row.col.f64.f64.f64.f64" + "{%0, %1, %2, %3}," + "{%4, %5}," + "{%6}," + "{%7, %8, %9, %10};\n" + : "=d"(d0), "=d"(d1), "=d"(d2), "=d"(d3) + : "d"(a0), "d"(a1), + "d"(b0), + "d"(c0), "d"(c1), "d"(c2), "d"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_16x8x4_F64F64F64F64_TN without CUTE_ARCH_MMA_SM90_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x8 TN +struct SM90_16x8x8_F64F64F64F64_TN +{ + using DRegisters = double[4]; + using ARegisters = double[4]; + using BRegisters = double[2]; + using CRegisters = double[4]; + + CUTE_HOST_DEVICE static void + fma(double & d0, double & d1, double & d2, double & d3, + double const& a0, double const& a1, double const& a2, double const& a3, + double const& b0, double const& b1, + double const& c0, double const& c1, double const& c2, double const& c3) + { +#if defined(CUTE_ARCH_MMA_SM90_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f64.f64.f64.f64" + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=d"(d0), "=d"(d1), "=d"(d2), "=d"(d3) + : "d"(a0), "d"(a1), "d"(a2), "d"(a3), + "d"(b0), "d"(b1), + "d"(c0), "d"(c1), "d"(c2), "d"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_16x8x8_F64F64F64F64_TN without CUTE_ARCH_MMA_SM90_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x16 TN +struct SM90_16x8x16_F64F64F64F64_TN +{ + using DRegisters = double[4]; + using ARegisters = double[8]; + using BRegisters = double[4]; + using CRegisters = double[4]; + + CUTE_HOST_DEVICE static void + fma(double & d0, double & d1, double & d2, double & d3, + double const& a0, double const& a1, double const& a2, double const& a3, + double const& a4, double const& a5, double const& a6, double const& a7, + double const& b0, double const& b1, double const& b2, double const& b3, + double const& c0, double const& c1, double const& c2, double const& c3) + { +#if defined(CUTE_ARCH_MMA_SM90_ENABLED) + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f64.f64.f64.f64" + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7, %8, %9, %10, %11}," + "{%12, %13, %14, %15}," + "{%16, %17, %18, %19};\n" + : "=d"(d0), "=d"(d1), "=d"(d2), "=d"(d3) + : "d"(a0), "d"(a1), "d"(a2), "d"(a3), + "d"(a4), "d"(a5), "d"(a6), "d"(a7), + "d"(b0), "d"(b1), "d"(b2), "d"(b3), + "d"(c0), "d"(c1), "d"(c2), "d"(c3)); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_16x8x16_F64F64F64F64_TN without CUTE_ARCH_MMA_SM90_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x4 TN +struct SM90_16x8x4_C64C64C64C64_TN +{ + using DRegisters = complex[4]; + using ARegisters = complex[2]; + using BRegisters = complex[1]; + using CRegisters = complex[4]; + + CUTE_HOST_DEVICE static void + fma(complex & d0, complex & d1, + complex & d2, complex & d3, + complex const& a0, complex const& a1, + complex const& b0, + complex const& c0, complex const& c1, + complex const& c2, complex const& c3) + { + // Because thrust::complex does not provide a mutable ref + double& rd0 = reinterpret_cast(d0)[0]; + double& id0 = reinterpret_cast(d0)[1]; + double& rd1 = reinterpret_cast(d1)[0]; + double& id1 = reinterpret_cast(d1)[1]; + double& rd2 = reinterpret_cast(d2)[0]; + double& id2 = reinterpret_cast(d2)[1]; + double& rd3 = reinterpret_cast(d3)[0]; + double& id3 = reinterpret_cast(d3)[1]; + + // d.real() = a.real() * b.real() + c.real(); + SM90_16x8x4_F64F64F64F64_TN::fma( + rd0, rd1, rd2, rd3, + a0.real(), a1.real(), + b0.real(), + c0.real(), c1.real(), c2.real(), c3.real()); + + // d.imag() = a.imag() * b.real() + c.imag(); + SM90_16x8x4_F64F64F64F64_TN::fma( + id0, id1, id2, id3, + a0.imag(), a1.imag(), + b0.real(), + c0.imag(), c1.imag(), c2.imag(), c3.imag()); + + // d.real() = -a.imag() * b.imag() + d.real(); + SM90_16x8x4_F64F64F64F64_TN::fma( + rd0, rd1, rd2, rd3, + -a0.imag(), -a1.imag(), + b0.imag(), + d0.real(), d1.real(), d2.real(), d3.real()); + + // d.imag() = a.real() * b.imag() + d.imag(); + SM90_16x8x4_F64F64F64F64_TN::fma( + id0, id1, id2, id3, + a0.real(), a1.real(), + b0.imag(), + d0.imag(), d1.imag(), d2.imag(), d3.imag()); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x8 TN +struct SM90_16x8x8_C64C64C64C64_TN +{ + using DRegisters = complex[4]; + using ARegisters = complex[4]; + using BRegisters = complex[2]; + using CRegisters = complex[4]; + + CUTE_HOST_DEVICE static void + fma(complex & d0, complex & d1, + complex & d2, complex & d3, + complex const& a0, complex const& a1, + complex const& a2, complex const& a3, + complex const& b0, complex const& b1, + complex const& c0, complex const& c1, + complex const& c2, complex const& c3) + { + // Because thrust::complex does not provide a mutable ref + double& rd0 = reinterpret_cast(d0)[0]; + double& id0 = reinterpret_cast(d0)[1]; + double& rd1 = reinterpret_cast(d1)[0]; + double& id1 = reinterpret_cast(d1)[1]; + double& rd2 = reinterpret_cast(d2)[0]; + double& id2 = reinterpret_cast(d2)[1]; + double& rd3 = reinterpret_cast(d3)[0]; + double& id3 = reinterpret_cast(d3)[1]; + + // d.real() = a.real() * b.real() + c.real(); + SM90_16x8x8_F64F64F64F64_TN::fma( + rd0, rd1, rd2, rd3, + a0.real(), a1.real(), a2.real(), a3.real(), + b0.real(), b1.real(), + c0.real(), c1.real(), c2.real(), c3.real()); + + // d.imag() = a.imag() * b.real() + c.imag(); + SM90_16x8x8_F64F64F64F64_TN::fma( + id0, id1, id2, id3, + a0.imag(), a1.imag(), a2.imag(), a3.imag(), + b0.real(), b1.real(), + c0.imag(), c1.imag(), c2.imag(), c3.imag()); + + // d.real() = -a.imag() * b.imag() + d.real(); + SM90_16x8x8_F64F64F64F64_TN::fma( + rd0, rd1, rd2, rd3, + -a0.imag(), -a1.imag(), -a2.imag(), -a3.imag(), + b0.imag(), b1.imag(), + d0.real(), d1.real(), d2.real(), d3.real()); + + // d.imag() = a.real() * b.imag() + d.imag(); + SM90_16x8x8_F64F64F64F64_TN::fma( + id0, id1, id2, id3, + a0.real(), a1.real(), a2.real(), a3.real(), + b0.imag(), b1.imag(), + d0.imag(), d1.imag(), d2.imag(), d3.imag()); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MMA 16x8x16 TN +struct SM90_16x8x16_C64C64C64C64_TN +{ + using DRegisters = complex[4]; + using ARegisters = complex[8]; + using BRegisters = complex[4]; + using CRegisters = complex[4]; + + CUTE_HOST_DEVICE static void + fma(complex & d0, complex & d1, + complex & d2, complex & d3, + complex const& a0, complex const& a1, + complex const& a2, complex const& a3, + complex const& a4, complex const& a5, + complex const& a6, complex const& a7, + complex const& b0, complex const& b1, + complex const& b2, complex const& b3, + complex const& c0, complex const& c1, + complex const& c2, complex const& c3) + { + // Because thrust::complex does not provide a mutable ref + double& rd0 = reinterpret_cast(d0)[0]; + double& id0 = reinterpret_cast(d0)[1]; + double& rd1 = reinterpret_cast(d1)[0]; + double& id1 = reinterpret_cast(d1)[1]; + double& rd2 = reinterpret_cast(d2)[0]; + double& id2 = reinterpret_cast(d2)[1]; + double& rd3 = reinterpret_cast(d3)[0]; + double& id3 = reinterpret_cast(d3)[1]; + + // d.real() = a.real() * b.real() + c.real(); + SM90_16x8x16_F64F64F64F64_TN::fma( + rd0, rd1, rd2, rd3, + a0.real(), a1.real(), a2.real(), a3.real(), + a4.real(), a5.real(), a6.real(), a7.real(), + b0.real(), b1.real(), b2.real(), b3.real(), + c0.real(), c1.real(), c2.real(), c3.real()); + + // d.imag() = a.imag() * b.real() + c.imag(); + SM90_16x8x16_F64F64F64F64_TN::fma( + id0, id1, id2, id3, + a0.imag(), a1.imag(), a2.imag(), a3.imag(), + a4.imag(), a5.imag(), a6.imag(), a7.imag(), + b0.real(), b1.real(), b2.real(), b3.real(), + c0.imag(), c1.imag(), c2.imag(), c3.imag()); + + // d.real() = -a.imag() * b.imag() + d.real(); + SM90_16x8x16_F64F64F64F64_TN::fma( + rd0, rd1, rd2, rd3, + -a0.imag(), -a1.imag(), -a2.imag(), -a3.imag(), + -a4.imag(), -a5.imag(), -a6.imag(), -a7.imag(), + b0.imag(), b1.imag(), b2.imag(), b3.imag(), + d0.real(), d1.real(), d2.real(), d3.real()); + + // d.imag() = a.real() * b.imag() + d.imag(); + SM90_16x8x16_F64F64F64F64_TN::fma( + id0, id1, id2, id3, + a0.real(), a1.real(), a2.real(), a3.real(), + a4.real(), a5.real(), a6.real(), a7.real(), + b0.imag(), b1.imag(), b2.imag(), b3.imag(), + d0.imag(), d1.imag(), d2.imag(), d3.imag()); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cute + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#include +#include + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cute { +namespace GMMA { + +template < + class ElementA, + class ElementB, + class ElementC, + class TileShape_MNK, + GMMA::Major MajorA = GMMA::Major::K, + GMMA::Major MajorB = GMMA::Major::K, + auto... Args // e.g. GMMA::ScaleOut::One, [GMMA::ScaleIn::One, GMMA::ScaleIn::One] + // But most commonly leave empty for defaults +> +CUTE_HOST_DEVICE constexpr +auto +ss_op_selector() +{ + static_assert(is_static::value, "TileShape_MNK must be static."); + static_assert(rank(TileShape_MNK{}) == 3, "TileShape_MNK must be rank 3."); + static_assert(size<0>(TileShape_MNK{}) % 64 == 0, "Tile_M must be a multiple of 64."); + auto Tile_N = size<1>(TileShape_MNK{}); + + // FP16 accumulator + if constexpr (is_same_v) { + static_assert(is_same_v, "Element types for AB must be half if ElementC is half."); + static_assert(is_same_v, "Element types for AB must be half if ElementC is half."); + static_assert(size<2>(TileShape_MNK{}) % 16 == 0, "Tile_K must be a multiple of 16."); + + // Dispatch against the Tile N mode size + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x16_F16F16F16_SS{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x16_F16F16F16_SS{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x16_F16F16F16_SS{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x16_F16F16F16_SS{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x16_F16F16F16_SS{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x16_F16F16F16_SS{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x16_F16F16F16_SS{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x16_F16F16F16_SS{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // FP32 accumulator + else if constexpr (is_same_v) { + + // FP16 inputs + if constexpr (is_same_v) { + static_assert(is_same_v, "ElementA and ElementB must be the same type for this config."); + static_assert(size<2>(TileShape_MNK{}) % 16 == 0, "Tile_K must be a multiple of 16."); + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x16_F32F16F16_SS{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x16_F32F16F16_SS{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x16_F32F16F16_SS{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x16_F32F16F16_SS{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x16_F32F16F16_SS{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x16_F32F16F16_SS{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x16_F32F16F16_SS{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x16_F32F16F16_SS{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // BF16 inputs + else if constexpr (is_same_v) { + static_assert(is_same_v, "ElementA and ElementB must be the same type for this config."); + static_assert(size<2>(TileShape_MNK{}) % 16 == 0, "Tile_K must be a multiple of 16."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x16_F32BF16BF16_SS{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x16_F32BF16BF16_SS{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x16_F32BF16BF16_SS{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x16_F32BF16BF16_SS{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x16_F32BF16BF16_SS{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x16_F32BF16BF16_SS{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x16_F32BF16BF16_SS{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x16_F32BF16BF16_SS{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // TF32 inputs + else if constexpr (is_same_v) { + static_assert(is_same_v, "ElementA and ElementB must be the same type for this config."); + static_assert(MajorA == GMMA::Major::K, "MajorA must be GMMA::Major::K for this config."); + static_assert(MajorB == GMMA::Major::K, "MajorB must be GMMA::Major::K for this config."); + static_assert(size<2>(TileShape_MNK{}) % 8 == 0, "Tile_K must be a multiple of 8."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x8_F32TF32TF32_SS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x8_F32TF32TF32_SS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x8_F32TF32TF32_SS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x8_F32TF32TF32_SS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x8_F32TF32TF32_SS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x8_F32TF32TF32_SS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x8_F32TF32TF32_SS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x8_F32TF32TF32_SS_TN{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + else { + static_assert(sizeof(ElementA) == 0, "No eligible GMMA operator for request configuration."); + } + } + + // S32 accumulator + else if constexpr (is_same_v) { + static_assert(MajorA == GMMA::Major::K, "MajorA must be GMMA::Major::K for this config."); + static_assert(MajorB == GMMA::Major::K, "MajorB must be GMMA::Major::K for this config."); + static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32."); + + // ElementA == int8_t && ElementB == int8_t + if constexpr (is_same_v && is_same_v) { + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x32_S32S8S8_SS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x32_S32S8S8_SS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x32_S32S8S8_SS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x32_S32S8S8_SS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x32_S32S8S8_SS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x32_S32S8S8_SS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x32_S32S8S8_SS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x32_S32S8S8_SS_TN{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // ElementA == int8_t && ElementB == uint8_t + else if constexpr (is_same_v && is_same_v) { + static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x32_S32S8U8_SS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x32_S32S8U8_SS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x32_S32S8U8_SS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x32_S32S8U8_SS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x32_S32S8U8_SS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x32_S32S8U8_SS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x32_S32S8U8_SS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x32_S32S8U8_SS_TN{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // ElementA == uint8_t && ElementB == int8_t + else if constexpr (is_same_v && is_same_v) { + static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x32_S32U8S8_SS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x32_S32U8S8_SS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x32_S32U8S8_SS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x32_S32U8S8_SS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x32_S32U8S8_SS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x32_S32U8S8_SS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x32_S32U8S8_SS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x32_S32U8S8_SS_TN{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // ElementA == uint8_t && ElementB == uint8_t + else if constexpr (is_same_v && is_same_v) { + static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x32_S32U8U8_SS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x32_S32U8U8_SS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x32_S32U8U8_SS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x32_S32U8U8_SS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x32_S32U8U8_SS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x32_S32U8U8_SS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x32_S32U8U8_SS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x32_S32U8U8_SS_TN{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + } + + // Unknown accumulator type + else { + static_assert(sizeof(ElementC) == 0, "Unknown ElementC accumulator type."); + } +} + +template < + class ElementA, + class ElementB, + class ElementC, + class TileShape_MNK, + GMMA::Major MajorA = GMMA::Major::K, + GMMA::Major MajorB = GMMA::Major::K, + auto... Args // e.g. GMMA::ScaleOut::One, [GMMA::ScaleIn::One, GMMA::ScaleIn::One] + // But most commonly leave empty for defaults +> +CUTE_HOST_DEVICE constexpr +auto +rs_op_selector() +{ + static_assert(is_static::value, "TileShape_MNK must be static."); + static_assert(rank(TileShape_MNK{}) == 3, "TileShape_MNK must be rank 3."); + static_assert(size<0>(TileShape_MNK{}) % 64 == 0, "Tile_M must be a multiple of 64."); + static_assert(MajorA == GMMA::Major::K, "Register source A operand GMMAs must have K-major A layout."); + auto Tile_N = size<1>(TileShape_MNK{}); + + // FP16 accumulator + if constexpr (is_same_v) { + static_assert(is_same_v, "Element types for AB must be half if ElementC is half."); + static_assert(is_same_v, "Element types for AB must be half if ElementC is half."); + static_assert(size<2>(TileShape_MNK{}) % 16 == 0, "Tile_K must be a multiple of 16."); + + // Dispatch against the Tile N mode size + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x16_F16F16F16_RS{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x16_F16F16F16_RS{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x16_F16F16F16_RS{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x16_F16F16F16_RS{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x16_F16F16F16_RS{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x16_F16F16F16_RS{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x16_F16F16F16_RS{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x16_F16F16F16_RS{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // FP32 accumulator + else if constexpr (is_same_v) { + static_assert(is_same_v, "ElementA and ElementB must be the same type for this config."); + static_assert(size<2>(TileShape_MNK{}) % 16 == 0, "Tile_K must be a multiple of 16."); + + // FP16 inputs + if constexpr (is_same_v) { + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x16_F32F16F16_RS{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x16_F32F16F16_RS{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x16_F32F16F16_RS{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x16_F32F16F16_RS{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x16_F32F16F16_RS{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x16_F32F16F16_RS{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x16_F32F16F16_RS{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x16_F32F16F16_RS{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // BF16 inputs + else if constexpr (is_same_v) { + static_assert(size<2>(TileShape_MNK{}) % 16 == 0, "Tile_K must be a multiple of 16."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x16_F32BF16BF16_RS{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x16_F32BF16BF16_RS{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x16_F32BF16BF16_RS{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x16_F32BF16BF16_RS{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x16_F32BF16BF16_RS{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x16_F32BF16BF16_RS{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x16_F32BF16BF16_RS{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x16_F32BF16BF16_RS{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // TF32 inputs + else if constexpr (is_same_v) { + static_assert(MajorB == GMMA::Major::K, "MajorB must be GMMA::Major::K for this config."); + static_assert(size<2>(TileShape_MNK{}) % 8 == 0, "Tile_K must be a multiple of 8."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x8_F32TF32TF32_RS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x8_F32TF32TF32_RS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x8_F32TF32TF32_RS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x8_F32TF32TF32_RS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x8_F32TF32TF32_RS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x8_F32TF32TF32_RS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x8_F32TF32TF32_RS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x8_F32TF32TF32_RS_TN{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + else { + static_assert(sizeof(ElementA) == 0, "No eligible GMMA operator for request configuration."); + } + } + + // S32 accumulator + else if constexpr (is_same_v) { + static_assert(MajorB == GMMA::Major::K, "MajorB must be GMMA::Major::K for this config."); + static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32."); + + // ElementA == int8_t && ElementB == int8_t + if constexpr (is_same_v && is_same_v) { + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x32_S32S8S8_RS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x32_S32S8S8_RS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x32_S32S8S8_RS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x32_S32S8S8_RS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x32_S32S8S8_RS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x32_S32S8S8_RS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x32_S32S8S8_RS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x32_S32S8S8_RS_TN{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // ElementA == int8_t && ElementB == uint8_t + else if constexpr (is_same_v && is_same_v) { + static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x32_S32S8U8_RS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x32_S32S8U8_RS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x32_S32S8U8_RS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x32_S32S8U8_RS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x32_S32S8U8_RS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x32_S32S8U8_RS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x32_S32S8U8_RS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x32_S32S8U8_RS_TN{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // ElementA == uint8_t && ElementB == int8_t + else if constexpr (is_same_v && is_same_v) { + static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x32_S32U8S8_RS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x32_S32U8S8_RS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x32_S32U8S8_RS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x32_S32U8S8_RS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x32_S32U8S8_RS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x32_S32U8S8_RS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x32_S32U8S8_RS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x32_S32U8S8_RS_TN{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + + // ElementA == uint8_t && ElementB == uint8_t + else if constexpr (is_same_v && is_same_v) { + static_assert(size<2>(TileShape_MNK{}) % 32 == 0, "Tile_K must be a multiple of 32."); + + if constexpr (Tile_N % 256 == 0) { + return SM90_64x256x32_S32U8U8_RS_TN{}; + } + else if constexpr (Tile_N % 192 == 0) { + return SM90_64x192x32_S32U8U8_RS_TN{}; + } + else if constexpr (Tile_N % 128 == 0) { + return SM90_64x128x32_S32U8U8_RS_TN{}; + } + else if constexpr (Tile_N % 96 == 0) { + return SM90_64x96x32_S32U8U8_RS_TN{}; + } + else if constexpr (Tile_N % 64 == 0) { + return SM90_64x64x32_S32U8U8_RS_TN{}; + } + else if constexpr (Tile_N % 32 == 0) { + return SM90_64x32x32_S32U8U8_RS_TN{}; + } + else if constexpr (Tile_N % 16 == 0) { + return SM90_64x16x32_S32U8U8_RS_TN{}; + } + else if constexpr (Tile_N % 8 == 0) { + return SM90_64x8x32_S32U8U8_RS_TN{}; + } + else { + static_assert(Tile_N % 8 == 0, "Tile_N must be a multiple of 8."); + } + } + } + + // Unknown accumulator type + else { + static_assert(sizeof(ElementC) == 0, "Unknown ElementC accumulator type."); + } +} +} // end namespace GMMA +} // end namespace cute + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/cute/arch/mma_sm90_desc.hpp b/csrc/mctlass/include/cute/arch/mma_sm90_desc.hpp new file mode 100644 index 0000000..b5c4512 --- /dev/null +++ b/csrc/mctlass/include/cute/arch/mma_sm90_desc.hpp @@ -0,0 +1,135 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#if !defined(__MACACC_RTC__) +#include +#endif + +#include + +#include + + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cute { + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// GMMA Descriptor and utilities + +// GMMA enums and utilities +namespace GMMA +{ + +enum class LayoutType : uint8_t { + INTERLEAVE = 0, + B128 = 1, + B64 = 2, + B32 = 3, +}; + +CUTE_HOST_DEVICE char const* to_string(LayoutType const& t) { + switch (t) { + case LayoutType::INTERLEAVE: return "INTERLEAVE"; + case LayoutType::B128: return "B128"; + case LayoutType::B64: return "B64"; + case LayoutType::B32: return "B32"; + } + return nullptr; +} + +#if !defined(__MACACC_RTC__) +// Output operator for all enums in this namespace +CUTE_HOST std::ostream& operator<<(std::ostream& os, LayoutType const& t) { + char const* s = to_string(t); + if (s) { + std::operator<<(os, s); // Explicit call to avoid ambiguity + } else { + os.setstate(std::ios_base::failbit); + } + return os; +} +#endif // !defined(__MACACC_RTC__) + +} // end namespace GMMA + +union GmmaDescriptor +{ + uint64_t desc_; + uint32_t reg32_[2]; + uint16_t reg16_[4]; + + // Bitfield implementation avoids the need for shifts in assignment + struct { + // start_address, bit [0,14), 4LSB not included + uint16_t start_address_ : 14, : 2; // 14 bits [0,14), 2 bits unused + // leading dimension byte offset, bit [16,30), 4LSB not included + // For N: This is the stride from the first col to the second col of the 8x2 brick in INTERLEAVED + // Unused for all SWIZZLE_* layouts (and assumed to be 1) + // For T: This is the stride from the first 8 rows to the next 8 rows. + uint16_t leading_byte_offset_ : 14, : 2; // 14 bits [0,14), 2 bits unused + // stride dimension byte offset, bit [32,46), 4LSB not included + // For N: This is the stride from the first 8 rows to the next 8 rows. + // For T: This is the stride fro mthe first 8 cols to the next 8 cols. + uint16_t stride_byte_offset_ : 14, : 2; // 14 bits [0,14), 2 bits unused + // base_offset, bit [49,52) + // Valid only for SWIZZLE_128B and SWIZZLE_64B + uint8_t : 1, base_offset_ : 3, : 4; // 1 bit unused, 3 bits [1,4), 4 bits unused + // layout type, bit [62,64) + // SWIZZLE_NONE = 0, SWIZZLE_32B = 3, SWIZZLE_64B = 2, SWIZZLE_128B = 1 + uint8_t : 6, layout_type_ : 2; // 6 bits unused, 2 bits [6,8) + }; + + // Decay to a uint64_t + CUTE_HOST_DEVICE constexpr + operator uint64_t() const noexcept { return desc_; } + + // Printer + CUTE_HOST_DEVICE friend void print(GmmaDescriptor const& t) + { + #if !defined(__MACACC_RTC__) + printf("GmmaDescriptor: 0x%016" PRIx64 "\n", t.desc_); + printf(" start_addr : 0x%04x\n", t.start_address_); + printf(" leading_off: 0x%04x (%d)\n", t.leading_byte_offset_, t.leading_byte_offset_); + printf(" stride_off : 0x%04x (%d)\n", t.stride_byte_offset_, t.stride_byte_offset_); + printf(" base_offset: 0x%01x\n", t.base_offset_); + printf(" layout_type: 0x%01x (%s)\n", t.layout_type_, to_string(static_cast(t.layout_type_))); + #endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cute + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/cute/arch/mma_sm90_gmma.hpp b/csrc/mctlass/include/cute/arch/mma_sm90_gmma.hpp new file mode 100644 index 0000000..4f00b9f --- /dev/null +++ b/csrc/mctlass/include/cute/arch/mma_sm90_gmma.hpp @@ -0,0 +1,12780 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +#include + + +namespace cute { + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Warpgroup sync primitives + +CUTE_HOST_DEVICE +void +warpgroup_arrive() +{ +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile ("wgmma.fence.sync.aligned;\n" ::: "memory"); +#else + CUTE_RUNTIME_ASSERT("Attempting to use wgmma.fence without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif +} + +template +CUTE_HOST_DEVICE +void +warpgroup_wait() +{ + static_assert(N >= 0 && N <= 7, "_warpgroup.wait {N}; must be in range [0, 7]"); +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile("wgmma.wait_group.sync.aligned %0;\n" :: "n"(N) : "memory"); +#else + CUTE_RUNTIME_ASSERT("Attempting to use wgmma.wait_group without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif +} + +// Marks the commit point for one or more sized batch of warpgroup MMAs. +CUTE_HOST_DEVICE +void +warpgroup_commit_batch() +{ +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory"); +#else + CUTE_RUNTIME_ASSERT("Attempting to use wgmma.commit_group without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif +} + +CUTE_HOST_DEVICE +void +warpgroup_fence_operand(uint32_t& reg) { + // MSVC emits a build error for 'asm volatile' + // even if it only occurs in a __device__ function. + // This prevents the error. +#if defined(__MACA_ARCH__) + asm volatile("" : "+r"(reg) :: "memory"); +#endif +} + +CUTE_HOST_DEVICE +void +warpgroup_fence_operand(float& reg) { +#if defined(__MACA_ARCH__) + asm volatile("" : "+f"(reg) :: "memory"); +#endif +} + +namespace GMMA { + +enum class Major { + K = 0, + MN = 1 +}; + +enum class ScaleOut { + Zero = 0, + One = 1 +}; + +enum class ScaleIn { + Neg = -1, + One = 1 +}; + +} // namespace GMMA + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// GMMA PTX definitions: C = (scaleA * A) * (scaleB * B) + (scaleD * C) +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x8x16_F16F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[2]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %4, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k16.f16.f16.f16 " + "{%0, %1}," + " %2," + " %3," + " p, %5, %6, %7, %8;\n" + "}\n" + : "+r"(d0), "+r"(d1) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x8x16_F16F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[2]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %7, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + " %6," + " p, %8, %9, %10;\n" + "}\n" + : "+r"(d0), "+r"(d1) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x16x16_F16F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k16.f16.f16.f16 " + "{%0, %1, %2, %3}," + " %4," + " %5," + " p, %7, %8, %9, %10;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x16x16_F16F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %9, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k16.f16.f16.f16 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + " %8," + " p, %10, %11, %12;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x32x16_F16F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f16.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + " %8," + " %9," + " p, %11, %12, %13, %14;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x32x16_F16F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %13, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f16.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + "{%8, %9, %10, %11}," + " %12," + " p, %14, %15, %16;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x64x16_F16F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %18, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f16.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + " %16," + " %17," + " p, %19, %20, %21, %22;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x64x16_F16F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %21, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f16.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + "{%16, %17, %18, %19}," + " %20," + " p, %22, %23, %24;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x96x16_F16F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[24]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %26, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k16.f16.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23}," + " %24," + " %25," + " p, %27, %28, %29, %30;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x96x16_F16F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[24]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %29, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k16.f16.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23}," + "{%24, %25, %26, %27}," + " %28," + " p, %30, %31, %32;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x128x16_F16F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %34, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f16.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " p, %35, %36, %37, %38;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x128x16_F16F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %37, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f16.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + "{%32, %33, %34, %35}," + " %36," + " p, %38, %39, %40;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x192x16_F16F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %50, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f16.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + " %48," + " %49," + " p, %51, %52, %53, %54;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x192x16_F16F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %53, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f16.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + "{%48, %49, %50, %51}," + " %52," + " p, %54, %55, %56;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x256x16_F16F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %66, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f16.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " p, %67, %68, %69, %70;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x16_F16F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x16 F16+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x256x16_F16F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %69, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f16.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + "{%64, %65, %66, %67}," + " %68," + " p, %70, %71, %72;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x16_F16F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x8x16_F32F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[4]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d0, float & d1, float & d2, float & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " + "{%0, %1, %2, %3}," + " %4," + " %5," + " p, %7, %8, %9, %10;\n" + "}\n" + : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x8x16_F32F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[4]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + float & d0, float & d1, float & d2, float & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %9, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k16.f32.f16.f16 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + " %8," + " p, %10, %11, %12;\n" + "}\n" + : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x16x16_F32F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[8]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d0, float & d1, float & d2, float & d3, + float & d4, float & d5, float & d6, float & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + " %8," + " %9," + " p, %11, %12, %13, %14;\n" + "}\n" + : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3), + "+f"(d4), "+f"(d5), "+f"(d6), "+f"(d7) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x16x16_F32F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[8]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + float & d0, float & d1, float & d2, float & d3, + float & d4, float & d5, float & d6, float & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %13, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + "{%8, %9, %10, %11}," + " %12," + " p, %14, %15, %16;\n" + "}\n" + : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3), + "+f"(d4), "+f"(d5), "+f"(d6), "+f"(d7) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x32x16_F32F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[16]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %18, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + " %16," + " %17," + " p, %19, %20, %21, %22;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x32x16_F32F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[16]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %21, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + "{%16, %17, %18, %19}," + " %20," + " p, %22, %23, %24;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x64x16_F32F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[32]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %34, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " p, %35, %36, %37, %38;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x64x16_F32F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[32]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %37, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + "{%32, %33, %34, %35}," + " %36," + " p, %38, %39, %40;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x96x16_F32F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[48]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %50, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + " %48," + " %49," + " p, %51, %52, %53, %54;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x96x16_F32F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[48]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %53, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + "{%48, %49, %50, %51}," + " %52," + " p, %54, %55, %56;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x128x16_F32F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[64]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + float & d48, float & d49, float & d50, float & d51, + float & d52, float & d53, float & d54, float & d55, + float & d56, float & d57, float & d58, float & d59, + float & d60, float & d61, float & d62, float & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %66, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " p, %67, %68, %69, %70;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47), + "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), + "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55), + "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), + "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x128x16_F32F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[64]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + float & d48, float & d49, float & d50, float & d51, + float & d52, float & d53, float & d54, float & d55, + float & d56, float & d57, float & d58, float & d59, + float & d60, float & d61, float & d62, float & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %69, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + "{%64, %65, %66, %67}," + " %68," + " p, %70, %71, %72;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47), + "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), + "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55), + "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), + "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x192x16_F32F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[96]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + float & d48, float & d49, float & d50, float & d51, + float & d52, float & d53, float & d54, float & d55, + float & d56, float & d57, float & d58, float & d59, + float & d60, float & d61, float & d62, float & d63, + float & d64, float & d65, float & d66, float & d67, + float & d68, float & d69, float & d70, float & d71, + float & d72, float & d73, float & d74, float & d75, + float & d76, float & d77, float & d78, float & d79, + float & d80, float & d81, float & d82, float & d83, + float & d84, float & d85, float & d86, float & d87, + float & d88, float & d89, float & d90, float & d91, + float & d92, float & d93, float & d94, float & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %98, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + " %96," + " %97," + " p, %99, %100, %101, %102;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47), + "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), + "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55), + "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), + "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63), + "+f"(d64), "+f"(d65), "+f"(d66), "+f"(d67), + "+f"(d68), "+f"(d69), "+f"(d70), "+f"(d71), + "+f"(d72), "+f"(d73), "+f"(d74), "+f"(d75), + "+f"(d76), "+f"(d77), "+f"(d78), "+f"(d79), + "+f"(d80), "+f"(d81), "+f"(d82), "+f"(d83), + "+f"(d84), "+f"(d85), "+f"(d86), "+f"(d87), + "+f"(d88), "+f"(d89), "+f"(d90), "+f"(d91), + "+f"(d92), "+f"(d93), "+f"(d94), "+f"(d95) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x192x16_F32F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[96]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + float & d48, float & d49, float & d50, float & d51, + float & d52, float & d53, float & d54, float & d55, + float & d56, float & d57, float & d58, float & d59, + float & d60, float & d61, float & d62, float & d63, + float & d64, float & d65, float & d66, float & d67, + float & d68, float & d69, float & d70, float & d71, + float & d72, float & d73, float & d74, float & d75, + float & d76, float & d77, float & d78, float & d79, + float & d80, float & d81, float & d82, float & d83, + float & d84, float & d85, float & d86, float & d87, + float & d88, float & d89, float & d90, float & d91, + float & d92, float & d93, float & d94, float & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %101, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + "{%96, %97, %98, %99}," + " %100," + " p, %102, %103, %104;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47), + "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), + "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55), + "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), + "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63), + "+f"(d64), "+f"(d65), "+f"(d66), "+f"(d67), + "+f"(d68), "+f"(d69), "+f"(d70), "+f"(d71), + "+f"(d72), "+f"(d73), "+f"(d74), "+f"(d75), + "+f"(d76), "+f"(d77), "+f"(d78), "+f"(d79), + "+f"(d80), "+f"(d81), "+f"(d82), "+f"(d83), + "+f"(d84), "+f"(d85), "+f"(d86), "+f"(d87), + "+f"(d88), "+f"(d89), "+f"(d90), "+f"(d91), + "+f"(d92), "+f"(d93), "+f"(d94), "+f"(d95) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x256x16_F32F16F16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[128]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d000, float & d001, float & d002, float & d003, + float & d004, float & d005, float & d006, float & d007, + float & d008, float & d009, float & d010, float & d011, + float & d012, float & d013, float & d014, float & d015, + float & d016, float & d017, float & d018, float & d019, + float & d020, float & d021, float & d022, float & d023, + float & d024, float & d025, float & d026, float & d027, + float & d028, float & d029, float & d030, float & d031, + float & d032, float & d033, float & d034, float & d035, + float & d036, float & d037, float & d038, float & d039, + float & d040, float & d041, float & d042, float & d043, + float & d044, float & d045, float & d046, float & d047, + float & d048, float & d049, float & d050, float & d051, + float & d052, float & d053, float & d054, float & d055, + float & d056, float & d057, float & d058, float & d059, + float & d060, float & d061, float & d062, float & d063, + float & d064, float & d065, float & d066, float & d067, + float & d068, float & d069, float & d070, float & d071, + float & d072, float & d073, float & d074, float & d075, + float & d076, float & d077, float & d078, float & d079, + float & d080, float & d081, float & d082, float & d083, + float & d084, float & d085, float & d086, float & d087, + float & d088, float & d089, float & d090, float & d091, + float & d092, float & d093, float & d094, float & d095, + float & d096, float & d097, float & d098, float & d099, + float & d100, float & d101, float & d102, float & d103, + float & d104, float & d105, float & d106, float & d107, + float & d108, float & d109, float & d110, float & d111, + float & d112, float & d113, float & d114, float & d115, + float & d116, float & d117, float & d118, float & d119, + float & d120, float & d121, float & d122, float & d123, + float & d124, float & d125, float & d126, float & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %130, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " p, %131, %132, %133, %134;\n" + "}\n" + : "+f"(d000), "+f"(d001), "+f"(d002), "+f"(d003), + "+f"(d004), "+f"(d005), "+f"(d006), "+f"(d007), + "+f"(d008), "+f"(d009), "+f"(d010), "+f"(d011), + "+f"(d012), "+f"(d013), "+f"(d014), "+f"(d015), + "+f"(d016), "+f"(d017), "+f"(d018), "+f"(d019), + "+f"(d020), "+f"(d021), "+f"(d022), "+f"(d023), + "+f"(d024), "+f"(d025), "+f"(d026), "+f"(d027), + "+f"(d028), "+f"(d029), "+f"(d030), "+f"(d031), + "+f"(d032), "+f"(d033), "+f"(d034), "+f"(d035), + "+f"(d036), "+f"(d037), "+f"(d038), "+f"(d039), + "+f"(d040), "+f"(d041), "+f"(d042), "+f"(d043), + "+f"(d044), "+f"(d045), "+f"(d046), "+f"(d047), + "+f"(d048), "+f"(d049), "+f"(d050), "+f"(d051), + "+f"(d052), "+f"(d053), "+f"(d054), "+f"(d055), + "+f"(d056), "+f"(d057), "+f"(d058), "+f"(d059), + "+f"(d060), "+f"(d061), "+f"(d062), "+f"(d063), + "+f"(d064), "+f"(d065), "+f"(d066), "+f"(d067), + "+f"(d068), "+f"(d069), "+f"(d070), "+f"(d071), + "+f"(d072), "+f"(d073), "+f"(d074), "+f"(d075), + "+f"(d076), "+f"(d077), "+f"(d078), "+f"(d079), + "+f"(d080), "+f"(d081), "+f"(d082), "+f"(d083), + "+f"(d084), "+f"(d085), "+f"(d086), "+f"(d087), + "+f"(d088), "+f"(d089), "+f"(d090), "+f"(d091), + "+f"(d092), "+f"(d093), "+f"(d094), "+f"(d095), + "+f"(d096), "+f"(d097), "+f"(d098), "+f"(d099), + "+f"(d100), "+f"(d101), "+f"(d102), "+f"(d103), + "+f"(d104), "+f"(d105), "+f"(d106), "+f"(d107), + "+f"(d108), "+f"(d109), "+f"(d110), "+f"(d111), + "+f"(d112), "+f"(d113), "+f"(d114), "+f"(d115), + "+f"(d116), "+f"(d117), "+f"(d118), "+f"(d119), + "+f"(d120), "+f"(d121), "+f"(d122), "+f"(d123), + "+f"(d124), "+f"(d125), "+f"(d126), "+f"(d127) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x16_F32F16F16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x16 F32+=F16*F16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x256x16_F32F16F16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[128]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a000, uint32_t const& a001, uint32_t const& a002, uint32_t const& a003, + uint64_t const& desc_b, + float & d000, float & d001, float & d002, float & d003, + float & d004, float & d005, float & d006, float & d007, + float & d008, float & d009, float & d010, float & d011, + float & d012, float & d013, float & d014, float & d015, + float & d016, float & d017, float & d018, float & d019, + float & d020, float & d021, float & d022, float & d023, + float & d024, float & d025, float & d026, float & d027, + float & d028, float & d029, float & d030, float & d031, + float & d032, float & d033, float & d034, float & d035, + float & d036, float & d037, float & d038, float & d039, + float & d040, float & d041, float & d042, float & d043, + float & d044, float & d045, float & d046, float & d047, + float & d048, float & d049, float & d050, float & d051, + float & d052, float & d053, float & d054, float & d055, + float & d056, float & d057, float & d058, float & d059, + float & d060, float & d061, float & d062, float & d063, + float & d064, float & d065, float & d066, float & d067, + float & d068, float & d069, float & d070, float & d071, + float & d072, float & d073, float & d074, float & d075, + float & d076, float & d077, float & d078, float & d079, + float & d080, float & d081, float & d082, float & d083, + float & d084, float & d085, float & d086, float & d087, + float & d088, float & d089, float & d090, float & d091, + float & d092, float & d093, float & d094, float & d095, + float & d096, float & d097, float & d098, float & d099, + float & d100, float & d101, float & d102, float & d103, + float & d104, float & d105, float & d106, float & d107, + float & d108, float & d109, float & d110, float & d111, + float & d112, float & d113, float & d114, float & d115, + float & d116, float & d117, float & d118, float & d119, + float & d120, float & d121, float & d122, float & d123, + float & d124, float & d125, float & d126, float & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %133, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.f16.f16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + "{%128, %129, %130, %131}," + " %132," + " p, %134, %135, %136;\n" + "}\n" + : "+f"(d000), "+f"(d001), "+f"(d002), "+f"(d003), + "+f"(d004), "+f"(d005), "+f"(d006), "+f"(d007), + "+f"(d008), "+f"(d009), "+f"(d010), "+f"(d011), + "+f"(d012), "+f"(d013), "+f"(d014), "+f"(d015), + "+f"(d016), "+f"(d017), "+f"(d018), "+f"(d019), + "+f"(d020), "+f"(d021), "+f"(d022), "+f"(d023), + "+f"(d024), "+f"(d025), "+f"(d026), "+f"(d027), + "+f"(d028), "+f"(d029), "+f"(d030), "+f"(d031), + "+f"(d032), "+f"(d033), "+f"(d034), "+f"(d035), + "+f"(d036), "+f"(d037), "+f"(d038), "+f"(d039), + "+f"(d040), "+f"(d041), "+f"(d042), "+f"(d043), + "+f"(d044), "+f"(d045), "+f"(d046), "+f"(d047), + "+f"(d048), "+f"(d049), "+f"(d050), "+f"(d051), + "+f"(d052), "+f"(d053), "+f"(d054), "+f"(d055), + "+f"(d056), "+f"(d057), "+f"(d058), "+f"(d059), + "+f"(d060), "+f"(d061), "+f"(d062), "+f"(d063), + "+f"(d064), "+f"(d065), "+f"(d066), "+f"(d067), + "+f"(d068), "+f"(d069), "+f"(d070), "+f"(d071), + "+f"(d072), "+f"(d073), "+f"(d074), "+f"(d075), + "+f"(d076), "+f"(d077), "+f"(d078), "+f"(d079), + "+f"(d080), "+f"(d081), "+f"(d082), "+f"(d083), + "+f"(d084), "+f"(d085), "+f"(d086), "+f"(d087), + "+f"(d088), "+f"(d089), "+f"(d090), "+f"(d091), + "+f"(d092), "+f"(d093), "+f"(d094), "+f"(d095), + "+f"(d096), "+f"(d097), "+f"(d098), "+f"(d099), + "+f"(d100), "+f"(d101), "+f"(d102), "+f"(d103), + "+f"(d104), "+f"(d105), "+f"(d106), "+f"(d107), + "+f"(d108), "+f"(d109), "+f"(d110), "+f"(d111), + "+f"(d112), "+f"(d113), "+f"(d114), "+f"(d115), + "+f"(d116), "+f"(d117), "+f"(d118), "+f"(d119), + "+f"(d120), "+f"(d121), "+f"(d122), "+f"(d123), + "+f"(d124), "+f"(d125), "+f"(d126), "+f"(d127) + : "r"(a000), "r"(a001), "r"(a002), "r"(a003), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x16_F32F16F16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x8x16_F32BF16BF16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[4]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d0, float & d1, float & d2, float & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k16.f32.bf16.bf16 " + "{%0, %1, %2, %3}," + " %4," + " %5," + " p, %7, %8, %9, %10;\n" + "}\n" + : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x8x16_F32BF16BF16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[4]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + float & d0, float & d1, float & d2, float & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %9, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k16.f32.bf16.bf16 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + " %8," + " p, %10, %11, %12;\n" + "}\n" + : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x16x16_F32BF16BF16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[8]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d0, float & d1, float & d2, float & d3, + float & d4, float & d5, float & d6, float & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + " %8," + " %9," + " p, %11, %12, %13, %14;\n" + "}\n" + : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3), + "+f"(d4), "+f"(d5), "+f"(d6), "+f"(d7) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x16x16_F32BF16BF16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[8]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + float & d0, float & d1, float & d2, float & d3, + float & d4, float & d5, float & d6, float & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %13, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + "{%8, %9, %10, %11}," + " %12," + " p, %14, %15, %16;\n" + "}\n" + : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3), + "+f"(d4), "+f"(d5), "+f"(d6), "+f"(d7) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x32x16_F32BF16BF16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[16]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %18, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + " %16," + " %17," + " p, %19, %20, %21, %22;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x32x16_F32BF16BF16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[16]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %21, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + "{%16, %17, %18, %19}," + " %20," + " p, %22, %23, %24;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x64x16_F32BF16BF16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[32]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %34, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " p, %35, %36, %37, %38;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x64x16_F32BF16BF16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[32]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %37, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + "{%32, %33, %34, %35}," + " %36," + " p, %38, %39, %40;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x96x16_F32BF16BF16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[48]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %50, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + " %48," + " %49," + " p, %51, %52, %53, %54;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x96x16_F32BF16BF16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[48]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %53, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + "{%48, %49, %50, %51}," + " %52," + " p, %54, %55, %56;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x128x16_F32BF16BF16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[64]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + float & d48, float & d49, float & d50, float & d51, + float & d52, float & d53, float & d54, float & d55, + float & d56, float & d57, float & d58, float & d59, + float & d60, float & d61, float & d62, float & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %66, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " p, %67, %68, %69, %70;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47), + "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), + "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55), + "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), + "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x128x16_F32BF16BF16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[64]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + float & d48, float & d49, float & d50, float & d51, + float & d52, float & d53, float & d54, float & d55, + float & d56, float & d57, float & d58, float & d59, + float & d60, float & d61, float & d62, float & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %69, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + "{%64, %65, %66, %67}," + " %68," + " p, %70, %71, %72;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47), + "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), + "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55), + "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), + "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x192x16_F32BF16BF16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[96]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + float & d48, float & d49, float & d50, float & d51, + float & d52, float & d53, float & d54, float & d55, + float & d56, float & d57, float & d58, float & d59, + float & d60, float & d61, float & d62, float & d63, + float & d64, float & d65, float & d66, float & d67, + float & d68, float & d69, float & d70, float & d71, + float & d72, float & d73, float & d74, float & d75, + float & d76, float & d77, float & d78, float & d79, + float & d80, float & d81, float & d82, float & d83, + float & d84, float & d85, float & d86, float & d87, + float & d88, float & d89, float & d90, float & d91, + float & d92, float & d93, float & d94, float & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %98, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + " %96," + " %97," + " p, %99, %100, %101, %102;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47), + "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), + "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55), + "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), + "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63), + "+f"(d64), "+f"(d65), "+f"(d66), "+f"(d67), + "+f"(d68), "+f"(d69), "+f"(d70), "+f"(d71), + "+f"(d72), "+f"(d73), "+f"(d74), "+f"(d75), + "+f"(d76), "+f"(d77), "+f"(d78), "+f"(d79), + "+f"(d80), "+f"(d81), "+f"(d82), "+f"(d83), + "+f"(d84), "+f"(d85), "+f"(d86), "+f"(d87), + "+f"(d88), "+f"(d89), "+f"(d90), "+f"(d91), + "+f"(d92), "+f"(d93), "+f"(d94), "+f"(d95) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x192x16_F32BF16BF16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[96]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + float & d48, float & d49, float & d50, float & d51, + float & d52, float & d53, float & d54, float & d55, + float & d56, float & d57, float & d58, float & d59, + float & d60, float & d61, float & d62, float & d63, + float & d64, float & d65, float & d66, float & d67, + float & d68, float & d69, float & d70, float & d71, + float & d72, float & d73, float & d74, float & d75, + float & d76, float & d77, float & d78, float & d79, + float & d80, float & d81, float & d82, float & d83, + float & d84, float & d85, float & d86, float & d87, + float & d88, float & d89, float & d90, float & d91, + float & d92, float & d93, float & d94, float & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %101, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + "{%96, %97, %98, %99}," + " %100," + " p, %102, %103, %104;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47), + "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), + "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55), + "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), + "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63), + "+f"(d64), "+f"(d65), "+f"(d66), "+f"(d67), + "+f"(d68), "+f"(d69), "+f"(d70), "+f"(d71), + "+f"(d72), "+f"(d73), "+f"(d74), "+f"(d75), + "+f"(d76), "+f"(d77), "+f"(d78), "+f"(d79), + "+f"(d80), "+f"(d81), "+f"(d82), "+f"(d83), + "+f"(d84), "+f"(d85), "+f"(d86), "+f"(d87), + "+f"(d88), "+f"(d89), "+f"(d90), "+f"(d91), + "+f"(d92), "+f"(d93), "+f"(d94), "+f"(d95) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x256x16_F32BF16BF16_SS +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[128]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d000, float & d001, float & d002, float & d003, + float & d004, float & d005, float & d006, float & d007, + float & d008, float & d009, float & d010, float & d011, + float & d012, float & d013, float & d014, float & d015, + float & d016, float & d017, float & d018, float & d019, + float & d020, float & d021, float & d022, float & d023, + float & d024, float & d025, float & d026, float & d027, + float & d028, float & d029, float & d030, float & d031, + float & d032, float & d033, float & d034, float & d035, + float & d036, float & d037, float & d038, float & d039, + float & d040, float & d041, float & d042, float & d043, + float & d044, float & d045, float & d046, float & d047, + float & d048, float & d049, float & d050, float & d051, + float & d052, float & d053, float & d054, float & d055, + float & d056, float & d057, float & d058, float & d059, + float & d060, float & d061, float & d062, float & d063, + float & d064, float & d065, float & d066, float & d067, + float & d068, float & d069, float & d070, float & d071, + float & d072, float & d073, float & d074, float & d075, + float & d076, float & d077, float & d078, float & d079, + float & d080, float & d081, float & d082, float & d083, + float & d084, float & d085, float & d086, float & d087, + float & d088, float & d089, float & d090, float & d091, + float & d092, float & d093, float & d094, float & d095, + float & d096, float & d097, float & d098, float & d099, + float & d100, float & d101, float & d102, float & d103, + float & d104, float & d105, float & d106, float & d107, + float & d108, float & d109, float & d110, float & d111, + float & d112, float & d113, float & d114, float & d115, + float & d116, float & d117, float & d118, float & d119, + float & d120, float & d121, float & d122, float & d123, + float & d124, float & d125, float & d126, float & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %130, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " p, %131, %132, %133, %134;\n" + "}\n" + : "+f"(d000), "+f"(d001), "+f"(d002), "+f"(d003), + "+f"(d004), "+f"(d005), "+f"(d006), "+f"(d007), + "+f"(d008), "+f"(d009), "+f"(d010), "+f"(d011), + "+f"(d012), "+f"(d013), "+f"(d014), "+f"(d015), + "+f"(d016), "+f"(d017), "+f"(d018), "+f"(d019), + "+f"(d020), "+f"(d021), "+f"(d022), "+f"(d023), + "+f"(d024), "+f"(d025), "+f"(d026), "+f"(d027), + "+f"(d028), "+f"(d029), "+f"(d030), "+f"(d031), + "+f"(d032), "+f"(d033), "+f"(d034), "+f"(d035), + "+f"(d036), "+f"(d037), "+f"(d038), "+f"(d039), + "+f"(d040), "+f"(d041), "+f"(d042), "+f"(d043), + "+f"(d044), "+f"(d045), "+f"(d046), "+f"(d047), + "+f"(d048), "+f"(d049), "+f"(d050), "+f"(d051), + "+f"(d052), "+f"(d053), "+f"(d054), "+f"(d055), + "+f"(d056), "+f"(d057), "+f"(d058), "+f"(d059), + "+f"(d060), "+f"(d061), "+f"(d062), "+f"(d063), + "+f"(d064), "+f"(d065), "+f"(d066), "+f"(d067), + "+f"(d068), "+f"(d069), "+f"(d070), "+f"(d071), + "+f"(d072), "+f"(d073), "+f"(d074), "+f"(d075), + "+f"(d076), "+f"(d077), "+f"(d078), "+f"(d079), + "+f"(d080), "+f"(d081), "+f"(d082), "+f"(d083), + "+f"(d084), "+f"(d085), "+f"(d086), "+f"(d087), + "+f"(d088), "+f"(d089), "+f"(d090), "+f"(d091), + "+f"(d092), "+f"(d093), "+f"(d094), "+f"(d095), + "+f"(d096), "+f"(d097), "+f"(d098), "+f"(d099), + "+f"(d100), "+f"(d101), "+f"(d102), "+f"(d103), + "+f"(d104), "+f"(d105), "+f"(d106), "+f"(d107), + "+f"(d108), "+f"(d109), "+f"(d110), "+f"(d111), + "+f"(d112), "+f"(d113), "+f"(d114), "+f"(d115), + "+f"(d116), "+f"(d117), "+f"(d118), "+f"(d119), + "+f"(d120), "+f"(d121), "+f"(d122), "+f"(d123), + "+f"(d124), "+f"(d125), "+f"(d126), "+f"(d127) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspA)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x16_F32BF16BF16_SS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x16 F32+=BF16*BF16 +template < + GMMA::Major tnspA, + GMMA::Major tnspB, + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x256x16_F32BF16BF16_RS +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[128]; + + static_assert(tnspA == GMMA::Major::K, + "Register source operand A must have K major layout."); + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a000, uint32_t const& a001, uint32_t const& a002, uint32_t const& a003, + uint64_t const& desc_b, + float & d000, float & d001, float & d002, float & d003, + float & d004, float & d005, float & d006, float & d007, + float & d008, float & d009, float & d010, float & d011, + float & d012, float & d013, float & d014, float & d015, + float & d016, float & d017, float & d018, float & d019, + float & d020, float & d021, float & d022, float & d023, + float & d024, float & d025, float & d026, float & d027, + float & d028, float & d029, float & d030, float & d031, + float & d032, float & d033, float & d034, float & d035, + float & d036, float & d037, float & d038, float & d039, + float & d040, float & d041, float & d042, float & d043, + float & d044, float & d045, float & d046, float & d047, + float & d048, float & d049, float & d050, float & d051, + float & d052, float & d053, float & d054, float & d055, + float & d056, float & d057, float & d058, float & d059, + float & d060, float & d061, float & d062, float & d063, + float & d064, float & d065, float & d066, float & d067, + float & d068, float & d069, float & d070, float & d071, + float & d072, float & d073, float & d074, float & d075, + float & d076, float & d077, float & d078, float & d079, + float & d080, float & d081, float & d082, float & d083, + float & d084, float & d085, float & d086, float & d087, + float & d088, float & d089, float & d090, float & d091, + float & d092, float & d093, float & d094, float & d095, + float & d096, float & d097, float & d098, float & d099, + float & d100, float & d101, float & d102, float & d103, + float & d104, float & d105, float & d106, float & d107, + float & d108, float & d109, float & d110, float & d111, + float & d112, float & d113, float & d114, float & d115, + float & d116, float & d117, float & d118, float & d119, + float & d120, float & d121, float & d122, float & d123, + float & d124, float & d125, float & d126, float & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %133, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k16.f32.bf16.bf16 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + "{%128, %129, %130, %131}," + " %132," + " p, %134, %135, %136;\n" + "}\n" + : "+f"(d000), "+f"(d001), "+f"(d002), "+f"(d003), + "+f"(d004), "+f"(d005), "+f"(d006), "+f"(d007), + "+f"(d008), "+f"(d009), "+f"(d010), "+f"(d011), + "+f"(d012), "+f"(d013), "+f"(d014), "+f"(d015), + "+f"(d016), "+f"(d017), "+f"(d018), "+f"(d019), + "+f"(d020), "+f"(d021), "+f"(d022), "+f"(d023), + "+f"(d024), "+f"(d025), "+f"(d026), "+f"(d027), + "+f"(d028), "+f"(d029), "+f"(d030), "+f"(d031), + "+f"(d032), "+f"(d033), "+f"(d034), "+f"(d035), + "+f"(d036), "+f"(d037), "+f"(d038), "+f"(d039), + "+f"(d040), "+f"(d041), "+f"(d042), "+f"(d043), + "+f"(d044), "+f"(d045), "+f"(d046), "+f"(d047), + "+f"(d048), "+f"(d049), "+f"(d050), "+f"(d051), + "+f"(d052), "+f"(d053), "+f"(d054), "+f"(d055), + "+f"(d056), "+f"(d057), "+f"(d058), "+f"(d059), + "+f"(d060), "+f"(d061), "+f"(d062), "+f"(d063), + "+f"(d064), "+f"(d065), "+f"(d066), "+f"(d067), + "+f"(d068), "+f"(d069), "+f"(d070), "+f"(d071), + "+f"(d072), "+f"(d073), "+f"(d074), "+f"(d075), + "+f"(d076), "+f"(d077), "+f"(d078), "+f"(d079), + "+f"(d080), "+f"(d081), "+f"(d082), "+f"(d083), + "+f"(d084), "+f"(d085), "+f"(d086), "+f"(d087), + "+f"(d088), "+f"(d089), "+f"(d090), "+f"(d091), + "+f"(d092), "+f"(d093), "+f"(d094), "+f"(d095), + "+f"(d096), "+f"(d097), "+f"(d098), "+f"(d099), + "+f"(d100), "+f"(d101), "+f"(d102), "+f"(d103), + "+f"(d104), "+f"(d105), "+f"(d106), "+f"(d107), + "+f"(d108), "+f"(d109), "+f"(d110), "+f"(d111), + "+f"(d112), "+f"(d113), "+f"(d114), "+f"(d115), + "+f"(d116), "+f"(d117), "+f"(d118), "+f"(d119), + "+f"(d120), "+f"(d121), "+f"(d122), "+f"(d123), + "+f"(d124), "+f"(d125), "+f"(d126), "+f"(d127) + : "r"(a000), "r"(a001), "r"(a002), "r"(a003), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB)), "n"(int32_t(tnspB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x16_F32BF16BF16_RS without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x8x8_F32TF32TF32_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[4]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d0, float & d1, float & d2, float & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k8.f32.tf32.tf32 " + "{%0, %1, %2, %3}," + " %4," + " %5," + " p, %7, %8;\n" + "}\n" + : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x8x8_F32TF32TF32_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + float & d0, float & d1, float & d2, float & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %9, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k8.f32.tf32.tf32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + " %8," + " p, %10, %11;\n" + "}\n" + : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x16x8_F32TF32TF32_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[8]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d0, float & d1, float & d2, float & d3, + float & d4, float & d5, float & d6, float & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + " %8," + " %9," + " p, %11, %12;\n" + "}\n" + : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3), + "+f"(d4), "+f"(d5), "+f"(d6), "+f"(d7) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x16x8_F32TF32TF32_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[8]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + float & d0, float & d1, float & d2, float & d3, + float & d4, float & d5, float & d6, float & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %13, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + "{%8, %9, %10, %11}," + " %12," + " p, %14, %15;\n" + "}\n" + : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3), + "+f"(d4), "+f"(d5), "+f"(d6), "+f"(d7) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x32x8_F32TF32TF32_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[16]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %18, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + " %16," + " %17," + " p, %19, %20;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x32x8_F32TF32TF32_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[16]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %21, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + "{%16, %17, %18, %19}," + " %20," + " p, %22, %23;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x64x8_F32TF32TF32_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[32]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %34, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " p, %35, %36;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x64x8_F32TF32TF32_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[32]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %37, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + "{%32, %33, %34, %35}," + " %36," + " p, %38, %39;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x96x8_F32TF32TF32_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[48]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %50, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + " %48," + " %49," + " p, %51, %52;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x96x8_F32TF32TF32_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[48]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %53, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + "{%48, %49, %50, %51}," + " %52," + " p, %54, %55;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x128x8_F32TF32TF32_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[64]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + float & d48, float & d49, float & d50, float & d51, + float & d52, float & d53, float & d54, float & d55, + float & d56, float & d57, float & d58, float & d59, + float & d60, float & d61, float & d62, float & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %66, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " p, %67, %68;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47), + "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), + "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55), + "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), + "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x128x8_F32TF32TF32_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[64]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + float & d48, float & d49, float & d50, float & d51, + float & d52, float & d53, float & d54, float & d55, + float & d56, float & d57, float & d58, float & d59, + float & d60, float & d61, float & d62, float & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %69, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + "{%64, %65, %66, %67}," + " %68," + " p, %70, %71;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47), + "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), + "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55), + "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), + "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x192x8_F32TF32TF32_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[96]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + float & d48, float & d49, float & d50, float & d51, + float & d52, float & d53, float & d54, float & d55, + float & d56, float & d57, float & d58, float & d59, + float & d60, float & d61, float & d62, float & d63, + float & d64, float & d65, float & d66, float & d67, + float & d68, float & d69, float & d70, float & d71, + float & d72, float & d73, float & d74, float & d75, + float & d76, float & d77, float & d78, float & d79, + float & d80, float & d81, float & d82, float & d83, + float & d84, float & d85, float & d86, float & d87, + float & d88, float & d89, float & d90, float & d91, + float & d92, float & d93, float & d94, float & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %98, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + " %96," + " %97," + " p, %99, %100;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47), + "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), + "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55), + "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), + "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63), + "+f"(d64), "+f"(d65), "+f"(d66), "+f"(d67), + "+f"(d68), "+f"(d69), "+f"(d70), "+f"(d71), + "+f"(d72), "+f"(d73), "+f"(d74), "+f"(d75), + "+f"(d76), "+f"(d77), "+f"(d78), "+f"(d79), + "+f"(d80), "+f"(d81), "+f"(d82), "+f"(d83), + "+f"(d84), "+f"(d85), "+f"(d86), "+f"(d87), + "+f"(d88), "+f"(d89), "+f"(d90), "+f"(d91), + "+f"(d92), "+f"(d93), "+f"(d94), "+f"(d95) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x192x8_F32TF32TF32_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[96]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + float & d00, float & d01, float & d02, float & d03, + float & d04, float & d05, float & d06, float & d07, + float & d08, float & d09, float & d10, float & d11, + float & d12, float & d13, float & d14, float & d15, + float & d16, float & d17, float & d18, float & d19, + float & d20, float & d21, float & d22, float & d23, + float & d24, float & d25, float & d26, float & d27, + float & d28, float & d29, float & d30, float & d31, + float & d32, float & d33, float & d34, float & d35, + float & d36, float & d37, float & d38, float & d39, + float & d40, float & d41, float & d42, float & d43, + float & d44, float & d45, float & d46, float & d47, + float & d48, float & d49, float & d50, float & d51, + float & d52, float & d53, float & d54, float & d55, + float & d56, float & d57, float & d58, float & d59, + float & d60, float & d61, float & d62, float & d63, + float & d64, float & d65, float & d66, float & d67, + float & d68, float & d69, float & d70, float & d71, + float & d72, float & d73, float & d74, float & d75, + float & d76, float & d77, float & d78, float & d79, + float & d80, float & d81, float & d82, float & d83, + float & d84, float & d85, float & d86, float & d87, + float & d88, float & d89, float & d90, float & d91, + float & d92, float & d93, float & d94, float & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %101, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + "{%96, %97, %98, %99}," + " %100," + " p, %102, %103;\n" + "}\n" + : "+f"(d00), "+f"(d01), "+f"(d02), "+f"(d03), + "+f"(d04), "+f"(d05), "+f"(d06), "+f"(d07), + "+f"(d08), "+f"(d09), "+f"(d10), "+f"(d11), + "+f"(d12), "+f"(d13), "+f"(d14), "+f"(d15), + "+f"(d16), "+f"(d17), "+f"(d18), "+f"(d19), + "+f"(d20), "+f"(d21), "+f"(d22), "+f"(d23), + "+f"(d24), "+f"(d25), "+f"(d26), "+f"(d27), + "+f"(d28), "+f"(d29), "+f"(d30), "+f"(d31), + "+f"(d32), "+f"(d33), "+f"(d34), "+f"(d35), + "+f"(d36), "+f"(d37), "+f"(d38), "+f"(d39), + "+f"(d40), "+f"(d41), "+f"(d42), "+f"(d43), + "+f"(d44), "+f"(d45), "+f"(d46), "+f"(d47), + "+f"(d48), "+f"(d49), "+f"(d50), "+f"(d51), + "+f"(d52), "+f"(d53), "+f"(d54), "+f"(d55), + "+f"(d56), "+f"(d57), "+f"(d58), "+f"(d59), + "+f"(d60), "+f"(d61), "+f"(d62), "+f"(d63), + "+f"(d64), "+f"(d65), "+f"(d66), "+f"(d67), + "+f"(d68), "+f"(d69), "+f"(d70), "+f"(d71), + "+f"(d72), "+f"(d73), "+f"(d74), "+f"(d75), + "+f"(d76), "+f"(d77), "+f"(d78), "+f"(d79), + "+f"(d80), "+f"(d81), "+f"(d82), "+f"(d83), + "+f"(d84), "+f"(d85), "+f"(d86), "+f"(d87), + "+f"(d88), "+f"(d89), "+f"(d90), "+f"(d91), + "+f"(d92), "+f"(d93), "+f"(d94), "+f"(d95) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x256x8_F32TF32TF32_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = float[128]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + float & d000, float & d001, float & d002, float & d003, + float & d004, float & d005, float & d006, float & d007, + float & d008, float & d009, float & d010, float & d011, + float & d012, float & d013, float & d014, float & d015, + float & d016, float & d017, float & d018, float & d019, + float & d020, float & d021, float & d022, float & d023, + float & d024, float & d025, float & d026, float & d027, + float & d028, float & d029, float & d030, float & d031, + float & d032, float & d033, float & d034, float & d035, + float & d036, float & d037, float & d038, float & d039, + float & d040, float & d041, float & d042, float & d043, + float & d044, float & d045, float & d046, float & d047, + float & d048, float & d049, float & d050, float & d051, + float & d052, float & d053, float & d054, float & d055, + float & d056, float & d057, float & d058, float & d059, + float & d060, float & d061, float & d062, float & d063, + float & d064, float & d065, float & d066, float & d067, + float & d068, float & d069, float & d070, float & d071, + float & d072, float & d073, float & d074, float & d075, + float & d076, float & d077, float & d078, float & d079, + float & d080, float & d081, float & d082, float & d083, + float & d084, float & d085, float & d086, float & d087, + float & d088, float & d089, float & d090, float & d091, + float & d092, float & d093, float & d094, float & d095, + float & d096, float & d097, float & d098, float & d099, + float & d100, float & d101, float & d102, float & d103, + float & d104, float & d105, float & d106, float & d107, + float & d108, float & d109, float & d110, float & d111, + float & d112, float & d113, float & d114, float & d115, + float & d116, float & d117, float & d118, float & d119, + float & d120, float & d121, float & d122, float & d123, + float & d124, float & d125, float & d126, float & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %130, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " p, %131, %132;\n" + "}\n" + : "+f"(d000), "+f"(d001), "+f"(d002), "+f"(d003), + "+f"(d004), "+f"(d005), "+f"(d006), "+f"(d007), + "+f"(d008), "+f"(d009), "+f"(d010), "+f"(d011), + "+f"(d012), "+f"(d013), "+f"(d014), "+f"(d015), + "+f"(d016), "+f"(d017), "+f"(d018), "+f"(d019), + "+f"(d020), "+f"(d021), "+f"(d022), "+f"(d023), + "+f"(d024), "+f"(d025), "+f"(d026), "+f"(d027), + "+f"(d028), "+f"(d029), "+f"(d030), "+f"(d031), + "+f"(d032), "+f"(d033), "+f"(d034), "+f"(d035), + "+f"(d036), "+f"(d037), "+f"(d038), "+f"(d039), + "+f"(d040), "+f"(d041), "+f"(d042), "+f"(d043), + "+f"(d044), "+f"(d045), "+f"(d046), "+f"(d047), + "+f"(d048), "+f"(d049), "+f"(d050), "+f"(d051), + "+f"(d052), "+f"(d053), "+f"(d054), "+f"(d055), + "+f"(d056), "+f"(d057), "+f"(d058), "+f"(d059), + "+f"(d060), "+f"(d061), "+f"(d062), "+f"(d063), + "+f"(d064), "+f"(d065), "+f"(d066), "+f"(d067), + "+f"(d068), "+f"(d069), "+f"(d070), "+f"(d071), + "+f"(d072), "+f"(d073), "+f"(d074), "+f"(d075), + "+f"(d076), "+f"(d077), "+f"(d078), "+f"(d079), + "+f"(d080), "+f"(d081), "+f"(d082), "+f"(d083), + "+f"(d084), "+f"(d085), "+f"(d086), "+f"(d087), + "+f"(d088), "+f"(d089), "+f"(d090), "+f"(d091), + "+f"(d092), "+f"(d093), "+f"(d094), "+f"(d095), + "+f"(d096), "+f"(d097), "+f"(d098), "+f"(d099), + "+f"(d100), "+f"(d101), "+f"(d102), "+f"(d103), + "+f"(d104), "+f"(d105), "+f"(d106), "+f"(d107), + "+f"(d108), "+f"(d109), "+f"(d110), "+f"(d111), + "+f"(d112), "+f"(d113), "+f"(d114), "+f"(d115), + "+f"(d116), "+f"(d117), "+f"(d118), "+f"(d119), + "+f"(d120), "+f"(d121), "+f"(d122), "+f"(d123), + "+f"(d124), "+f"(d125), "+f"(d126), "+f"(d127) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x8_F32TF32TF32_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x8 TN F32+=TF32*TF32 +template < + GMMA::ScaleIn scaleA = GMMA::ScaleIn::One, + GMMA::ScaleIn scaleB = GMMA::ScaleIn::One +> +struct SM90_64x256x8_F32TF32TF32_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = float[128]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a000, uint32_t const& a001, uint32_t const& a002, uint32_t const& a003, + uint64_t const& desc_b, + float & d000, float & d001, float & d002, float & d003, + float & d004, float & d005, float & d006, float & d007, + float & d008, float & d009, float & d010, float & d011, + float & d012, float & d013, float & d014, float & d015, + float & d016, float & d017, float & d018, float & d019, + float & d020, float & d021, float & d022, float & d023, + float & d024, float & d025, float & d026, float & d027, + float & d028, float & d029, float & d030, float & d031, + float & d032, float & d033, float & d034, float & d035, + float & d036, float & d037, float & d038, float & d039, + float & d040, float & d041, float & d042, float & d043, + float & d044, float & d045, float & d046, float & d047, + float & d048, float & d049, float & d050, float & d051, + float & d052, float & d053, float & d054, float & d055, + float & d056, float & d057, float & d058, float & d059, + float & d060, float & d061, float & d062, float & d063, + float & d064, float & d065, float & d066, float & d067, + float & d068, float & d069, float & d070, float & d071, + float & d072, float & d073, float & d074, float & d075, + float & d076, float & d077, float & d078, float & d079, + float & d080, float & d081, float & d082, float & d083, + float & d084, float & d085, float & d086, float & d087, + float & d088, float & d089, float & d090, float & d091, + float & d092, float & d093, float & d094, float & d095, + float & d096, float & d097, float & d098, float & d099, + float & d100, float & d101, float & d102, float & d103, + float & d104, float & d105, float & d106, float & d107, + float & d108, float & d109, float & d110, float & d111, + float & d112, float & d113, float & d114, float & d115, + float & d116, float & d117, float & d118, float & d119, + float & d120, float & d121, float & d122, float & d123, + float & d124, float & d125, float & d126, float & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %133, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k8.f32.tf32.tf32 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + "{%128, %129, %130, %131}," + " %132," + " p, %134, %135;\n" + "}\n" + : "+f"(d000), "+f"(d001), "+f"(d002), "+f"(d003), + "+f"(d004), "+f"(d005), "+f"(d006), "+f"(d007), + "+f"(d008), "+f"(d009), "+f"(d010), "+f"(d011), + "+f"(d012), "+f"(d013), "+f"(d014), "+f"(d015), + "+f"(d016), "+f"(d017), "+f"(d018), "+f"(d019), + "+f"(d020), "+f"(d021), "+f"(d022), "+f"(d023), + "+f"(d024), "+f"(d025), "+f"(d026), "+f"(d027), + "+f"(d028), "+f"(d029), "+f"(d030), "+f"(d031), + "+f"(d032), "+f"(d033), "+f"(d034), "+f"(d035), + "+f"(d036), "+f"(d037), "+f"(d038), "+f"(d039), + "+f"(d040), "+f"(d041), "+f"(d042), "+f"(d043), + "+f"(d044), "+f"(d045), "+f"(d046), "+f"(d047), + "+f"(d048), "+f"(d049), "+f"(d050), "+f"(d051), + "+f"(d052), "+f"(d053), "+f"(d054), "+f"(d055), + "+f"(d056), "+f"(d057), "+f"(d058), "+f"(d059), + "+f"(d060), "+f"(d061), "+f"(d062), "+f"(d063), + "+f"(d064), "+f"(d065), "+f"(d066), "+f"(d067), + "+f"(d068), "+f"(d069), "+f"(d070), "+f"(d071), + "+f"(d072), "+f"(d073), "+f"(d074), "+f"(d075), + "+f"(d076), "+f"(d077), "+f"(d078), "+f"(d079), + "+f"(d080), "+f"(d081), "+f"(d082), "+f"(d083), + "+f"(d084), "+f"(d085), "+f"(d086), "+f"(d087), + "+f"(d088), "+f"(d089), "+f"(d090), "+f"(d091), + "+f"(d092), "+f"(d093), "+f"(d094), "+f"(d095), + "+f"(d096), "+f"(d097), "+f"(d098), "+f"(d099), + "+f"(d100), "+f"(d101), "+f"(d102), "+f"(d103), + "+f"(d104), "+f"(d105), "+f"(d106), "+f"(d107), + "+f"(d108), "+f"(d109), "+f"(d110), "+f"(d111), + "+f"(d112), "+f"(d113), "+f"(d114), "+f"(d115), + "+f"(d116), "+f"(d117), "+f"(d118), "+f"(d119), + "+f"(d120), "+f"(d121), "+f"(d122), "+f"(d123), + "+f"(d124), "+f"(d125), "+f"(d126), "+f"(d127) + : "r"(a000), "r"(a001), "r"(a002), "r"(a003), + "l"(desc_b), + "r"(int32_t(scale_D)), "n"(int32_t(scaleA)), "n"(int32_t(scaleB))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x8_F32TF32TF32_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=S8*S8 +struct SM90_64x8x32_S32S8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.s8 " + "{%0, %1, %2, %3}," + " %4," + " %5," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=S8*S8 +struct SM90_64x8x32_S32S8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3}," + " %4," + " %5," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=S8*S8 +struct SM90_64x16x32_S32S8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + " %8," + " %9," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=S8*S8 +struct SM90_64x16x32_S32S8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + " %8," + " %9," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=S8*S8 +struct SM90_64x32x32_S32S8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %18, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + " %16," + " %17," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=S8*S8 +struct SM90_64x32x32_S32S8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %18, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + " %16," + " %17," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=S8*S8 +struct SM90_64x64x32_S32S8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %34, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=S8*S8 +struct SM90_64x64x32_S32S8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %34, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=S8*S8 +struct SM90_64x96x32_S32S8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %50, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + " %48," + " %49," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=S8*S8 +struct SM90_64x96x32_S32S8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %50, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + " %48," + " %49," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=S8*S8 +struct SM90_64x128x32_S32S8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %66, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=S8*S8 +struct SM90_64x128x32_S32S8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %66, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=S8*S8 +struct SM90_64x192x32_S32S8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %98, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + " %96," + " %97," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=S8*S8 +struct SM90_64x192x32_S32S8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %98, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + " %96," + " %97," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=S8*S8 +struct SM90_64x256x32_S32S8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %130, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32S8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=S8*S8 +struct SM90_64x256x32_S32S8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %130, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32S8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=S8*S8 +struct SM90_64x8x32_S32S8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %9, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.s8 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + " %8," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=S8*S8 +struct SM90_64x8x32_S32S8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %9, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + " %8," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=S8*S8 +struct SM90_64x16x32_S32S8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %13, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + "{%8, %9, %10, %11}," + " %12," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=S8*S8 +struct SM90_64x16x32_S32S8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %13, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + "{%8, %9, %10, %11}," + " %12," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=S8*S8 +struct SM90_64x32x32_S32S8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %21, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + "{%16, %17, %18, %19}," + " %20," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=S8*S8 +struct SM90_64x32x32_S32S8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %21, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + "{%16, %17, %18, %19}," + " %20," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=S8*S8 +struct SM90_64x64x32_S32S8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %37, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + "{%32, %33, %34, %35}," + " %36," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=S8*S8 +struct SM90_64x64x32_S32S8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %37, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + "{%32, %33, %34, %35}," + " %36," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=S8*S8 +struct SM90_64x96x32_S32S8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %53, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + "{%48, %49, %50, %51}," + " %52," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=S8*S8 +struct SM90_64x96x32_S32S8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %53, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + "{%48, %49, %50, %51}," + " %52," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=S8*S8 +struct SM90_64x128x32_S32S8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %69, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + "{%64, %65, %66, %67}," + " %68," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=S8*S8 +struct SM90_64x128x32_S32S8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %69, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + "{%64, %65, %66, %67}," + " %68," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=S8*S8 +struct SM90_64x192x32_S32S8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %101, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + "{%96, %97, %98, %99}," + " %100," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=S8*S8 +struct SM90_64x192x32_S32S8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %101, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + "{%96, %97, %98, %99}," + " %100," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=S8*S8 +struct SM90_64x256x32_S32S8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a000, uint32_t const& a001, uint32_t const& a002, uint32_t const& a003, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %133, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.s8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + "{%128, %129, %130, %131}," + " %132," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "r"(a000), "r"(a001), "r"(a002), "r"(a003), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32S8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=S8*S8 +struct SM90_64x256x32_S32S8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a000, uint32_t const& a001, uint32_t const& a002, uint32_t const& a003, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %133, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.s8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + "{%128, %129, %130, %131}," + " %132," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "r"(a000), "r"(a001), "r"(a002), "r"(a003), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32S8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=S8*U8 +struct SM90_64x8x32_S32S8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.u8 " + "{%0, %1, %2, %3}," + " %4," + " %5," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=S8*U8 +struct SM90_64x8x32_S32S8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3}," + " %4," + " %5," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=S8*U8 +struct SM90_64x16x32_S32S8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + " %8," + " %9," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=S8*U8 +struct SM90_64x16x32_S32S8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + " %8," + " %9," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=S8*U8 +struct SM90_64x32x32_S32S8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %18, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + " %16," + " %17," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=S8*U8 +struct SM90_64x32x32_S32S8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %18, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + " %16," + " %17," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=S8*U8 +struct SM90_64x64x32_S32S8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %34, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=S8*U8 +struct SM90_64x64x32_S32S8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %34, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=S8*U8 +struct SM90_64x96x32_S32S8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %50, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + " %48," + " %49," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=S8*U8 +struct SM90_64x96x32_S32S8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %50, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + " %48," + " %49," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=S8*U8 +struct SM90_64x128x32_S32S8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %66, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=S8*U8 +struct SM90_64x128x32_S32S8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %66, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=S8*U8 +struct SM90_64x192x32_S32S8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %98, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + " %96," + " %97," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=S8*U8 +struct SM90_64x192x32_S32S8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %98, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + " %96," + " %97," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=S8*U8 +struct SM90_64x256x32_S32S8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %130, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32S8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=S8*U8 +struct SM90_64x256x32_S32S8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %130, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32S8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=S8*U8 +struct SM90_64x8x32_S32S8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %9, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.u8 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + " %8," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=S8*U8 +struct SM90_64x8x32_S32S8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %9, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + " %8," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=S8*U8 +struct SM90_64x16x32_S32S8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %13, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + "{%8, %9, %10, %11}," + " %12," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=S8*U8 +struct SM90_64x16x32_S32S8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %13, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + "{%8, %9, %10, %11}," + " %12," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=S8*U8 +struct SM90_64x32x32_S32S8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %21, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + "{%16, %17, %18, %19}," + " %20," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=S8*U8 +struct SM90_64x32x32_S32S8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %21, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + "{%16, %17, %18, %19}," + " %20," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=S8*U8 +struct SM90_64x64x32_S32S8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %37, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + "{%32, %33, %34, %35}," + " %36," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=S8*U8 +struct SM90_64x64x32_S32S8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %37, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + "{%32, %33, %34, %35}," + " %36," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=S8*U8 +struct SM90_64x96x32_S32S8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %53, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + "{%48, %49, %50, %51}," + " %52," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=S8*U8 +struct SM90_64x96x32_S32S8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %53, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + "{%48, %49, %50, %51}," + " %52," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=S8*U8 +struct SM90_64x128x32_S32S8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %69, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + "{%64, %65, %66, %67}," + " %68," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=S8*U8 +struct SM90_64x128x32_S32S8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %69, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + "{%64, %65, %66, %67}," + " %68," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=S8*U8 +struct SM90_64x192x32_S32S8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %101, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + "{%96, %97, %98, %99}," + " %100," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=S8*U8 +struct SM90_64x192x32_S32S8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %101, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + "{%96, %97, %98, %99}," + " %100," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=S8*U8 +struct SM90_64x256x32_S32S8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a000, uint32_t const& a001, uint32_t const& a002, uint32_t const& a003, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %133, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.s8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + "{%128, %129, %130, %131}," + " %132," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "r"(a000), "r"(a001), "r"(a002), "r"(a003), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32S8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=S8*U8 +struct SM90_64x256x32_S32S8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a000, uint32_t const& a001, uint32_t const& a002, uint32_t const& a003, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %133, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.s8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + "{%128, %129, %130, %131}," + " %132," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "r"(a000), "r"(a001), "r"(a002), "r"(a003), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32S8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=U8*S8 +struct SM90_64x8x32_S32U8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.u8.s8 " + "{%0, %1, %2, %3}," + " %4," + " %5," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=U8*S8 +struct SM90_64x8x32_S32U8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3}," + " %4," + " %5," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=U8*S8 +struct SM90_64x16x32_S32U8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + " %8," + " %9," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=U8*S8 +struct SM90_64x16x32_S32U8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + " %8," + " %9," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=U8*S8 +struct SM90_64x32x32_S32U8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %18, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + " %16," + " %17," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=U8*S8 +struct SM90_64x32x32_S32U8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %18, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + " %16," + " %17," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=U8*S8 +struct SM90_64x64x32_S32U8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %34, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=U8*S8 +struct SM90_64x64x32_S32U8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %34, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=U8*S8 +struct SM90_64x96x32_S32U8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %50, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + " %48," + " %49," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=U8*S8 +struct SM90_64x96x32_S32U8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %50, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + " %48," + " %49," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=U8*S8 +struct SM90_64x128x32_S32U8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %66, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=U8*S8 +struct SM90_64x128x32_S32U8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %66, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=U8*S8 +struct SM90_64x192x32_S32U8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %98, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + " %96," + " %97," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=U8*S8 +struct SM90_64x192x32_S32U8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %98, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + " %96," + " %97," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=U8*S8 +struct SM90_64x256x32_S32U8S8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %130, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32U8S8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=U8*S8 +struct SM90_64x256x32_S32U8S8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %130, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32U8S8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=U8*S8 +struct SM90_64x8x32_S32U8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %9, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.u8.s8 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + " %8," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=U8*S8 +struct SM90_64x8x32_S32U8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %9, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + " %8," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=U8*S8 +struct SM90_64x16x32_S32U8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %13, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + "{%8, %9, %10, %11}," + " %12," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=U8*S8 +struct SM90_64x16x32_S32U8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %13, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + "{%8, %9, %10, %11}," + " %12," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=U8*S8 +struct SM90_64x32x32_S32U8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %21, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + "{%16, %17, %18, %19}," + " %20," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=U8*S8 +struct SM90_64x32x32_S32U8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %21, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + "{%16, %17, %18, %19}," + " %20," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=U8*S8 +struct SM90_64x64x32_S32U8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %37, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + "{%32, %33, %34, %35}," + " %36," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=U8*S8 +struct SM90_64x64x32_S32U8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %37, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + "{%32, %33, %34, %35}," + " %36," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=U8*S8 +struct SM90_64x96x32_S32U8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %53, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + "{%48, %49, %50, %51}," + " %52," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=U8*S8 +struct SM90_64x96x32_S32U8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %53, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + "{%48, %49, %50, %51}," + " %52," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=U8*S8 +struct SM90_64x128x32_S32U8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %69, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + "{%64, %65, %66, %67}," + " %68," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=U8*S8 +struct SM90_64x128x32_S32U8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %69, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + "{%64, %65, %66, %67}," + " %68," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=U8*S8 +struct SM90_64x192x32_S32U8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %101, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + "{%96, %97, %98, %99}," + " %100," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=U8*S8 +struct SM90_64x192x32_S32U8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %101, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + "{%96, %97, %98, %99}," + " %100," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=U8*S8 +struct SM90_64x256x32_S32U8S8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a000, uint32_t const& a001, uint32_t const& a002, uint32_t const& a003, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %133, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.u8.s8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + "{%128, %129, %130, %131}," + " %132," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "r"(a000), "r"(a001), "r"(a002), "r"(a003), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32U8S8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=U8*S8 +struct SM90_64x256x32_S32U8S8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a000, uint32_t const& a001, uint32_t const& a002, uint32_t const& a003, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %133, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.u8.s8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + "{%128, %129, %130, %131}," + " %132," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "r"(a000), "r"(a001), "r"(a002), "r"(a003), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32U8S8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=U8*U8 +struct SM90_64x8x32_S32U8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.u8.u8 " + "{%0, %1, %2, %3}," + " %4," + " %5," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=U8*U8 +struct SM90_64x8x32_S32U8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %6, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3}," + " %4," + " %5," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=U8*U8 +struct SM90_64x16x32_S32U8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + " %8," + " %9," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=U8*U8 +struct SM90_64x16x32_S32U8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %10, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + " %8," + " %9," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=U8*U8 +struct SM90_64x32x32_S32U8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %18, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + " %16," + " %17," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=U8*U8 +struct SM90_64x32x32_S32U8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %18, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + " %16," + " %17," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=U8*U8 +struct SM90_64x64x32_S32U8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %34, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=U8*U8 +struct SM90_64x64x32_S32U8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %34, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + " %32," + " %33," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=U8*U8 +struct SM90_64x96x32_S32U8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %50, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + " %48," + " %49," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=U8*U8 +struct SM90_64x96x32_S32U8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %50, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + " %48," + " %49," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=U8*U8 +struct SM90_64x128x32_S32U8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %66, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=U8*U8 +struct SM90_64x128x32_S32U8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %66, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + " %64," + " %65," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=U8*U8 +struct SM90_64x192x32_S32U8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %98, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + " %96," + " %97," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=U8*U8 +struct SM90_64x192x32_S32U8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %98, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + " %96," + " %97," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=U8*U8 +struct SM90_64x256x32_S32U8U8_SS_TN +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %130, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32U8U8_SS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=U8*U8 +struct SM90_64x256x32_S32U8U8_SS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint64_t[1]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %130, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + " %128," + " %129," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "l"(desc_a), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32U8U8_SS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=U8*U8 +struct SM90_64x8x32_S32U8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %9, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.u8.u8 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + " %8," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x8x32 TN S32+=U8*U8 +struct SM90_64x8x32_S32U8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[4]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %9, 0;\n" + "wgmma.mma_async.sync.aligned.m64n8k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + " %8," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x8x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=U8*U8 +struct SM90_64x16x32_S32U8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %13, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + "{%8, %9, %10, %11}," + " %12," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x16x32 TN S32+=U8*U8 +struct SM90_64x16x32_S32U8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[8]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3, + uint64_t const& desc_b, + uint32_t & d0, uint32_t & d1, uint32_t & d2, uint32_t & d3, + uint32_t & d4, uint32_t & d5, uint32_t & d6, uint32_t & d7, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %13, 0;\n" + "wgmma.mma_async.sync.aligned.m64n16k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7}," + "{%8, %9, %10, %11}," + " %12," + " p;\n" + "}\n" + : "+r"(d0), "+r"(d1), "+r"(d2), "+r"(d3), + "+r"(d4), "+r"(d5), "+r"(d6), "+r"(d7) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x16x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=U8*U8 +struct SM90_64x32x32_S32U8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %21, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + "{%16, %17, %18, %19}," + " %20," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x32x32 TN S32+=U8*U8 +struct SM90_64x32x32_S32U8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[16]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %21, 0;\n" + "wgmma.mma_async.sync.aligned.m64n32k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15}," + "{%16, %17, %18, %19}," + " %20," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x32x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=U8*U8 +struct SM90_64x64x32_S32U8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %37, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + "{%32, %33, %34, %35}," + " %36," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x64x32 TN S32+=U8*U8 +struct SM90_64x64x32_S32U8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[32]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %37, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31}," + "{%32, %33, %34, %35}," + " %36," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x64x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=U8*U8 +struct SM90_64x96x32_S32U8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %53, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + "{%48, %49, %50, %51}," + " %52," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x96x32 TN S32+=U8*U8 +struct SM90_64x96x32_S32U8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[48]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %53, 0;\n" + "wgmma.mma_async.sync.aligned.m64n96k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47}," + "{%48, %49, %50, %51}," + " %52," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x96x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=U8*U8 +struct SM90_64x128x32_S32U8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %69, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + "{%64, %65, %66, %67}," + " %68," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x128x32 TN S32+=U8*U8 +struct SM90_64x128x32_S32U8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[64]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %69, 0;\n" + "wgmma.mma_async.sync.aligned.m64n128k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63}," + "{%64, %65, %66, %67}," + " %68," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x128x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=U8*U8 +struct SM90_64x192x32_S32U8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %101, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + "{%96, %97, %98, %99}," + " %100," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x192x32 TN S32+=U8*U8 +struct SM90_64x192x32_S32U8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[96]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a00, uint32_t const& a01, uint32_t const& a02, uint32_t const& a03, + uint64_t const& desc_b, + uint32_t & d00, uint32_t & d01, uint32_t & d02, uint32_t & d03, + uint32_t & d04, uint32_t & d05, uint32_t & d06, uint32_t & d07, + uint32_t & d08, uint32_t & d09, uint32_t & d10, uint32_t & d11, + uint32_t & d12, uint32_t & d13, uint32_t & d14, uint32_t & d15, + uint32_t & d16, uint32_t & d17, uint32_t & d18, uint32_t & d19, + uint32_t & d20, uint32_t & d21, uint32_t & d22, uint32_t & d23, + uint32_t & d24, uint32_t & d25, uint32_t & d26, uint32_t & d27, + uint32_t & d28, uint32_t & d29, uint32_t & d30, uint32_t & d31, + uint32_t & d32, uint32_t & d33, uint32_t & d34, uint32_t & d35, + uint32_t & d36, uint32_t & d37, uint32_t & d38, uint32_t & d39, + uint32_t & d40, uint32_t & d41, uint32_t & d42, uint32_t & d43, + uint32_t & d44, uint32_t & d45, uint32_t & d46, uint32_t & d47, + uint32_t & d48, uint32_t & d49, uint32_t & d50, uint32_t & d51, + uint32_t & d52, uint32_t & d53, uint32_t & d54, uint32_t & d55, + uint32_t & d56, uint32_t & d57, uint32_t & d58, uint32_t & d59, + uint32_t & d60, uint32_t & d61, uint32_t & d62, uint32_t & d63, + uint32_t & d64, uint32_t & d65, uint32_t & d66, uint32_t & d67, + uint32_t & d68, uint32_t & d69, uint32_t & d70, uint32_t & d71, + uint32_t & d72, uint32_t & d73, uint32_t & d74, uint32_t & d75, + uint32_t & d76, uint32_t & d77, uint32_t & d78, uint32_t & d79, + uint32_t & d80, uint32_t & d81, uint32_t & d82, uint32_t & d83, + uint32_t & d84, uint32_t & d85, uint32_t & d86, uint32_t & d87, + uint32_t & d88, uint32_t & d89, uint32_t & d90, uint32_t & d91, + uint32_t & d92, uint32_t & d93, uint32_t & d94, uint32_t & d95, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %101, 0;\n" + "wgmma.mma_async.sync.aligned.m64n192k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95}," + "{%96, %97, %98, %99}," + " %100," + " p;\n" + "}\n" + : "+r"(d00), "+r"(d01), "+r"(d02), "+r"(d03), + "+r"(d04), "+r"(d05), "+r"(d06), "+r"(d07), + "+r"(d08), "+r"(d09), "+r"(d10), "+r"(d11), + "+r"(d12), "+r"(d13), "+r"(d14), "+r"(d15), + "+r"(d16), "+r"(d17), "+r"(d18), "+r"(d19), + "+r"(d20), "+r"(d21), "+r"(d22), "+r"(d23), + "+r"(d24), "+r"(d25), "+r"(d26), "+r"(d27), + "+r"(d28), "+r"(d29), "+r"(d30), "+r"(d31), + "+r"(d32), "+r"(d33), "+r"(d34), "+r"(d35), + "+r"(d36), "+r"(d37), "+r"(d38), "+r"(d39), + "+r"(d40), "+r"(d41), "+r"(d42), "+r"(d43), + "+r"(d44), "+r"(d45), "+r"(d46), "+r"(d47), + "+r"(d48), "+r"(d49), "+r"(d50), "+r"(d51), + "+r"(d52), "+r"(d53), "+r"(d54), "+r"(d55), + "+r"(d56), "+r"(d57), "+r"(d58), "+r"(d59), + "+r"(d60), "+r"(d61), "+r"(d62), "+r"(d63), + "+r"(d64), "+r"(d65), "+r"(d66), "+r"(d67), + "+r"(d68), "+r"(d69), "+r"(d70), "+r"(d71), + "+r"(d72), "+r"(d73), "+r"(d74), "+r"(d75), + "+r"(d76), "+r"(d77), "+r"(d78), "+r"(d79), + "+r"(d80), "+r"(d81), "+r"(d82), "+r"(d83), + "+r"(d84), "+r"(d85), "+r"(d86), "+r"(d87), + "+r"(d88), "+r"(d89), "+r"(d90), "+r"(d91), + "+r"(d92), "+r"(d93), "+r"(d94), "+r"(d95) + : "r"(a00), "r"(a01), "r"(a02), "r"(a03), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x192x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=U8*U8 +struct SM90_64x256x32_S32U8U8_RS_TN +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a000, uint32_t const& a001, uint32_t const& a002, uint32_t const& a003, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %133, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.u8.u8 " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + "{%128, %129, %130, %131}," + " %132," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "r"(a000), "r"(a001), "r"(a002), "r"(a003), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32U8U8_RS_TN without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA 64x256x32 TN S32+=U8*U8 +struct SM90_64x256x32_S32U8U8_RS_TN_SATURATE +{ + using DRegisters = void; + using ARegisters = uint32_t[4]; + using BRegisters = uint64_t[1]; + using CRegisters = uint32_t[128]; + + CUTE_HOST_DEVICE static void + fma(uint32_t const& a000, uint32_t const& a001, uint32_t const& a002, uint32_t const& a003, + uint64_t const& desc_b, + uint32_t & d000, uint32_t & d001, uint32_t & d002, uint32_t & d003, + uint32_t & d004, uint32_t & d005, uint32_t & d006, uint32_t & d007, + uint32_t & d008, uint32_t & d009, uint32_t & d010, uint32_t & d011, + uint32_t & d012, uint32_t & d013, uint32_t & d014, uint32_t & d015, + uint32_t & d016, uint32_t & d017, uint32_t & d018, uint32_t & d019, + uint32_t & d020, uint32_t & d021, uint32_t & d022, uint32_t & d023, + uint32_t & d024, uint32_t & d025, uint32_t & d026, uint32_t & d027, + uint32_t & d028, uint32_t & d029, uint32_t & d030, uint32_t & d031, + uint32_t & d032, uint32_t & d033, uint32_t & d034, uint32_t & d035, + uint32_t & d036, uint32_t & d037, uint32_t & d038, uint32_t & d039, + uint32_t & d040, uint32_t & d041, uint32_t & d042, uint32_t & d043, + uint32_t & d044, uint32_t & d045, uint32_t & d046, uint32_t & d047, + uint32_t & d048, uint32_t & d049, uint32_t & d050, uint32_t & d051, + uint32_t & d052, uint32_t & d053, uint32_t & d054, uint32_t & d055, + uint32_t & d056, uint32_t & d057, uint32_t & d058, uint32_t & d059, + uint32_t & d060, uint32_t & d061, uint32_t & d062, uint32_t & d063, + uint32_t & d064, uint32_t & d065, uint32_t & d066, uint32_t & d067, + uint32_t & d068, uint32_t & d069, uint32_t & d070, uint32_t & d071, + uint32_t & d072, uint32_t & d073, uint32_t & d074, uint32_t & d075, + uint32_t & d076, uint32_t & d077, uint32_t & d078, uint32_t & d079, + uint32_t & d080, uint32_t & d081, uint32_t & d082, uint32_t & d083, + uint32_t & d084, uint32_t & d085, uint32_t & d086, uint32_t & d087, + uint32_t & d088, uint32_t & d089, uint32_t & d090, uint32_t & d091, + uint32_t & d092, uint32_t & d093, uint32_t & d094, uint32_t & d095, + uint32_t & d096, uint32_t & d097, uint32_t & d098, uint32_t & d099, + uint32_t & d100, uint32_t & d101, uint32_t & d102, uint32_t & d103, + uint32_t & d104, uint32_t & d105, uint32_t & d106, uint32_t & d107, + uint32_t & d108, uint32_t & d109, uint32_t & d110, uint32_t & d111, + uint32_t & d112, uint32_t & d113, uint32_t & d114, uint32_t & d115, + uint32_t & d116, uint32_t & d117, uint32_t & d118, uint32_t & d119, + uint32_t & d120, uint32_t & d121, uint32_t & d122, uint32_t & d123, + uint32_t & d124, uint32_t & d125, uint32_t & d126, uint32_t & d127, + GMMA::ScaleOut const scale_D = GMMA::ScaleOut::One) + { +#if defined(CUTE_ARCH_MMA_SM90A_ENABLED) + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %133, 0;\n" + "wgmma.mma_async.sync.aligned.m64n256k32.s32.u8.u8.satfinite " + "{%0, %1, %2, %3, %4, %5, %6, %7, " + " %8, %9, %10, %11, %12, %13, %14, %15, " + " %16, %17, %18, %19, %20, %21, %22, %23, " + " %24, %25, %26, %27, %28, %29, %30, %31, " + " %32, %33, %34, %35, %36, %37, %38, %39, " + " %40, %41, %42, %43, %44, %45, %46, %47, " + " %48, %49, %50, %51, %52, %53, %54, %55, " + " %56, %57, %58, %59, %60, %61, %62, %63, " + " %64, %65, %66, %67, %68, %69, %70, %71, " + " %72, %73, %74, %75, %76, %77, %78, %79, " + " %80, %81, %82, %83, %84, %85, %86, %87, " + " %88, %89, %90, %91, %92, %93, %94, %95, " + " %96, %97, %98, %99, %100, %101, %102, %103, " + " %104, %105, %106, %107, %108, %109, %110, %111, " + " %112, %113, %114, %115, %116, %117, %118, %119, " + " %120, %121, %122, %123, %124, %125, %126, %127}," + "{%128, %129, %130, %131}," + " %132," + " p;\n" + "}\n" + : "+r"(d000), "+r"(d001), "+r"(d002), "+r"(d003), + "+r"(d004), "+r"(d005), "+r"(d006), "+r"(d007), + "+r"(d008), "+r"(d009), "+r"(d010), "+r"(d011), + "+r"(d012), "+r"(d013), "+r"(d014), "+r"(d015), + "+r"(d016), "+r"(d017), "+r"(d018), "+r"(d019), + "+r"(d020), "+r"(d021), "+r"(d022), "+r"(d023), + "+r"(d024), "+r"(d025), "+r"(d026), "+r"(d027), + "+r"(d028), "+r"(d029), "+r"(d030), "+r"(d031), + "+r"(d032), "+r"(d033), "+r"(d034), "+r"(d035), + "+r"(d036), "+r"(d037), "+r"(d038), "+r"(d039), + "+r"(d040), "+r"(d041), "+r"(d042), "+r"(d043), + "+r"(d044), "+r"(d045), "+r"(d046), "+r"(d047), + "+r"(d048), "+r"(d049), "+r"(d050), "+r"(d051), + "+r"(d052), "+r"(d053), "+r"(d054), "+r"(d055), + "+r"(d056), "+r"(d057), "+r"(d058), "+r"(d059), + "+r"(d060), "+r"(d061), "+r"(d062), "+r"(d063), + "+r"(d064), "+r"(d065), "+r"(d066), "+r"(d067), + "+r"(d068), "+r"(d069), "+r"(d070), "+r"(d071), + "+r"(d072), "+r"(d073), "+r"(d074), "+r"(d075), + "+r"(d076), "+r"(d077), "+r"(d078), "+r"(d079), + "+r"(d080), "+r"(d081), "+r"(d082), "+r"(d083), + "+r"(d084), "+r"(d085), "+r"(d086), "+r"(d087), + "+r"(d088), "+r"(d089), "+r"(d090), "+r"(d091), + "+r"(d092), "+r"(d093), "+r"(d094), "+r"(d095), + "+r"(d096), "+r"(d097), "+r"(d098), "+r"(d099), + "+r"(d100), "+r"(d101), "+r"(d102), "+r"(d103), + "+r"(d104), "+r"(d105), "+r"(d106), "+r"(d107), + "+r"(d108), "+r"(d109), "+r"(d110), "+r"(d111), + "+r"(d112), "+r"(d113), "+r"(d114), "+r"(d115), + "+r"(d116), "+r"(d117), "+r"(d118), "+r"(d119), + "+r"(d120), "+r"(d121), "+r"(d122), "+r"(d123), + "+r"(d124), "+r"(d125), "+r"(d126), "+r"(d127) + : "r"(a000), "r"(a001), "r"(a002), "r"(a003), + "l"(desc_b), + "r"(int32_t(scale_D))); +#else + CUTE_RUNTIME_ASSERT("Attempting to use SM90_64x256x32_S32U8U8_RS_TN_SATURATE without CUTE_ARCH_MMA_SM90A_ENABLED"); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cute diff --git a/csrc/mctlass/include/cute/arch/util.hpp b/csrc/mctlass/include/cute/arch/util.hpp new file mode 100644 index 0000000..b783c2b --- /dev/null +++ b/csrc/mctlass/include/cute/arch/util.hpp @@ -0,0 +1,249 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +#if defined(__clang__) && defined(__MACA__) + // __cvta_generic_to_shared was added in Clang 14: https://reviews.llvm.org/D111665 + #if __clang_major__ >= 14 + #define CUTE_CLANG_SUPPORTS_CVTA_GENERIC_TO_SHARED 1 + #endif + + // __nvvm_get_smem_pointer added in Clang 14: https://reviews.llvm.org/D111665 + // ... but will not work on Windows until Clang 15: https://reviews.llvm.org/D122897 + #if (!defined(_WIN32) && __clang_major__ >= 14) || __clang_major__ >= 15 + #define CUTE_CLANG_SUPPORTS_NVVM_GET_SMEM_POINTER 1 + #endif +#endif + +#if defined(__MXCC__) || defined(__MACACC_RTC__) + #if 1 + #define CUTE_MACA_SUPPORTS_CVTA_GENERIC_TO_SHARED 1 + #endif + + #if 1 + #define CUTE_MACA_SUPPORTS_NVVM_GET_SMEM_POINTER 1 + #endif +#endif + +#if CUTE_MACA_SUPPORTS_CVTA_GENERIC_TO_SHARED || CUTE_CLANG_SUPPORTS_CVTA_GENERIC_TO_SHARED + #define CUTE_CVTA_GENERIC_TO_SHARED_SUPPORTED 1 +#endif + +#if !defined(CUTE_CVTA_GENERIC_TO_SHARED_ACTIVATED) && CUTE_CVTA_GENERIC_TO_SHARED_SUPPORTED && defined(__MACA_ARCH__) + #define CUTE_CVTA_GENERIC_TO_SHARED_ACTIVATED 1 +#endif + +#if CUTE_MACA_SUPPORTS_NVVM_GET_SMEM_POINTER || CUTE_CLANG_SUPPORTS_NVVM_GET_SMEM_POINTER + #define CUTE_NVVM_GET_SMEM_POINTER_SUPPORTED 1 +#endif + +#if !defined(CUTE_NVVM_GET_SMEM_POINTER_ACTIVATED) && CUTE_NVVM_GET_SMEM_POINTER_SUPPORTED && defined(__MACA_ARCH__) + #define CUTE_NVVM_GET_SMEM_POINTER_ACTIVATED 1 +#endif + +// Clang 14+ provides a declaration of __nvvm_get_smem_pointer, so we only need +// to provide one for NVCC +#if CUTE_MACA_SUPPORTS_NVVM_GET_SMEM_POINTER + extern "C" { + // This NVVM intrinsic is subject to change in future versions of CUDA. + // Clients should not call it directly. + CUTE_DEVICE uint32_t __nvvm_get_smem_pointer(void*); + } +#endif + +namespace cute +{ + +/// CUTE helper to cast SMEM pointer to unsigned +CUTE_DEVICE +uint32_t +cast_smem_ptr_to_uint(void const* const ptr) +{ +// We prefer to use the new CVTA intrinsics if they are available, otherwise we will fall back to +// the previous internal intrinsics if they are available. +#if CUTE_CVTA_GENERIC_TO_SHARED_ACTIVATED + // + // This NVVM intrinsic converts an address in shared memory to a plain + // unsigned integer. This is necessary to pass to shared memory instructions + // in inline PTX. + // + // In CUDA 11 and beyond, this replaces __nvvm_get_smem_pointer() [only available in 10.2]. + // + //__device__ size_t __cvta_generic_to_shared(void* ptr); + + /// CUTE helper to get SMEM pointer + return static_cast(__cvta_generic_to_shared(ptr)); + +#elif CUTE_NVVM_GET_SMEM_POINTER_ACTIVATED + + return __nvvm_get_smem_pointer(ptr); + +#elif defined(__MACA_ARCH__) + + uint32_t smem_ptr; + + return smem_ptr; + +#else + + + (void) ptr; + printf("ERROR: cast_smem_ptr_to_uint not supported but used.\n"); + return 0; + +#endif +} + +// +// Utility for pointer interfaces +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +void +explode(Fn fn, + PtrS&& s, int_sequence, + PtrD&& d, int_sequence) +{ + return fn(s[Is]..., d[Id]...); +} + +template +CUTE_HOST_DEVICE constexpr +void +explode(Fn fn, + PtrA&& a, int_sequence, + PtrB&& b, int_sequence, + PtrC&& c, int_sequence) +{ + return fn(a[Ia]..., b[Ib]..., c[Ic]...); +} + +template +CUTE_HOST_DEVICE constexpr +void +explode(Fn fn, + PtrD&& d, int_sequence, + PtrA&& a, int_sequence, + PtrB&& b, int_sequence, + PtrC&& c, int_sequence) +{ + return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]...); +} + +template +CUTE_HOST_DEVICE constexpr +void +explode_with_d_scaling(Fn fn, + PtrA&& a, int_sequence, + PtrB&& b, int_sequence, + PtrC&& c, int_sequence, + ParamType&& p0) +{ + return fn(a[Ia]..., b[Ib]..., c[Ic]..., p0); +} + +template +CUTE_HOST_DEVICE constexpr +void +explode_with_d_scaling(Fn fn, + PtrD&& d, int_sequence, + PtrA&& a, int_sequence, + PtrB&& b, int_sequence, + PtrC&& c, int_sequence, + ParamType&& p0) +{ + return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., p0); +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +void +explode(Fn fn, PtrS&& s, PtrD&& d) +{ + return detail::explode(fn, + s, make_int_sequence{}, + d, make_int_sequence{}); +} + +template +CUTE_HOST_DEVICE constexpr +void +explode(Fn fn, PtrA&& a, PtrB&& b, PtrC&& c) +{ + return detail::explode(fn, + a, make_int_sequence{}, + b, make_int_sequence{}, + c, make_int_sequence{}); +} + +template +CUTE_HOST_DEVICE constexpr +void +explode(Fn fn, PtrD&& d, PtrA&& a, PtrB&& b, PtrC&& c) +{ + return detail::explode(fn, + d, make_int_sequence{}, + a, make_int_sequence{}, + b, make_int_sequence{}, + c, make_int_sequence{}); +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/atom/copy_atom.hpp b/csrc/mctlass/include/cute/atom/copy_atom.hpp new file mode 100644 index 0000000..9180756 --- /dev/null +++ b/csrc/mctlass/include/cute/atom/copy_atom.hpp @@ -0,0 +1,707 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +#include + +#include + +#include + +#include + +namespace cute +{ + +template +struct Copy_Atom; + +template +struct Copy_Atom : Copy_Atom, T> +{}; + +template +struct Copy_Atom, T> + : Copy_Traits +{ + using Traits = Copy_Traits; + + // Bit and Thr layouts from the Copy_Traits + using ThrID = typename Traits::ThrID; + using BitLayoutSrc = typename Traits::SrcLayout; + using BitLayoutDst = typename Traits::DstLayout; + using BitLayoutRef = typename Traits::RefLayout; + + using ValType = T; + + using ValLayoutSrc = decltype(upcast::value>(BitLayoutSrc{})); + using ValLayoutDst = decltype(upcast::value>(BitLayoutDst{})); + using ValLayoutRef = decltype(upcast::value>(BitLayoutRef{})); + + CUTE_STATIC_ASSERT_V(size<0>(ValLayoutSrc{}) == size(ThrID{}), "CopyOperation is not valid for Src of ValType."); + CUTE_STATIC_ASSERT_V(size<0>(ValLayoutDst{}) == size(ThrID{}), "CopyOperation is not valid for Dst of ValType."); + CUTE_STATIC_ASSERT_V(size<0>(ValLayoutRef{}) == size(ThrID{}), "CopyOperation is not valid for Ref of ValType."); + + static constexpr int NumValSrc = size<1>(ValLayoutSrc{}); + static constexpr int NumValDst = size<1>(ValLayoutDst{}); + + // Additional Trait parameters/transformations + template + CUTE_HOST_DEVICE + auto + with(TraitsArgs&&... args) const { + auto traits = Traits::with(std::forward(args)...); + return Copy_Atom{traits}; + } + + // + // Tensor call interfaces + // + + // Check and call instruction, or recurse + template + CUTE_HOST_DEVICE + void + call(Tensor const& src, + Tensor & dst) const + { + static_assert(SLayout::rank == 1, "Expected rank-1 src tensor"); + static_assert(DLayout::rank == 1, "Expected rank-1 dst tensor"); + + if constexpr (is_constant::value || + is_constant::value) { + // Dispatch to unpack for instruction + return copy_unpack(*this, src, dst); + } else + if constexpr (is_tuple::value && + is_tuple::value) { + // If the size of the src/dst doesn't match the instruction, + // recurse this rank-1 layout by peeling off the mode + // ((A,B,C,...)) -> (A,B,C,...) + return copy(*this, tensor<0>(src), tensor<0>(dst)); + } else { + static_assert(sizeof(TS) < 0, "No instruction match and no recursion possible."); + } + } + + // Accept mutable temporaries + template + CUTE_HOST_DEVICE + void + call(Tensor const& src, + Tensor && dst) const + { + return call(src, dst); + } +}; + +// +// A tiling of copy atoms +// + +template +struct ThrCopy; + +template coord [Need not be 2D...] + class ShapeTile_MN> // coord space +struct TiledCopy : Copy_Atom +{ + // Layout information from the CopyAtom + using AtomThrID = typename Copy_Atom::ThrID; // thrid -> thr_idx + using AtomLayoutSrc = typename Copy_Atom::ValLayoutSrc; // (thr,val) -> offset + using AtomLayoutDst = typename Copy_Atom::ValLayoutDst; // (thr,val) -> offset + using AtomLayoutRef = typename Copy_Atom::ValLayoutRef; // (thr,val) -> offset + + using AtomNumThr = decltype(size<0>(AtomLayoutRef{})); + using AtomNumVal = decltype(size<1>(AtomLayoutRef{})); + + // Layout information for the TiledCopy + using Tiler_MN = ShapeTile_MN; + using TiledShape_MN = decltype(shape(ShapeTile_MN{})); + using TiledLayout_TV = LayoutCopy_TV; + using TiledNumThr = decltype(size<0>(TiledLayout_TV{})); + using TiledNumVal = decltype(size<1>(TiledLayout_TV{})); + + CUTE_STATIC_ASSERT_V(TiledNumThr{} % AtomNumThr{} == Int<0>{}, "TiledCopy uses too few thrs for selected CopyAtom"); + CUTE_STATIC_ASSERT_V(TiledNumVal{} % AtomNumVal{} == Int<0>{}, "TiledCopy uses too few vals for selected CopyAtom"); + + // Tile a tensor or a layout from shape + // (M,N,...) + // to shape + // ((ThrV,ThrX),FrgV,(RestM,RestN,...)) + // where + // ThrV: The threads local to a COPY_ATOM Src. + // ThrX: The threads tiled across COPY_ATOMs Src. + // FrgV: The values local to a COPY_ATOM Src. + // RestM: The values tiled in M. + // RestN: The values tiled in N. + template + CUTE_HOST_DEVICE constexpr static + auto + tidfrg_S(STensor&& stensor) + { + constexpr int R = remove_cvref_t::rank; + static_assert(R >= rank_v, "Rank of tensor to be partitioned too small."); + // Generalize the dimension checks for arbitrary rank + //CUTE_STATIC_ASSERT_V(size<0>(stensor) % size<0>(TiledShape_MNK{}) == Int<0>{}); + //CUTE_STATIC_ASSERT_V(size<1>(stensor) % size<1>(TiledShape_MNK{}) == Int<0>{}); + + return tile2thrfrg(zipped_divide(stensor,Tiler_MN{}), right_inverse(AtomLayoutRef{}).compose(AtomLayoutSrc{})); + } + + // Tile a tensor or a layout from shape + // (M,N,...) + // to shape + // ((ThrV,ThrX),FrgV,(RestM,RestN,...)) + // where + // ThrV: The threads local to a COPY_ATOM Dst. + // ThrX: The threads tiled across COPY_ATOMs Dst. + // FrgV: The values local to a COPY_ATOM Dst. + // RestM: The values tiled in M. + // RestN: The values tiled in N. + template + CUTE_HOST_DEVICE constexpr static + auto + tidfrg_D(DTensor&& dtensor) + { + constexpr int R = remove_cvref_t::rank; + static_assert(R >= rank_v, "Rank of tensor to be partitioned too small."); + // Generalize the dimension checks for arbitrary rank + //CUTE_STATIC_ASSERT_V(size<0>(stensor) % size<0>(TiledShape_MNK{}) == Int<0>{}); + //CUTE_STATIC_ASSERT_V(size<1>(stensor) % size<1>(TiledShape_MNK{}) == Int<0>{}); + + return tile2thrfrg(zipped_divide(dtensor,Tiler_MN{}), right_inverse(AtomLayoutRef{}).compose(AtomLayoutDst{})); + } + + // Tile a tensor or a layout from shape + // (Tile,(RestM,RestN,...)) + // to shape + // ((ThrV,ThrX),FrgV,(RestM,RestN,...)) + template + CUTE_HOST_DEVICE constexpr static + auto + tile2thrfrg(Tensor&& tensor, Ref2TrgLayout const& ref2trg) + { + // Take the thrs/vals that the atom is interested in + // NOTE: Assumes the AtomNumThr are contiguous and identity within TiledThrID + auto atom_layout_TV = zipped_divide(TiledLayout_TV{}, make_shape(AtomNumThr{}, AtomNumVal{})); + // ((atom_tid,atom_val),(rest_tid,rest_val)) -> (m,n) + + // Transform to the trg layout + auto trg_layout_TV = atom_layout_TV.compose(ref2trg, _); + // ((trg_tid,trg_val),(rest_tid,rest_val)) -> (m,n) + + // Transform the thrs mode from thrid to thr_idx + // NOTE: Assumes the AtomNumThr are contiguous and identity within TiledThrID + auto thrval2mn = coalesce(zip(trg_layout_TV), Shape<_1,Shape<_1,_1>>{}); + // ((trg_tid,rest_tid),(trg_val,rest_val)) -> (m,n) + + /// ================== + + // Transform the tile mode + auto tv_tensor = tensor.compose(thrval2mn, _); + // ((thrid,val),(RM,RN,...)) + + // Unfold and return + return tv_tensor(make_coord(_,_), _); + } + + // retile_S and retile_D assume they are working with the reference layout -- they are the same + template + CUTE_HOST_DEVICE constexpr static + auto + retile(Tensor&& tensor) + { + constexpr int R = remove_cvref_t::rank; + // Assert that AtomLayoutSrc|Dst is identity so we can skip the Ref transformation + + // Assume the first size<0>(tensor) elements are the first val_ids in TiledLayout_TV. + // Then, we only need the shape+layout of those size<0>(tensor) elements in TiledLayout_TV + // and that shape is what we gather from the other modes of tensor + + auto V = size<0>(tensor); + + auto frg_layout_mn = upcast(right_inverse(TiledLayout_TV{}).with_shape(TiledShape_MN{})); + // (m,n) -> v_idx -- The shape and order of the V inside of TiledLayout_TV + + auto frg_layout_v = zipped_divide(logical_product(make_layout(V), right_inverse(frg_layout_mn)), make_layout(AtomNumVal{})); + // (atom_vals,rest_vals) -> (v,m,n) + + /// ======= + + // Tile the tensor for TileFrg + auto t_tensor = zipped_divide(tensor, prepend(product_each(shape(frg_layout_mn)), V)); + // ((TileV,TileM,TileN,...),(1,RestM,RestN,...)) + + // Transform the tile mode + auto v_tensor = t_tensor.compose(frg_layout_v, _); + // ((atom_vals,rest_vals),(1,RM,RN,...)) + + // Unfold and return + return v_tensor(_, append(Int<0>{},_)); + } + + CUTE_HOST_DEVICE constexpr static + auto + get_layoutS_TV() + { + // (M,N) -> (M,N) + auto ref_S = make_layout(make_shape(TiledShape_MN{}, Int<1>{})); + // (thr_idx,val_idx) -> (M,N) + return tile2thrfrg(ref_S, right_inverse(AtomLayoutRef{}).compose(AtomLayoutSrc{}))(_,_,Int<0>{}); + } + + CUTE_HOST_DEVICE constexpr static + auto + get_layoutS_MN() + { + // (thr_idx,val_idx) -> (M,N) + auto layoutS_TV = get_layoutS_TV(); + // (M,K) -> (thr_idx,val_idx) + auto layoutS_MK = right_inverse(layoutS_TV).with_shape(TiledShape_MN{}); + + // athrid = (v,m,k) -> thr_idx + auto thrID_S = make_layout(size<0>(TiledLayout_TV{})); + + return cute::make_tuple(layoutS_MK, thrID_S); + } + + CUTE_HOST_DEVICE constexpr static + auto + get_layoutD_TV() + { + // (M,N) -> (M,N) + auto ref_D = make_layout(make_shape(TiledShape_MN{}, Int<1>{})); + // (thr_idx,val_idx) -> (M,N) + return tile2thrfrg(ref_D, right_inverse(AtomLayoutRef{}).compose(AtomLayoutDst{}))(_,_,Int<0>{}); + } + + CUTE_HOST_DEVICE constexpr static + auto + get_layoutD_MN() + { + // (thr_idx,val_idx) -> (M,N) + auto layoutD_TV = get_layoutD_TV(); + // (M,K) -> (thr_idx,val_idx) + auto layoutD_MK = right_inverse(layoutD_TV).with_shape(TiledShape_MN{}); + + // athrid = (v,m,k) -> thr_idx + auto thrID_D = make_layout(size<0>(TiledLayout_TV{})); + + return cute::make_tuple(layoutD_MK, thrID_D); + } + + template ::value)> + CUTE_HOST_DEVICE static + auto + get_slice(ThrIdx const& thr_idx) + { + return ThrCopy(thr_idx); + } + + template ::value)> + CUTE_HOST_DEVICE static + auto + get_thread_slice(ThrIdx const& thr_idx) + { + return get_slice(thr_idx); + } +}; + +template +struct ThrCopy +{ + ThrIdx thr_idx_; + + CUTE_HOST_DEVICE + ThrCopy(ThrIdx const& thr_idx) : thr_idx_(thr_idx) {} + + template + CUTE_HOST_DEVICE + auto + partition_S(STensor&& stensor) { + //static_assert(sizeof(typename remove_cvref_t::value_type) == sizeof(typename TiledCopy::ValType), + // "Expected ValType for tiling SrcTensor."); + auto thr_tensor = make_tensor(std::forward(stensor).data(), TiledCopy::tidfrg_S(stensor.layout())); + return thr_tensor(thr_idx_, _, repeat>(_)); + } + + template + CUTE_HOST_DEVICE + auto + partition_D(DTensor&& dtensor) { + //static_assert(sizeof(typename remove_cvref_t::value_type) == sizeof(typename TiledCopy::ValType), + // "Expected ValType for tiling DstTensor."); + auto thr_tensor = make_tensor(std::forward(dtensor).data(), TiledCopy::tidfrg_D(dtensor.layout())); + return thr_tensor(thr_idx_, _, repeat>(_)); + } + + template + CUTE_HOST_DEVICE static + auto + retile_S(STensor&& stensor) { + // static_assert(sizeof(typename remove_cvref_t::value_type) == sizeof(typename TiledCopy::ValType), + // "Expected ValType for tiling SrcTensor."); + return make_tensor(std::forward(stensor).data(), TiledCopy::retile(stensor.layout())); + } + + template + CUTE_HOST_DEVICE static + auto + retile_D(DTensor&& dtensor) { + // static_assert(sizeof(typename remove_cvref_t::value_type) == sizeof(typename TiledCopy::ValType), + // "Expected ValType for tiling DstTensor."); + return make_tensor(std::forward(dtensor).data(), TiledCopy::retile(dtensor.layout())); + } +}; + + +template +CUTE_HOST_DEVICE +auto +make_tiled_copy_impl(Copy_Atom const& atom, + LayoutCopy_TV const&, + Tiler const&) +{ + return TiledCopy, LayoutCopy_TV, Tiler>{atom}; +} + +// +// These tile the Copy_Atom as a whole +// + +template +CUTE_HOST_DEVICE +auto +make_tiled_copy_A(Copy_Atom const& copy_atom, + TiledMMA const& tiled_mma) +{ + using MNK = typename TiledMMA::TiledShape_MNK; + return make_tiled_copy_impl(copy_atom, tiled_mma.get_layoutA_TV(), make_shape(size<0>(MNK{}),size<2>(MNK{}))); +} + +template +CUTE_HOST_DEVICE +auto +make_tiled_copy_B(Copy_Atom const& copy_atom, + TiledMMA const& tiled_mma) +{ + using MNK = typename TiledMMA::TiledShape_MNK; + return make_tiled_copy_impl(copy_atom, tiled_mma.get_layoutB_TV(), make_shape(size<1>(MNK{}),size<2>(MNK{}))); +} + +template +CUTE_HOST_DEVICE +auto +make_tiled_copy_C(Copy_Atom const& copy_atom, + TiledMMA const& tiled_mma) +{ + using MNK = typename TiledMMA::TiledShape_MNK; + return make_tiled_copy_impl(copy_atom, tiled_mma.get_layoutC_TV(), make_shape(size<0>(MNK{}),size<1>(MNK{}))); +} + +// returns the smallest tiled copy that can retile LayoutC_TV +// for use with pipelined epilogues with subtiled stores +template +CUTE_HOST_DEVICE +auto +make_tiled_copy_C_atom(Copy_Atom const& copy_atom, + TiledMMA const& tiled_mma) +{ + // Truncate the V-layout to just the Copy_Atom, keep the V-order + auto layoutC_TV = tiled_mma.get_layoutC_TV(); + auto copy_V = Int::NumValSrc>{}; + CUTE_STATIC_ASSERT_V(copy_V <= size<1>(layoutC_TV)); + auto layout_TV = composition(layoutC_TV, make_layout(make_shape(size<0>(layoutC_TV), copy_V))); + + // Recompute tiler and restride the TV layout for the new tiler + + // Tiler -- Find the active elements in the MMA tensor and generate a tiler to extract them + // Convert to the awkward by-mode tiler to preserve the modes of the tiled MMA + using MNK = typename TiledMMA::TiledShape_MNK; + auto mma_tiler = make_shape(size<0>(MNK{}),size<1>(MNK{})); + auto mma_zeros = repeat_like(mma_tiler, Int<0>{}); + + auto tiler = transform(make_seq{}, [&](auto i) { + return filter(composition(make_layout(mma_tiler, replace(mma_zeros, Int<1>{})), layout_TV)); + }); + + // Layout_TV -- Find the (tid,vid) -> tile coord transformation + // Apply the tiler to a reference and transform the codomain + // tile_coord -> mma_coord + auto tile2mma = composition(make_layout(mma_tiler), tiler); + + // (tid,vid) -> tile_coord + auto layout_tv = composition(left_inverse(tile2mma), layout_TV); + + + using MNK = typename TiledMMA::TiledShape_MNK; + return make_tiled_copy_impl(copy_atom, layout_tv, tiler); +} + +template > +CUTE_HOST_DEVICE +auto +make_tiled_copy(Copy_Atom const& copy_atom, + ThrLayout const& thr_layout = {}, // (m,n) -> thr_idx + ValLayout const& val_layout = {}) +{ + constexpr int R = cute::max(rank_v, rank_v); + + auto thr_layout_mn = append(thr_layout, Layout<_1>{}); + auto val_layout_mn = append(val_layout, Layout<_1>{}); + + // Take the raked_products to compute the Layout_MN + auto layout_mn = raked_product(thr_layout_mn, val_layout_mn); + auto layout_tv = right_inverse(layout_mn).with_shape(make_shape(size(thr_layout), size(val_layout))); + // print("thr_layout: "); print(thr_layout_mn); print("\n"); + // print("val_layout: "); print(val_layout_mn); print("\n"); + // print("layout_mn : "); print(layout_mn); print("\n"); + // print("layout_tv : "); print(layout_tv); print("\n"); + + return make_tiled_copy_impl(copy_atom, layout_tv, product_each(shape(layout_mn))); +} + +// Make a TiledCopy out of the copy_atom that matches the Src-Layout of tiled_copy +template +CUTE_HOST_DEVICE +auto +make_tiled_copy_S(Copy_Atom const& copy_atom, + TiledCopy const& tiled_copy) +{ + return make_tiled_copy_impl(copy_atom, tiled_copy.get_layoutS_TV(), typename TiledCopy::Tiler_MN{}); +} + +// Make a TiledCopy out of the copy_atom that matches the Dst-Layout of tiled_copy +template +CUTE_HOST_DEVICE +auto +make_tiled_copy_D(Copy_Atom const& copy_atom, + TiledCopy const& tiled_copy) +{ + return make_tiled_copy_impl(copy_atom, tiled_copy.get_layoutD_TV(), typename TiledCopy::Tiler_MN{}); +} + +// +// Size +// + +// The logical size of a TileCopy +template +CUTE_HOST_DEVICE constexpr +auto +tile_size(TiledCopy const&) +{ + return size(typename TiledCopy::TiledShape_MN{}); +} + +// The number of threads involved in a TiledCopy +template +CUTE_HOST_DEVICE constexpr +auto +size(TiledCopy const&) +{ + return typename TiledCopy::TiledNumThr{}; +} + +// +// Display utilities +// + +template +CUTE_HOST_DEVICE +void +print(Copy_Atom, T> const&) +{ + using Atom = Copy_Atom, T>; + print("Copy_Atom\n"); + print(" ThrID: "); print(typename Atom::ThrID{}); print("\n"); + print(" ValLayoutSrc: "); print(typename Atom::ValLayoutSrc{}); print("\n"); + print(" ValLayoutDst: "); print(typename Atom::ValLayoutDst{}); print("\n"); + print(" ValLayoutRef: "); print(typename Atom::ValLayoutRef{}); print("\n"); + print(" ValueType: %db\n", int(sizeof_bits::value)); +} + +template +CUTE_HOST_DEVICE +void +print(TiledCopy const& copy, char const* pad = "") +{ + using Copy = TiledCopy; + print("TiledCopy\n"); + print(" Tiler_MN: "); print(typename Copy::Tiler_MN{}); print("\n"); + print(" TiledLayout_TV: "); print(typename Copy::TiledLayout_TV{}); print("\n"); + print(static_cast(copy)); +} + +template +CUTE_HOST_DEVICE +void +print(ThrCopy const&) +{ + print(TiledCopy{}); +} + +template +CUTE_HOST_DEVICE +auto +print_latex(TiledCopy const& copy) +{ + auto [layoutS_MN, thrID_S] = copy.get_layoutS_MN(); + auto [layoutD_MN, thrID_D] = copy.get_layoutD_MN(); + + print_latex_copy(layoutS_MN, thrID_S, + layoutD_MN, thrID_D); +} + +// MNK Copy Layout to Latex TIKZ -- 8-value color coded by thread +template +CUTE_HOST_DEVICE +void +print_latex_copy(LayoutS const& S, ThrIDS const& TS, // (m,n) -> (tid,vid) and tid -> thr_idx + LayoutD const& D, ThrIDD const& TD) // (m,n) -> (tid,vid) and tid -> thr_idx +{ + CUTE_STATIC_ASSERT_V(rank(S) == Int<2>{}); + CUTE_STATIC_ASSERT_V(rank(D) == Int<2>{}); + + assert(size<0>(S) == size<0>(D)); + assert(size<1>(S) == size<1>(D)); + + char const* latex_header = + "\\documentclass{standalone}\n" + "\\usepackage{tikz}\n" + "\\usetikzlibrary{external}\n" + "\\tikzexternalize\n" + "\\begin{document}\n" + "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},box/.style={rectangle,draw=black,thick,minimum size=1cm,anchor=center}]\n\n"; + char const* latex_footer = + "\\end{tikzpicture}\n" + "\\end{document}\n"; + + char const* color_map[8] = {"{rgb,255:red,175;green,175;blue,255}", + "{rgb,255:red,175;green,255;blue,175}", + "{rgb,255:red,255;green,255;blue,175}", + "{rgb,255:red,255;green,175;blue,175}", + "{rgb,255:red,210;green,210;blue,255}", + "{rgb,255:red,210;green,255;blue,210}", + "{rgb,255:red,255;green,255;blue,210}", + "{rgb,255:red,255;green,210;blue,210}",}; + + // Header + printf("%% LayoutS: "); print(S); printf("\n"); + printf("%% ThrIDS : "); print(TS); printf("\n"); + printf("%% LayoutD: "); print(D); printf("\n"); + printf("%% ThrIDD : "); print(TD); printf("\n\n"); + + printf(latex_header); + + // S starting at 0,0 + for (int i = 0; i < size<0>(S); ++i) { + for (int j = 0; j < size<1>(S); ++j) { + int thrid = S(i,j) % size(TS); + int val_idx = S(i,j) / size(TS); + int thr_idx = TS(thrid); + + printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n", + color_map[thr_idx % 8], + i, j, + thr_idx, val_idx); + } + } + + // D starting at 0,size<1>(S)+3 + for (int i = 0; i < size<0>(D); ++i) { + for (int j = 0; j < size<1>(D); ++j) { + int thrid = D(i,j) % size(TD); + int val_idx = D(i,j) / size(TD); + int thr_idx = TD(thrid); + + printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n", + color_map[thr_idx % 8], + i, j + size<1>(S) + 3, + thr_idx, val_idx); + } + } + + // S Labels + for (int i = 0, j = -1; i < size<0>(S); ++i) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, i); + } + for (int j = 0, i = -1; j < size<1>(S); ++j) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, j); + } + // D Labels + for (int i = 0, j = size<1>(D); i < size<0>(S); ++i) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j + size<1>(S) + 3, i); + } + for (int j = 0, i = -1; j < size<1>(D); ++j) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j + size<1>(S) + 3, j); + } + + // Footer + printf(latex_footer); +} + +} // end namespace cute + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#include +#include +// #include +// Config +// #if (__CUDACC_VER_MAJOR__ >= 12) +#if 0 +# define CUTE_COPY_ATOM_TMA_SM90_ENABLED +#endif + +#if defined(CUTE_COPY_ATOM_TMA_SM90_ENABLED) +#include +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/cute/atom/copy_traits.hpp b/csrc/mctlass/include/cute/atom/copy_traits.hpp new file mode 100644 index 0000000..cea03c0 --- /dev/null +++ b/csrc/mctlass/include/cute/atom/copy_traits.hpp @@ -0,0 +1,131 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +namespace cute +{ + +/** + * concept Copy_Traits + * { + * using ThrID = // Logical thread id (tid) -> tidx + * + * using SrcLayout = // (Logical src thread id (tid), Logical src value id (vid)) -> bit + * using DstLayout = // (Logical dst thread id (tid), Logical dst value id (vid)) -> bit + * using RefLayout = // (Logical ref thread id (tid), Logical ref value id (vid)) -> bit + * }; + * + * The abstract bit ordering of the Copy_Traits (the codomain of SrcLayout, DstLayout, and RefLayout) + * is arbitrary and only used to construct maps + * (ref-tid,ref-vid) -> (src-tid,src-vid) + * (ref-tid,ref-vid) -> (dst-tid,dst-vid) + * in TiledCopy. The Layout_TV in TiledCopy is in accordance with the RefLayout of a Traits, then mapped to + * the Src or Dst (tid,vid) representation on demand. + * + */ + +template +struct Copy_Traits +{ + static_assert(sizeof(CopyOperation) == 0, "Copy_Traits not implemented for this Copy_Operation."); +}; + +template +struct Copy_Traits> +{ + // Logical thread id to thread idx (one-thread) + using ThrID = Layout<_1>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout::value>>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout::value>>>; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; +}; + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (one-thread) + using ThrID = Layout<_1>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout, Stride<_0,_0>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout, Stride<_0,_0>>; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; +}; + +// +// Generic copy_unpack for any Copy_Traits +// +template +CUTE_HOST_DEVICE constexpr +void +copy_unpack(Copy_Traits const&, + Tensor const& src, + Tensor & dst) +{ + // Specializations can generalize on these checks + //static_assert(is_smem::value, "Expected smem for this Copy_Traits"); + //static_assert(is_rmem::value, "Expected rmem for this Copy_Traits"); + + using RegistersSrc = typename Operation::SRegisters; + using RegistersDst = typename Operation::DRegisters; + using RegTypeSrc = typename remove_extent::type; + using RegTypeDst = typename remove_extent::type; + constexpr int RegNumSrc = extent::value; + constexpr int RegNumDst = extent::value; + + Tensor rS = recast(src); + Tensor rD = recast(dst); + + CUTE_STATIC_ASSERT_V(size(rS) == Int{}, + "In CopyAtom, src layout doesn't vectorize into registers. This src layout is incompatible with this tiled copy."); + CUTE_STATIC_ASSERT_V(size(rD) == Int{}, + "In CopyAtom, dst layout doesn't vectorize into registers. This dst layout is incompatible with this tiled copy."); + + detail::explode(Operation::copy, + rS, make_int_sequence{}, + rD, make_int_sequence{}); +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/atom/copy_traits_sm75.hpp b/csrc/mctlass/include/cute/atom/copy_traits_sm75.hpp new file mode 100644 index 0000000..6130396 --- /dev/null +++ b/csrc/mctlass/include/cute/atom/copy_traits_sm75.hpp @@ -0,0 +1,160 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +#include + +#include + +namespace cute +{ + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (warp) + using ThrID = Layout<_32>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout,_128>, + Stride, _1>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout, + Stride<_32, _1>>; + + // Reference map from (thr,val) to bit + using RefLayout = DstLayout; +}; + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (warp) + using ThrID = Layout<_32>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout,_128>, + Stride, _1>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout>, + Stride<_32,Stride< _1,_1024>>>; + + // Reference map from (thr,val) to bit + using RefLayout = DstLayout; +}; + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (warp) + using ThrID = Layout<_32>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout, + Stride<_128, _1>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout>, + Stride<_32,Stride< _1,_1024>>>; + + // Reference map from (thr,val) to bit + using RefLayout = DstLayout; +}; + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (warp) + using ThrID = Layout<_32>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout, + Stride<_128, _1>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout>, + Stride<_32,Stride< _1,_1024>>>; + + // Reference map from (thr,val) to bit + using RefLayout = DstLayout; +}; + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (warp) + using ThrID = Layout<_32>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout,_128>, + Stride, _1>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout,Shape <_16, _2>>, + Stride,Stride< _1,_128>>>; + + // Reference map from (thr,val) to bit + using RefLayout = DstLayout; +}; + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (warp) + using ThrID = Layout<_32>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout,_128>, + Stride, _1>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout,Shape <_16, _2, _2>>, + Stride,Stride< _1,_128,_1024>>>; + + // Reference map from (thr,val) to bit + using RefLayout = DstLayout; +}; + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (warp) + using ThrID = Layout<_32>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout, + Stride<_128, _1>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout,Shape <_16, _2, _4>>, + Stride,Stride< _1,_128,_1024>>>; + + // Reference map from (thr,val) to bit + using RefLayout = DstLayout; +}; + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/atom/copy_traits_sm80.hpp b/csrc/mctlass/include/cute/atom/copy_traits_sm80.hpp new file mode 100644 index 0000000..d42a4b7 --- /dev/null +++ b/csrc/mctlass/include/cute/atom/copy_traits_sm80.hpp @@ -0,0 +1,130 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +#include + +#include + +namespace cute +{ + +template +struct Copy_Traits> +{ + // Logical thread id to thread idx (one-thread) + using ThrID = Layout<_1>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout::value>>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout::value>>>; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; +}; + +template +struct Copy_Traits> +{ + // Logical thread id to thread idx (one-thread) + using ThrID = Layout<_1>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout::value>>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout::value>>>; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; +}; + +template +struct Copy_Traits> +{ + // Logical thread id to thread idx (one-thread) + using ThrID = Layout<_1>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout::value>>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout::value>>>; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (16 threads) + using ThrID = Layout<_16>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout, _64>, + Stride, _1>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout>, + Stride<_16, Stride<_1, _256>>>; + + // Reference map from (thr,val) to bit + using RefLayout = DstLayout; +}; + +// Element copy selector +template +CUTE_HOST_DEVICE constexpr +auto +select_elementwise_copy(SrcTensor const&, DstTensor const&) +{ + using SrcType = typename SrcTensor::value_type; + using DstType = typename DstTensor::value_type; + +#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED) + if constexpr (is_gmem::value && is_smem::value && + sizeof(SrcType) == sizeof(DstType) && + (sizeof(SrcType) == 4 || sizeof(SrcType) == 8 || sizeof(SrcType) == 16)) + { + return SM80_CP_ASYNC_CACHEALWAYS{}; + } else { + return UniversalCopy{}; + } + + CUTE_GCC_UNREACHABLE; +#else + return UniversalCopy{}; +#endif +} + +} diff --git a/csrc/mctlass/include/cute/atom/copy_traits_sm90.hpp b/csrc/mctlass/include/cute/atom/copy_traits_sm90.hpp new file mode 100644 index 0000000..8c5e843 --- /dev/null +++ b/csrc/mctlass/include/cute/atom/copy_traits_sm90.hpp @@ -0,0 +1,132 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +#include +#include + +#include + +namespace cute +{ + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (warp) + using ThrID = Layout<_32>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = typename Copy_Traits::DstLayout; + // Map from (dst-thr,dst-val) to bit + using DstLayout = typename Copy_Traits::SrcLayout; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; +}; + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (warp) + using ThrID = Layout<_32>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = typename Copy_Traits::DstLayout; + // Map from (dst-thr,dst-val) to bit + using DstLayout = typename Copy_Traits::SrcLayout; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; +}; + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (warp) + using ThrID = Layout<_32>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = typename Copy_Traits::DstLayout; + // Map from (dst-thr,dst-val) to bit + using DstLayout = typename Copy_Traits::SrcLayout; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; +}; + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (warp) + using ThrID = Layout<_32>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = typename Copy_Traits::DstLayout; + // Map from (dst-thr,dst-val) to bit + using DstLayout = typename Copy_Traits::SrcLayout; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; +}; + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (warp) + using ThrID = Layout<_32>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = typename Copy_Traits::DstLayout; + // Map from (dst-thr,dst-val) to bit + using DstLayout = typename Copy_Traits::SrcLayout; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; +}; + +template <> +struct Copy_Traits +{ + // Logical thread id to thread idx (warp) + using ThrID = Layout<_32>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = typename Copy_Traits::DstLayout; + // Map from (dst-thr,dst-val) to bit + using DstLayout = typename Copy_Traits::SrcLayout; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; +}; + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/atom/copy_traits_sm90_tma.hpp b/csrc/mctlass/include/cute/atom/copy_traits_sm90_tma.hpp new file mode 100644 index 0000000..ae33b24 --- /dev/null +++ b/csrc/mctlass/include/cute/atom/copy_traits_sm90_tma.hpp @@ -0,0 +1,973 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#if !defined(__MACACC_RTC__) +#include +#endif + +#include +#include + +#include + +#include +#include + +namespace cute +{ + +////////////////////////////////////////////////////////////////////////////// +///////////////////////////// TMA_LOAD /////////////////////////////////////// +////////////////////////////////////////////////////////////////////////////// + +struct SM90_TMA_LOAD_OP : SM90_TMA_LOAD {}; + +// The executable SM90_TMA_LOAD with tma_desc and tma_mbar +template +struct Copy_Traits +{ + using ThrID = Layout<_1>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout>; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; + + // SM90_TMA_LOAD arguments + TmaDescriptor const& tma_desc_; + uint64_t& tma_load_mbar_; + + template + CUTE_HOST_DEVICE constexpr + void + copy_unpack_(void const* const dst_ptr, + Coord const& src_coord, seq) const + { +#if 0 + print("THR (%d,%d,%d) BLK (%d,%d,%d)\n", + threadIdx.x, threadIdx.y, threadIdx.z, + blockIdx.x, blockIdx.y, blockIdx.z); + print(" TMA Coord "); print(src_coord); print("\n"); + print(" TMA Shape "); print(make_tuple(uint64_t(tma_desc_.size0_), + uint64_t(tma_desc_.size1_), + uint64_t(tma_desc_.size2_), + uint64_t(tma_desc_.size3_))); print("\n"); +#endif + + SM90_TMA_LOAD::copy(&tma_desc_, + tma_load_mbar_, + dst_ptr, + get(src_coord)...); + } + + // This is the copy_unpack dispatch for this Copy_Traits + // Src needs to be a gmem tensor with TmaCoordIterator .data() + // Dst needs to be a smem tensor + template + CUTE_HOST_DEVICE friend constexpr + void + copy_unpack(Copy_Traits const& traits, + Tensor const& src, + Tensor & dst) + { + //static_assert(is_gmem::value, "Expected gmem src for SM90_TMA_LOAD"); // TMA spoofed src tensor + static_assert(is_smem::value, "Expected smem dst for SM90_TMA_LOAD"); + + traits.copy_unpack_(dst.data().get(), src.data().coord_, tuple_seq{}); + } +}; + +// The non-executable SM90_TMA_LOAD with tma_desc and no tma_mbar +// Use .with(tma_mbar) to construct an executable version +template +struct Copy_Traits +{ + using ThrID = Layout<_1>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout>; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; + + // SM90_TMA_LOAD arguments + TmaDescriptor tma_desc_; + GmemStrides g_stride_; + + // Return TmaDescriptor/TensorMap + CUTE_HOST_DEVICE constexpr + TmaDescriptor const* + get_tma_descriptor() const { + return &tma_desc_; + } + + // Construct an executable SM90_TMA_LOAD with tma_mbar + CUTE_HOST_DEVICE constexpr + Copy_Traits + with(uint64_t& tma_mbar, uint16_t const& multicast_mask = 0) const { + // We accept multicast_mask here to keep the API for both atoms consistent + // assert(multicast_mask == 0); + (void) multicast_mask; + return {tma_desc_, tma_mbar}; + } + + // Generate the TMA coord tensor + template + CUTE_HOST_DEVICE constexpr + auto + get_tma_tensor(GShape const& g_shape) const { + static_assert(is_congruent::value); + constexpr int tma_rank = decltype(cute::min(rank(flatten(g_stride_)), Int<5>{}))::value; + return make_tensor(ArithmeticTupleIterator(as_arithmetic_tuple(repeat(Int<0>{}))), + g_shape, + g_stride_); + } + + // Don't try to execute a copy with SM90_TMA_LOAD before calling .with() + template + CUTE_HOST_DEVICE friend constexpr void + copy_unpack(Copy_Traits const& traits, + Tensor const& src, + Tensor & dst) = delete; +}; + +////////////////////////////////////////////////////////////////////////////// +///////////////////////////// TMA_LOAD_MULTICAST ///////////////////////////// +////////////////////////////////////////////////////////////////////////////// + +struct SM90_TMA_LOAD_MULTICAST_OP : SM90_TMA_LOAD_MULTICAST {}; + +template +struct Copy_Traits +{ + using ThrID = Layout<_1>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout>; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; + + // SM90_TMA_LOAD_MULTICAST arguments + TmaDescriptor const& tma_desc_; + uint64_t& tma_load_mbar_; + uint16_t const& multicast_mask_; + + template + CUTE_HOST_DEVICE constexpr + void + copy_unpack_(void const* const dst_ptr, + Coord const& src_coord, seq) const + { +#if 0 + print("THR (%d,%d,%d) BLK (%d,%d,%d)\n", + threadIdx.x, threadIdx.y, threadIdx.z, + blockIdx.x, blockIdx.y, blockIdx.z); + print(" TMA Coord "); print(src_coord); print("\n"); + print(" TMA Shape "); print(make_tuple(uint64_t(tma_desc_.size0_), + uint64_t(tma_desc_.size1_), + uint64_t(tma_desc_.size2_), + uint64_t(tma_desc_.size3_))); print("\n"); +#endif + + SM90_TMA_LOAD_MULTICAST::copy(&tma_desc_, + tma_load_mbar_, + multicast_mask_, + dst_ptr, + get(src_coord)...); + } + + template + CUTE_HOST_DEVICE friend constexpr + void + copy_unpack(Copy_Traits const& traits, + Tensor const& src, + Tensor & dst) + { + //static_assert(is_gmem::value, "Expected gmem src for SM90_TMA_LOAD"); // TMA spoofed src tensor + static_assert(is_smem::value, "Expected smem dst for SM90_TMA_LOAD_MULTICAST"); + + traits.copy_unpack_(dst.data().get(), src.data().coord_, tuple_seq{}); + } +}; + +template +struct Copy_Traits +{ + using ThrID = Layout<_1>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout>; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; + + // SM90_TMA_LOAD_MULTICAST arguments + TmaDescriptor tma_desc_; + GmemStrides g_stride_; + + // Return TmaDescriptor/TensorMap + CUTE_HOST_DEVICE constexpr + TmaDescriptor const* + get_tma_descriptor() const { + return &tma_desc_; + } + + // Construct an executable SM90_TMA_LOAD_MULTICAST with tma_mbar + CUTE_HOST_DEVICE constexpr + Copy_Traits + with(uint64_t& tma_load_mbar, uint16_t const& multicast_mask) const { + return {tma_desc_, tma_load_mbar, multicast_mask}; + } + + // Generate the TMA coord tensor + template + CUTE_HOST_DEVICE constexpr + auto + get_tma_tensor(GShape const& g_shape) const { + static_assert(is_congruent::value); + constexpr int tma_rank = decltype(cute::min(rank(flatten(g_stride_)), Int<5>{}))::value; + return make_tensor(ArithmeticTupleIterator(as_arithmetic_tuple(repeat(Int<0>{}))), + g_shape, + g_stride_); + } + + // Don't try to execute a copy with SM90_TMA_LOAD_MULTICAST before calling .with() + template + CUTE_HOST_DEVICE friend constexpr void + copy_unpack(Copy_Traits const& traits, + Tensor const& src, + Tensor & dst) = delete; +}; + +////////////////////////////////////////////////////////////////////////////// +///////////////////////////// TMA_STORE ////////////////////////////////////// +////////////////////////////////////////////////////////////////////////////// + +// The executable SM90_TMA_STORE with tma_desc +template +struct Copy_Traits +{ + using ThrID = Layout<_1>; + + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout>; + + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; + + // SM90_TMA_STORE arguments + TmaDescriptor tma_desc_; + GmemStrides g_stride_; + + // Return TmaDescriptor/TensorMap + CUTE_HOST_DEVICE constexpr + TmaDescriptor const* + get_tma_descriptor() const { + return &tma_desc_; + } + + // Generate the TMA coord tensor + template + CUTE_HOST_DEVICE constexpr + auto + get_tma_tensor(GShape const& g_shape) const { + static_assert(is_congruent::value); + constexpr int tma_rank = decltype(cute::min(rank(flatten(g_stride_)), Int<5>{}))::value; + return make_tensor(ArithmeticTupleIterator(as_arithmetic_tuple(repeat(Int<0>{}))), + g_shape, + g_stride_); + } + + template + CUTE_HOST_DEVICE constexpr + void + copy_unpack_(void const* const src_ptr, + Coord const& dst_coord, seq) const + { +#if 0 + print("THR (%d,%d,%d) BLK (%d,%d,%d)\n", + threadIdx.x, threadIdx.y, threadIdx.z, + blockIdx.x, blockIdx.y, blockIdx.z); + print(" TMA Coord "); print(dst_coord); print("\n"); + print(" TMA Shape "); print(make_tuple(uint64_t(tma_desc_.size0_), + uint64_t(tma_desc_.size1_), + uint64_t(tma_desc_.size2_), + uint64_t(tma_desc_.size3_))); print("\n"); +#endif + + SM90_TMA_STORE::copy(&tma_desc_, + src_ptr, + get(dst_coord)...); + } + + // This is the copy_unpack dispatch for this Copy_Traits + // Src needs to be a smem tensor + // Dst needs to be a gmem tensor with TmaCoordIterator .data() + template + CUTE_HOST_DEVICE friend constexpr + void + copy_unpack(Copy_Traits const& traits, + Tensor const& src, + Tensor & dst) + { + static_assert(is_smem::value, "Expected smem src for SM90_TMA_STORE"); + //static_assert(is_gmem::value, "Expected gmem dst for SM90_TMA_STORE"); // TMA spoofed src tensor + + traits.copy_unpack_(src.data().get(), dst.data().coord_, tuple_seq{}); + } +}; + +////////////////////////////////////////////////////////////////////////////// +///////////////////////////// BULK COPY ////////////////////////////////////// +////////////////////////////////////////////////////////////////////////////// + +template +struct Copy_Traits +{ + static_assert(int32_t(NumBits::value / 8) % 16 == 0, + "Bulk Copy requires copy vector size align to 16B."); + + using ThrID = Layout<_1>; + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout>; + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; + + // SM90_BULK_COPY_G2S arguments + // 0: uint64_t* bulk_load_memory_barrier + cute::tuple bulk_load_mbar_; + + template + CUTE_HOST_DEVICE friend constexpr + void + copy_unpack(Copy_Traits const& traits, + Tensor const& src, + Tensor & dst) + { + static_assert(is_same, cute::tuple>::value, + "Extra arguments not set. Set .with() before use."); + static_assert(is_gmem::value, "Expected gmem src for SM90_BULK_COPY_G2S"); + static_assert(is_smem::value, "Expected smem dst for SM90_BULK_COPY_G2S"); + SM90_BULK_COPY_G2S::copy(src.data().get(), *get<0>(traits.bulk_load_mbar_), + dst.data().get(), int32_t(NumBits::value / 8)); + } + + // Record the memory barrier for the instruction + CUTE_HOST_DEVICE constexpr + Copy_Traits + with(uint64_t& bulk_mbar) const { + return {{&bulk_mbar}}; + } +}; + +template +struct Copy_Traits +{ + static_assert(int32_t(NumBits::value / 8) % 16 == 0, + "Bulk Copy requires copy vector size align to 16B."); + + using ThrID = Layout<_1>; + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout>; + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; + + template + CUTE_HOST_DEVICE friend constexpr + void + copy_unpack(Copy_Traits const& traits, + Tensor const& src, + Tensor & dst) + { + static_assert(is_smem::value, "Expected smem src for SM90_BULK_COPY_S2G"); + static_assert(is_gmem::value, "Expected gmem dst for SM90_BULK_COPY_S2G"); + SM90_BULK_COPY_S2G::copy(src.data().get(), dst.data().get(), int32_t(NumBits::value / 8)); + } +}; + +// +// Placeholder for the bulk copy algorithm's default, auto-vectorizing behavior +// + +template +struct Copy_Traits +{ + // Logical thread id to thread idx (one-thread) + using ThrID = Layout<_1>; + // Map from (src-thr,src-val) to bit + using SrcLayout = Layout, Stride<_0,_0>>; + // Map from (dst-thr,dst-val) to bit + using DstLayout = Layout, Stride<_0,_0>>; + // Reference map from (thr,val) to bit + using RefLayout = SrcLayout; + + // SM90_UBULK_COPY arguments + // 0: uint64_t* bulk_load_memory_barrier [if this is a BULK_LOAD_G2S] + cute::tuple opargs_; + + // Record the memory barrier for the instruction + CUTE_HOST_DEVICE constexpr + Copy_Traits + with(uint64_t& bulk_mbar) const { + return {{&bulk_mbar}}; + } +}; + +// +// MAKE_TMA_COPY and related +// + +namespace detail +{ + +template +auto +get_swizzle_portion(ComposedLayout,Offset,SLayout>) +{ + return Swizzle{}; +} + +template +auto +get_swizzle_portion(Layout) +{ + return Swizzle<0,4,3>{}; +} + +template +auto +get_nonswizzle_portion(ComposedLayout,Offset,SLayout> const& slayout) +{ + return slayout.layout_fn(); +} + +template +auto +get_nonswizzle_portion(Layout const& slayout) +{ + return slayout; +} + +template +TMA::SmemSwizzleBits +get_tma_swizzle_bits(Swizzle) +{ + if constexpr (M == 4) { + switch (B) { + default: static_assert(0 <= B && B <= 3, "Expected B = 0,1,2, or 3 when M == 4. Unsupported layout swizzle."); + case 3: return TMA::SmemSwizzleBits::B128; + case 2: return TMA::SmemSwizzleBits::B64; + case 1: return TMA::SmemSwizzleBits::B32; + case 0: return TMA::SmemSwizzleBits::DISABLE; + } + } else + { + static_assert(M < 0, "Unsupported layout swizzle."); + } +} + +template +TMA::SmemSwizzleBits +get_tma_swizzle_bits(Layout const& layout) +{ + return get_tma_swizzle_bits(get_swizzle_portion(layout)); +} + +#if !defined(__MACACC_RTC__) +// Use a smem2gmode map to read through the GMEM tensor +// and construct a TMA Descriptor for the resulting instruction +template +CUTE_HOST +auto +make_tma_copy_desc(Tensor const& gtensor, // The original GMEM Tensor + Layout const& smem_inv, // smem_idx to flat gmode + Swizzle const& swizzle) // Swizzle fn on smem_idx +{ + using T = typename GEngine::value_type; + + auto flat_glayout = flatten(gtensor.layout()); + CUTE_STATIC_ASSERT_V(rank(flat_glayout) == rank(smem_inv)); + constexpr int rank_smem_inv = decltype(rank(smem_inv))::value; + + auto tma_multimode = rank(flat_glayout) > Int<5>{}; + constexpr uint32_t tma_dim = cute::min(rank(flat_glayout), 5);; + + // + // TMA gmem desc info + // + + void* gmem_address = (void*) gtensor.data(); + + cute::array gmem_prob_shape = {1,1,1,1,1}; + cute::array gmem_prob_stride = {0,0,0,0,0}; + for_each(make_seq{}, [&](auto i) { + auto e = stride(smem_inv); // For g++-7.5, let it deduce e rather than fuse with below + constexpr int j = decltype(e.mode())::value; + constexpr int tma_i = i < 5 ? i : 4; + + // Problem stride + uint64_t stride_j = stride(flat_glayout) * sizeof(T); + uint64_t old_stride = gmem_prob_stride[tma_i]; + gmem_prob_stride[tma_i] = gcd(gmem_prob_stride[tma_i], stride_j); + + // Problem shape + uint64_t shape_j = shape(flat_glayout); + if (gmem_prob_stride[tma_i] != 0) { + // We're "resetting" this TMA mode and using it as a "multimode" + // Recurrence: g_shape = (s_i - 1) * (d_i / gcd_j d_j) + 1 + gmem_prob_shape[tma_i] = (gmem_prob_shape[tma_i]-1) * (old_stride / gmem_prob_stride[tma_i]) + + (shape_j-1) * (stride_j / gmem_prob_stride[tma_i]) + + 1; + } else { + gmem_prob_shape[tma_i] = shape_j; + } + }); + + assert((reinterpret_cast(gmem_address) & 0b1111) == 0); // Address must be 16B-aligned + + assert(gmem_prob_shape[0] >= (uint64_t(1))); // Size must be min 1 + assert(gmem_prob_shape[0] <= (uint64_t(1) << 32)); // Size must be max 2^32 + assert(gmem_prob_shape[1] >= (uint64_t(1))); // Size must be min 1 + assert(gmem_prob_shape[1] <= (uint64_t(1) << 32)); // Size must be max 2^32 + assert(gmem_prob_shape[2] >= (uint64_t(1))); // Size must be min 1 + assert(gmem_prob_shape[2] <= (uint64_t(1) << 32)); // Size must be max 2^32 + assert(gmem_prob_shape[3] >= (uint64_t(1))); // Size must be min 1 + assert(gmem_prob_shape[3] <= (uint64_t(1) << 32)); // Size must be max 2^32 + assert(gmem_prob_shape[4] >= (uint64_t(1))); // Size must be min 1 + assert(gmem_prob_shape[4] <= (uint64_t(1) << 32)); // Size must be max 2^32 + + assert((gmem_prob_stride[0]) == sizeof(T)); // First stride is implicitly 1 + assert((gmem_prob_stride[1]) < (uint64_t(1) << 40)); // Stride must be max 2^40 + assert((gmem_prob_stride[1] & 0b1111) == 0); // Stride must be multiple of 16B (128b) + assert((gmem_prob_stride[2]) < (uint64_t(1) << 40)); // Stride must be max 2^40 + assert((gmem_prob_stride[2] & 0b1111) == 0); // Stride must be multiple of 16B (128b) + assert((gmem_prob_stride[3]) < (uint64_t(1) << 40)); // Stride must be max 2^40 + assert((gmem_prob_stride[3] & 0b1111) == 0); // Stride must be multiple of 16B (128b) + assert((gmem_prob_stride[4]) < (uint64_t(1) << 40)); // Stride must be max 2^40 + assert((gmem_prob_stride[4] & 0b1111) == 0); // Stride must be multiple of 16B (128b) + + // + // TMA smem desc info + // + + cute::array smem_box_shape = {1,1,1,1,1}; + cute::array smem_box_stride = {1,1,1,1,1}; + for_each(make_seq{}, [&](auto i) { + uint32_t shape_i = shape(smem_inv); + constexpr int tma_i = i < 5 ? i : 4; + if (tma_multimode && tma_i == 4) { + // We're "reusing" this TMA mode and using it as a "multimode" + smem_box_shape[tma_i] = 1; + } else { + smem_box_shape[tma_i] = shape_i; + } + }); + + assert(smem_box_shape[0] >= (uint64_t(1))); // Size must be min 1 + assert(smem_box_shape[0] <= (uint64_t(1) << 8)); // Size must be max 2^8 + assert(smem_box_shape[0] >= (uint64_t(1))); // Size must be min 1 + assert(smem_box_shape[0] <= (uint64_t(1) << 8)); // Size must be max 2^8 + assert(smem_box_shape[0] >= (uint64_t(1))); // Size must be min 1 + assert(smem_box_shape[0] <= (uint64_t(1) << 8)); // Size must be max 2^8 + assert(smem_box_shape[0] >= (uint64_t(1))); // Size must be min 1 + assert(smem_box_shape[0] <= (uint64_t(1) << 8)); // Size must be max 2^8 + + assert(smem_box_stride[0] >= (uint32_t(1))); // Stride must be min 1 + assert(smem_box_stride[0] <= (uint32_t(8))); // Stride must be max 2^3 + assert(smem_box_stride[1] >= (uint32_t(1))); // Stride must be min 1 + assert(smem_box_stride[1] <= (uint32_t(8))); // Stride must be max 2^3 + assert(smem_box_stride[2] >= (uint32_t(1))); // Stride must be min 1 + assert(smem_box_stride[2] <= (uint32_t(8))); // Stride must be max 2^3 + assert(smem_box_stride[3] >= (uint32_t(1))); // Stride must be min 1 + assert(smem_box_stride[3] <= (uint32_t(8))); // Stride must be max 2^3 + assert(smem_box_stride[4] >= (uint32_t(1))); // Stride must be min 1 + assert(smem_box_stride[4] <= (uint32_t(8))); // Stride must be max 2^3 + + // + // Construct the descriptor + // + + TmaDescriptor tma_desc = {0}; + + // + // TMA general info + // + +// #if (__CUDACC_VER_MAJOR__ >= 12) +#if 0 + CUtensorMapDataType tma_format = TMA::to_CUtensorMapDataType(); + CUtensorMapInterleave tma_interleave = CU_TENSOR_MAP_INTERLEAVE_NONE; + CUtensorMapL2promotion tma_l2Promotion = CU_TENSOR_MAP_L2_PROMOTION_NONE; + CUtensorMapFloatOOBfill tma_oobFill = CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE; + + // TMA smem swizzle type + CUtensorMapSwizzle smem_swizzle = TMA::to_CUtensorMapSwizzle(get_tma_swizzle_bits(swizzle)); + CUresult result = cuTensorMapEncodeTiled( + &tma_desc, + tma_format, + tma_dim, + gmem_address, + gmem_prob_shape.data(), + gmem_prob_stride.data() + 1, // gmem_prob_stride[0] implicitly 1 + smem_box_shape.data(), + smem_box_stride.data(), + tma_interleave, + smem_swizzle, + tma_l2Promotion, + tma_oobFill); + + if (result != CUDA_SUCCESS) { + std::cerr << "TMA Desc Addr: " << &tma_desc + << "\nformat " << tma_format + << "\ndim " << tma_dim + << "\ngmem_address " << gmem_address + << "\nglobalDim " << gmem_prob_shape + << "\nglobalStrides " << gmem_prob_stride + << "\nboxDim " << smem_box_shape + << "\nelementStrides " << smem_box_stride + << "\ninterleave " << tma_interleave + << "\nswizzle " << smem_swizzle + << "\nl2Promotion " << tma_l2Promotion + << "\noobFill " << tma_oobFill << std::endl; + std::cerr << "Error: Failed to initialize the TMA descriptor " << result << std::endl; + assert(false); + } + +#endif // (__CUDACC_VER_MAJOR__ >= 12) + // Finally, get the inverse permutation of the E bases for the mocked gmem stride + auto gmem_stride_bases_flat = transform(make_seq{}, [&](auto i) { + auto k = find(stride(smem_inv), E{}); + // For gcc 7.5 -- avoid 'if constexpr' + int32_t tma_coord_stride = int32_t(stride(flat_glayout) * sizeof(T) / (gmem_prob_stride[4] != 0 ? gmem_prob_stride[4] : 16)); + return conditional_return(tma_multimode && (k >= Int<4>{}), + E<4>{} * tma_coord_stride, // The 4th TMA mode is the multimode, use int32_t coord stride + E{}); + }); + + // Give that the profile of gtensor and fold it + // NOTE: This is the only reason we want the original gtensor shape rather than the more intuitive flattened shape + auto gmem_stride_bases = stride(composition(make_layout(repeat_like(shape(flat_glayout), Int<2>{}), gmem_stride_bases_flat), + make_layout(repeat_like(shape(gtensor), Int<2>{})))); + + return make_tuple(tma_desc, gmem_stride_bases); +} + +template +CUTE_HOST +auto +make_tma_copy_tiled(CopyOp, + Tensor const& gtensor, // Full GMEM Tensor + SLayout const& slayout, // CTA Tile of SMEM + Layout const& cta_t_map, // T: CTA thr idx -> logical TMA tid + Layout const& cta_v_map) // V: CTA val idx -> gmem coord +{ + // + // TMA parameter checking + // + + CUTE_STATIC_ASSERT_V(product_each(shape(slayout)) == product_each(shape(cta_v_map)), + "TMA requires CTA_Tile and SLayout top-level shape equivalence."); + CUTE_STATIC_ASSERT_V(size(slayout) % cosize(cta_t_map) == Int<0>{}, + "Number of active CTAs in TMA must divide domain size of slayout."); + + // + // TMA slayout manipulation + // + + auto flat_glayout = flatten(gtensor.layout()); + + // Invert the smem to get the largest contiguous vector in the smem layout + auto inv_smem_layout = right_inverse(get_nonswizzle_portion(slayout)); + // trunc_smem_idx -> trunc_smem_coord + + // Map from smem idx to a gmem mode + auto sidx_to_gmode = coalesce(composition(cta_v_map, inv_smem_layout)); + + // Truncate any incompatibilities + auto smem_rank = find_if(stride(sidx_to_gmode), [](auto e) { + auto v = basis_value(e); + return not is_constant<1,decltype(v)>{}; + }); + static_assert(smem_rank > 0, "Could not find a common smem-gmem vectorization for TMA. Do they have a common majorness?"); + // TMA uses a maximum of 5 modes + // If the gtensor has more than 5 modes, we need to reserve the last TMA-mode as a "multimode" + constexpr int smem_tma_rank = cute::min(int(smem_rank), (rank(flat_glayout) > Int<5>{} ? 4 : 5)); + + // Keep only the static-1 basis modes into gmem + auto sidx_to_gmode_trunc = take<0,smem_tma_rank>(sidx_to_gmode); + + // Split according to the portion each multicast CTA will be responsible for + auto sidx_to_gmode_vt = logical_divide(sidx_to_gmode_trunc, shape_div(size(sidx_to_gmode_trunc), cosize(cta_t_map))); + +#if 0 + print("g_layout : "); print(gtensor.layout()); print("\n"); + print("s_layout : "); print(slayout); print("\n"); + print("cta_t_map : "); print(cta_t_map); print("\n"); + print("cta_v_map : "); print(cta_v_map); print("\n"); + print("inv_smem : "); print(inv_smem_layout); print("\n"); + print("sidx_to_gmode : "); print(sidx_to_gmode); print("\n"); + + print("sidx_to_gmode_trunc : "); print(sidx_to_gmode_trunc); print("\n"); + print("sidx_to_gmode_vt : "); print(sidx_to_gmode_vt); print("\n"); +#endif + + // + // TMA gtensor manipulation + // + + // Generate a TupleBasis for the gtensor + auto flat_gbasis = make_basis_like(shape(flat_glayout)); + + // Fold the flat_gbasis into the glayout + auto glayout_basis = make_layout(shape(gtensor), + stride(composition(make_layout(repeat_like(shape(flat_glayout), Int<2>{}), flat_gbasis), + make_layout(repeat_like(shape(gtensor), Int<2>{}))))); + + // Tile the modes of gtensor with the truncated cta_v_map o inv_smem_layout_trunc + auto tma_layout_v_trunc = flatten(composition(glayout_basis, layout<0>(sidx_to_gmode_vt))); + + // Append any missing basis on the end as size-1 modes b/c they got truncated + // NOTE This is essentially ArithmeticTuple complement... + auto missing_basis = fold(stride(tma_layout_v_trunc), flat_gbasis, [](auto init, auto e) { + auto k = find(init, e); + return remove(init); + }); + + // The appended map from truncated smem codomain to gmem mode: trunc_smem_idx -> gmem_mode + auto tma_layout_v = make_layout(flatten(cute::make_tuple(tma_layout_v_trunc.shape(), repeat(Int<1>{}))), + flatten(cute::make_tuple(tma_layout_v_trunc.stride(), missing_basis))); + +#if 0 + print("flat_gbasis : "); print(flat_gbasis); print("\n"); + print("missing_b : "); print(missing_basis); print("\n"); + print("tma_layout_v : "); print(tma_layout_v); print("\n"); +#endif + + // + // Construct the TMA Desc and GMEM mode ordering + // + + auto [tma_desc, gmem_stride_bases] = detail::make_tma_copy_desc(gtensor, tma_layout_v, get_swizzle_portion(slayout)); + + // + // Construct the Copy_Traits + // + + using T = typename GEngine::value_type; + constexpr int num_bits = decltype(size<0>(sidx_to_gmode_vt))::value * sizeof(T) * 8; + using Traits = Copy_Traits, decltype(gmem_stride_bases)>; + +#if 0 + print("num_bits : "); print(num_bits); print("\n"); + print("g_stride_bases: "); print(gmem_stride_bases); print("\n"); +#endif + + Traits tma_traits{tma_desc, gmem_stride_bases}; + + // + // Construct the TiledCopy + // + + auto cta_tiler = product_each(shape(cta_v_map)); + + // (CTA V, CTA T) -> smem_coord + auto layout_vt = composition(inv_smem_layout, make_layout(shape(sidx_to_gmode_vt))); + // Scale that up to cover all of the smem_coords + auto layout_VT = tile_to_shape(layout_vt, make_shape(size(cta_v_map)/size<1>(layout_vt), size<1>(layout_vt))); + // Flip it and change the domain of the T from logical thr to thr_idx + auto layout_TV = make_layout(composition(layout<1>(layout_VT), cta_t_map), layout<0>(layout_VT)); + +#if 0 + print("cta_tiler : "); print(cta_tiler); print("\n"); + print("layout_VT : "); print(layout_VT); print("\n"); + print("layout_TV : "); print(layout_TV); print("\n"); +#endif + + using T = typename GEngine::value_type; + return TiledCopy, decltype(layout_TV), decltype(cta_tiler)>{tma_traits}; +} +#endif // !defined(__MACACC_RTC__) + +} // end namespace detail + +/** Make a CuTe CTA-collective TiledCopy for a TMA operation. + * + * @param CopyOp The target copy operation: SM90_TMA_LOAD, SM90_TMA_LOAD_MULTICAST, SM90_TMA_STORE + * @param gtensor The GMEM Tensor to be involved in the TMA. + * @param slayout The SMEM Layout to be involved in the TMA. + * @param cta_tile The CTA-local tile that each CTA will be tiling GMEM with. + * This is often the blk_shape that is used to tile the GMEM for CTAs: + * local_tile(gtensor, blk_shape, blk_coord) -> CTA-local tile of gtensor + * @param cluster_size When using SM90_TMA_LOAD_MULTICAST, this can be a (static) power-of-2 <= 16 + * defining the multicast size (used to further partition the SMEM) + * Else, static-1 + * + * This code attempts to maximize the TMA box size. It does this by tracing + * the SMEM "vector" -- the inverse of the smem layout -- to find the largest + * contiguous array of smem that can be written to/from global memory given + * the constraints that the TMA instruction imposes. + * + * This is accomplished by assigning "basis" strides to the GMEM to track which + * modes of SMEM map to which modes of GMEM, then reorder the modes of GMEM according + * to the SMEM vector, and then using those GMEM/SMEM modes to fill in the desc. + * + * Examples: + using T = float; + T* gptr = nullptr; + + { + // Simple 2D + Tensor gtensor = make_tensor(gptr, make_shape(1024, 256), GenRowMajor{}); // K-Major GMEM + auto slayout = make_layout(make_shape(_64{}, _32{}), GenRowMajor{}); // K-Major SMEM + auto tma = make_tma_copy(SM90_TMA_LOAD{}, gtensor, slayout); + } + + { + // GMMA 2D + Tensor gtensor = make_tensor(gptr, make_shape(1024, 256)); // MN-Major GMEM + auto slayout = tile_to_shape(GMMA::Layout_MN_SW128_Atom{}, make_shape(_128{},_64{})); // MN-Major Swizzled+Tiled 128x64 SMEM + auto tma = make_tma_copy(SM90_TMA_LOAD{}, gtensor, slayout); + } + + { + // 3D + Tensor gtensor = make_tensor(gptr, make_shape(1024, 32, 512), make_stride(64, Int<1>{}, 65536)); // GMEM + auto slayout = make_layout(make_shape(_16{}, _8{}, _2{}), make_stride(_16{}, _1{}, _8{})); // SMEM w/ same major-mode + auto tma = make_tma_copy(SM90_TMA_LOAD{}, gtensor, slayout); + } + + { + // cuTENSOR 4D + auto layout = make_shape(make_shape(32,40),make_shape(make_shape(8,8),656)); // GMEM + auto cta_tile = make_shape(_128{},make_shape(_32{},_2{})); // GMEM Tiling: + // Take 128-elem from m: m0 must divide 128, + // m-last may be predicated + // Take 32-elem from k0, 2-elem from k1 + auto slayout = make_layout(cta_tile); // Col-Major SMEM + auto tma = make_tma_copy(SM90_TMA_LOAD{}, gtensor, slayout, cta_tile, Int<1>{}); + } + * + * Check the TMA box size and desc: + print("TMA Box size: "); print(typename decltype(tma)::Tiler_MN{}); print("\n"); + print("TMA desc : "); print(tma.tma_desc_); print("\n"); + * + * Usage: + Tensor mA = tma_a.get_tma_tensor(make_shape(M,N)); // (M,N) TMA coord tensor + Tensor gA = local_tile(mA, cta_tile, cta_coord); // (BLK_M,BLK_N) TMA coord tensor for this CTA + Tensor sA = make_tensor(make_smem_ptr(sptr), slayout); // (BLK_M,BLK_N) SMEM tensor + + auto cta_tma = tma.get_slice(cta_idx_in_cluster); // Slice for multicast partitioning + Tensor tAgA = cta_tma.partition_S(gA); // Partition for src + Tensor tAsA = cta_tma.partition_D(sA); // Partition for dst + + copy(tma.with(barrier, mcast_mask), tAgA, tAsA); // copy with supporting TMA params + */ +#if !defined(__MACACC_RTC__) +template +CUTE_HOST +auto +make_tma_copy(CopyOp const& copy_op, + Tensor const& gtensor, + SLayout const& slayout, + CTA_Tile const& cta_tile, + Cluster_Size const& cluster_size) +{ + + return detail::make_tma_copy_tiled(copy_op, + gtensor, + slayout, + make_layout(cluster_size), + make_identity_layout(cta_tile)); +} + +// Explicit defaulting +template +CUTE_HOST +auto +make_tma_copy(CopyOp const& copy_op, + Tensor const& gtensor, + SLayout const& slayout) +{ + return make_tma_copy(copy_op, gtensor, slayout, product_each(shape(slayout)), Int<1>{}); +} + +template +CUTE_HOST +auto +make_tma_copy(CopyOp const& copy_op, + Tensor const& gtensor, + SLayout const& slayout, + Cluster_Size const& cluster_size) +{ + return make_tma_copy(copy_op, gtensor, slayout, product_each(shape(slayout)), cluster_size); +} +#endif // !defined(__MACACC_RTC__) + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/atom/mma_atom.hpp b/csrc/mctlass/include/cute/atom/mma_atom.hpp new file mode 100644 index 0000000..0c11e7c --- /dev/null +++ b/csrc/mctlass/include/cute/atom/mma_atom.hpp @@ -0,0 +1,1063 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +#include + +#include + +#include + +#include + +namespace cute { + +template +struct MMA_Atom; + +template +struct MMA_Atom : MMA_Atom> +{}; + +template +struct MMA_Atom> + : MMA_Traits +{ + using Traits = MMA_Traits; + + // Element value types from the MMA_Traits + using ValTypeD = typename Traits::ElementDVal; + using ValTypeA = typename Traits::ElementAVal; + using ValTypeB = typename Traits::ElementBVal; + using ValTypeC = typename Traits::ElementCVal; + + // Thr-Val layouts from the MMA_Traits + using Shape_MNK = typename Traits::Shape_MNK; + using ThrID = typename Traits::ThrID; + using LayoutC_TV = typename Traits::CLayout; + using LayoutA_TV = typename Traits::ALayout; + using LayoutB_TV = typename Traits::BLayout; + + // Fragment value types from the MMA_Traits (optional, defaults to Val type) + using FrgTypeD = typename detail::FrgTypeC_or_Default::type; + using FrgTypeA = typename detail::FrgTypeA_or_Default::type; + using FrgTypeB = typename detail::FrgTypeB_or_Default::type; + using FrgTypeC = typename detail::FrgTypeC_or_Default::type; + + // Additional Trait parameters/transformations + template + CUTE_HOST_DEVICE + auto + with(TraitsArgs&&... args) const { + auto traits = Traits::with(std::forward(args)...); + return MMA_Atom{traits}; + } + + // + // Tensor call interfaces + // + + // Cast, check, and call fma + template + CUTE_HOST_DEVICE constexpr + void + call(Tensor & D, + Tensor const& A, + Tensor const& B, + Tensor const& C) const + { + static_assert(DLayout::rank == 1, "Expected rank-1 D tensor"); + static_assert(ALayout::rank == 1, "Expected rank-1 A tensor"); + static_assert(BLayout::rank == 1, "Expected rank-1 B tensor"); + static_assert(CLayout::rank == 1, "Expected rank-1 C tensor"); + + return mma_unpack(*this, D, A, B, C); + } + + // Three arguments reproduces C + template + CUTE_HOST_DEVICE constexpr + void + call(Tensor const& A, + Tensor const& B, + Tensor & C) const + { + return call(C, A, B, C); + } + + // + // make_fragment_A|B|C + // These functions are awkward as they expect already-partitioned tensors + // resulting from a previous call to partition_A|B|C + // The reasoning is that we can inspect the layout of the partitioned data + // and attempt to match it in generated fragment to promote vectorization + // when copying from partition to fragment. + // + + template + CUTE_HOST_DEVICE static constexpr + auto + make_fragment_C(CTensor&& ctensor) + { + // Check that this tensor is likely already partitioned + CUTE_STATIC_ASSERT_V(rank(ctensor) >= Int<3>{}); // VMN + CUTE_STATIC_ASSERT_V(size<0>(ctensor) == size<1>(LayoutC_TV{})); + // C is a bit special because we are after accumulators here + // The input/output type doesn't have to match the accumulator type + //static_assert(std::is_same::value_type>::value, "Expecting ValTypeC type"); + + // We'll never base the accumulator layout on the input tensor layout, so just return a FrgTypeC tensor + return make_tensor(shape(ctensor)); + } + + template + CUTE_HOST_DEVICE static constexpr + auto + make_fragment_A(ATensor&& atensor) + { + // Check that this tensor is likely already partitioned + CUTE_STATIC_ASSERT_V(rank(atensor) >= Int<3>{}); // VMK + CUTE_STATIC_ASSERT_V(size<0>(atensor) == size<1>(LayoutA_TV{})); + + if constexpr (has_dereference::value) { + // If the intended FrgTypeA is a view (of the current tensor), forward the whole + static_assert(is_same::value_type>::value, "Expecting ValTypeA type"); + return make_tensor(std::forward(atensor)); + } else { + // Else, the intended FrgTypeA is a value type, construct a new tensor with a fragment layout + return make_fragment_like(atensor); + } + + CUTE_GCC_UNREACHABLE; + } + + template + CUTE_HOST_DEVICE static constexpr + auto + make_fragment_B(BTensor&& btensor) + { + // Check that this tensor is likely already partitioned + CUTE_STATIC_ASSERT_V(rank(btensor) >= Int<3>{}); // VNK + CUTE_STATIC_ASSERT_V(size<0>(btensor) == size<1>(LayoutB_TV{})); + + if constexpr (has_dereference::value) { + // If the intended FrgTypeB is a view (of the current tensor), forward the whole + static_assert(is_same::value_type>::value, "Expecting ValTypeB type"); + return make_tensor(std::forward(btensor)); + } else { + // Else, the intended FrgTypeB is a value type, construct a new tensor with a fragment layout + return make_fragment_like(btensor); + } + + CUTE_GCC_UNREACHABLE; + } +}; + +// +// A tiling of mma atoms +// + +template +struct ThrMMA; + +template >, + class ValLayoutMNK = Layout>, + class PermutationsMNK = Tile> +struct TiledMMA : MMA_Atom +{ + static_assert(rank_v == 3, "TiledMMA requires rank-3 AtomLayoutMNK"); + static_assert(rank_v == 3, "TiledMMA requires rank-3 ValLayoutMNK"); + static_assert(rank_v == 3, "TiledMMA requires rank-3 PermutationsMNK"); + + using AtomShape_MNK = typename MMA_Atom::Shape_MNK; + + using AtomLayoutC_TV = typename MMA_Atom::LayoutC_TV; + using AtomLayoutA_TV = typename MMA_Atom::LayoutA_TV; + using AtomLayoutB_TV = typename MMA_Atom::LayoutB_TV; + + // ThrV -> thread_idx + using AtomThrID = typename MMA_Atom::ThrID; + + // (M,N,K) + using TiledShape_MNK = decltype(make_shape(size<0>(AtomShape_MNK{})*size<0>(AtomLayoutMNK{})*size<0>(ValLayoutMNK{}), + size<1>(AtomShape_MNK{})*size<1>(AtomLayoutMNK{})*size<1>(ValLayoutMNK{}), + size<2>(AtomShape_MNK{})*size<2>(AtomLayoutMNK{})*size<2>(ValLayoutMNK{}))); + + // thrid = (ThrV,ThrM,ThrN,ThrK) -> thr_idx + using ThrLayoutVMNK = decltype(tiled_product(AtomThrID{}, AtomLayoutMNK{})); + + // thr_idx -> (ThrV,ThrM,ThrN,ThrK) + using TidLayout = decltype(right_inverse(ThrLayoutVMNK{})); + + // Tile a tensor or a layout from shape + // (M,N,...) + // to shape + // ((ThrV,(ThrM,ThrN)),(FrgV,(RestM,RestN,...))) + // where + // ThrV: The threads local to an MMA. layout<0>(ThrLayoutVMNK): ThrV -> thread_idx + // ThrM: The threads tiled in M. layout<1>(ThrLayoutVMNK): ThrM -> thread_idx + // ThrN: The threads tiled in N. layout<2>(ThrLayoutVMNK): ThrN -> thread_idx + // FrgV: The values local to an MMA. + // RestM: The values tiled in M. + // RestN: The values tiled in N. + template + CUTE_HOST_DEVICE constexpr static + auto + thrfrg_C(CTensor&& ctensor) + { + CUTE_STATIC_ASSERT_V(rank(ctensor) >= Int<2>{}); + CUTE_STATIC_ASSERT_V(size<0>(ctensor) % size<0>(TiledShape_MNK{}) == Int<0>{}); + CUTE_STATIC_ASSERT_V(size<1>(ctensor) % size<1>(TiledShape_MNK{}) == Int<0>{}); + + // Reorder the tensor for the TiledAtom + auto t_tile = make_tile(left_inverse(get<0>(PermutationsMNK{})), + left_inverse(get<1>(PermutationsMNK{}))); + auto t_tensor = logical_divide(ctensor, t_tile); // (PermM,PermN) + + // Tile the tensor for the Atom + auto a_tile = make_tile(make_layout(size<0>(AtomShape_MNK{})), + make_layout(size<1>(AtomShape_MNK{}))); + auto a_tensor = zipped_divide(t_tensor, a_tile); // ((AtomM,AtomN),(RestM,RestN)) + + // Transform the Atom mode from (M,K) to (Thr,Val) + auto tv_tensor = a_tensor.compose(AtomLayoutC_TV{},_); // ((ThrV,FrgV),(RestM,RestN)) + + // Tile the tensor for the C-threads + auto thr_tile = make_tile(_, + make_tile(make_layout(size<1>(ThrLayoutVMNK{})), + make_layout(size<2>(ThrLayoutVMNK{})))); + auto thr_tensor = zipped_divide(tv_tensor, thr_tile); // ((ThrV,(ThrM,ThrN)),(FrgV,(RestM,RestN))) + + return thr_tensor; + } + + // Tile from (M,N,...) + // to (thr_idx,(FrgV,(RestM,RestN,...))) + template + CUTE_HOST_DEVICE constexpr static + auto + tidfrg_C(CTensor&& ctensor) + { + // Don't need a ctile composition because ThrK is last mode in TidLayout + + return thrfrg_C(ctensor).compose(TidLayout{}, _); + } + + // Tile a tensor or a layout from shape + // (M,K,...) + // to shape + // ((ThrV,(ThrM,ThrK)),(FrgV,(RestM,RestK,...))) + // where + // ThrV: The threads local to an MMA. layout<0>(ThrLayoutVMNK): ThrV -> thread_idx + // ThrM: The threads tiled in M. layout<1>(ThrLayoutVMNK): ThrM -> thread_idx + // ThrK: The threads tiled in K. layout<3>(ThrLayoutVMNK): ThrK -> thread_idx + // FrgV: The values local to an MMA. + // RestM: The values tiled in M. + // RestK: The values tiled in K. + template + CUTE_HOST_DEVICE constexpr static + auto + thrfrg_A(ATensor&& atensor) + { + CUTE_STATIC_ASSERT_V(rank(atensor) >= Int<2>{}); + CUTE_STATIC_ASSERT_V(size<0>(atensor) % size<0>(TiledShape_MNK{}) == Int<0>{}); + CUTE_STATIC_ASSERT_V(size<1>(atensor) % size<2>(TiledShape_MNK{}) == Int<0>{}); + + // Reorder the tensor for the TiledAtom + auto t_tile = make_tile(left_inverse(get<0>(PermutationsMNK{})), + left_inverse(get<2>(PermutationsMNK{}))); + auto t_tensor = logical_divide(atensor, t_tile); // (PermM,PermK) + + // Tile the tensor for the Atom + auto a_tile = make_tile(make_layout(size<0>(AtomShape_MNK{})), + make_layout(size<2>(AtomShape_MNK{}))); + auto a_tensor = zipped_divide(t_tensor, a_tile); // ((AtomM,AtomK),(RestM,RestK)) + + // Transform the Atom mode from (M,K) to (Thr,Val) + auto tv_tensor = a_tensor.compose(AtomLayoutA_TV{},_); // ((ThrV,FrgV),(RestM,RestK)) + + // Tile the tensor for the Thread + auto thr_tile = make_tile(_, + make_tile(make_layout(size<1>(ThrLayoutVMNK{})), + make_layout(size<3>(ThrLayoutVMNK{})))); + auto thr_tensor = zipped_divide(tv_tensor, thr_tile); // ((ThrV,(ThrM,ThrK)),(FrgV,(RestM,RestK))) + + return thr_tensor; + } + + // Tile from (M,K,...) + // to (thr_idx,(FrgV,(RestM,RestK,...))) + template + CUTE_HOST_DEVICE constexpr static + auto + tidfrg_A(ATensor&& atensor) + { + auto atile = make_tile(_, + make_tile(make_layout(make_shape (size<1>(ThrLayoutVMNK{}), size<2>(ThrLayoutVMNK{})), + make_stride( Int<1>{} , Int<0>{} )), + _)); + // (ThrV,(ThrM,ThrK)) -> (ThrV,(ThrM,ThrN,ThrK)) + + return thrfrg_A(atensor).compose(atile, _).compose(TidLayout{}, _); + } + + // Tile a tensor or a layout from shape + // (N,K,...) + // to shape + // ((ThrV,(ThrN,ThrK)),(FrgV,(RestN,RestK,...))) + // where + // ThrV: The threads local to an MMA. layout<0>(ThrLayoutVMNK): ThrV -> thread_idx + // ThrN: The threads tiled in N. layout<2>(ThrLayoutVMNK): ThrN -> thread_idx + // ThrK: The threads tiled in K. layout<3>(ThrLayoutVMNK): ThrK -> thread_idx + // FrgV: The values local to an MMA. + // RestN: The values tiled in N. + // RestK: The values tiled in K. + template + CUTE_HOST_DEVICE constexpr static + auto + thrfrg_B(BTensor&& btensor) + { + CUTE_STATIC_ASSERT_V(rank(btensor) >= Int<2>{}); + CUTE_STATIC_ASSERT_V(size<0>(btensor) % size<1>(TiledShape_MNK{}) == Int<0>{}); + CUTE_STATIC_ASSERT_V(size<1>(btensor) % size<2>(TiledShape_MNK{}) == Int<0>{}); + + // Reorder the tensor for the TiledAtom + auto t_tile = make_tile(left_inverse(get<1>(PermutationsMNK{})), + left_inverse(get<2>(PermutationsMNK{}))); + auto t_tensor = logical_divide(btensor, t_tile); // (PermN,PermK) + + // Tile the tensor for the Atom + auto a_tile = make_tile(make_layout(size<1>(AtomShape_MNK{})), + make_layout(size<2>(AtomShape_MNK{}))); + auto a_tensor = zipped_divide(t_tensor, a_tile); // ((AtomN,AtomK),(RestN,RestK)) + + // Transform the Atom mode from (M,K) to (Thr,Val) + auto tv_tensor = a_tensor.compose(AtomLayoutB_TV{},_); // ((ThrV,FrgV),(RestN,RestK)) + + // Tile the tensor for the Thread + auto thr_tile = make_tile(_, + make_tile(make_layout(size<2>(ThrLayoutVMNK{})), + make_layout(size<3>(ThrLayoutVMNK{})))); + auto thr_tensor = zipped_divide(tv_tensor, thr_tile); // ((ThrV,(ThrN,ThrK)),(FrgV,(RestN,RestK))) + + return thr_tensor; + } + + // Tile from (N,K,...) + // to (thr_idx,(FrgV,(RestN,RestK,...))) + template + CUTE_HOST_DEVICE constexpr static + auto + tidfrg_B(BTensor&& btensor) + { + auto btile = make_tile(_, + make_tile(make_layout(make_shape (size<1>(ThrLayoutVMNK{}), size<2>(ThrLayoutVMNK{})), + make_stride( Int<0>{} , Int<1>{} )), + _)); + // (ThrV,(ThrN,ThrK)) -> (ThrV,(ThrM,ThrN,ThrK)) + + return thrfrg_B(btensor).compose(btile, _).compose(TidLayout{}, _); + } + + template ::value)> + CUTE_HOST_DEVICE static constexpr + auto + get_slice(ThrIdx const& thr_idx) + { + auto thr_vmnk = ThrLayoutVMNK{}.get_flat_coord(thr_idx); + return ThrMMA(thr_vmnk); + } + + template ::value)> + CUTE_HOST_DEVICE static constexpr + auto + get_thread_slice(ThrIdx const& thr_idx) + { + return get_slice(thr_idx); + } + + // + // Utility for printing and visualization + // + + CUTE_HOST_DEVICE constexpr static + auto + get_layoutC_MN() + { + // (M,N) -> (M,N) + auto ref_C = make_layout(make_shape(size<0>(TiledShape_MNK{}), size<1>(TiledShape_MNK{}))); + // (cthrid,val) -> (M,N) + auto layoutC_TV = thrfrg_C(ref_C); + // (M,N) -> (cthrid,frg) + auto layoutC_MN = right_inverse(layoutC_TV).with_shape(shape(ref_C)); + + // cthrid = (v,m,n) -> thr_idx + auto thrID_C = ThrLayoutVMNK{}(_,_,_,Int<0>{}); + + return cute::make_tuple(layoutC_MN, thrID_C); + } + + CUTE_HOST_DEVICE constexpr static + auto + get_layoutC_TV() + { + // (M,N) -> (M,N) + auto ref_C = make_layout(make_shape(size<0>(TiledShape_MNK{}), size<1>(TiledShape_MNK{}))); + + return tidfrg_C(ref_C); + } + + CUTE_HOST_DEVICE constexpr static + auto + get_layoutA_MK() + { + // (M,K) -> (M,K) + auto ref_A = make_layout(make_shape(size<0>(TiledShape_MNK{}), size<2>(TiledShape_MNK{}))); + // (athrid,val) -> (M,K) + auto layoutA_TV = thrfrg_A(ref_A); + // (M,K) -> (athrid,frg) + auto layoutA_MK = right_inverse(layoutA_TV).with_shape(shape(ref_A)); + + // athrid = (v,m,k) -> thr_idx + auto thrID_A = ThrLayoutVMNK{}(_,_,Int<0>{},_); + + return cute::make_tuple(layoutA_MK, thrID_A); + } + + CUTE_HOST_DEVICE constexpr static + auto + get_layoutA_TV() + { + // (M,K) -> (M,K) + auto ref_A = make_layout(make_shape(size<0>(TiledShape_MNK{}), size<2>(TiledShape_MNK{}))); + + return tidfrg_A(ref_A); + } + + CUTE_HOST_DEVICE constexpr static + auto + get_layoutB_NK() + { + // (N,K) -> (N,K) + auto ref_B = make_layout(make_shape(size<1>(TiledShape_MNK{}), size<2>(TiledShape_MNK{}))); + // (bthrid,val) -> (N,K) + auto layoutB_TV = thrfrg_B(ref_B); + // (N,K) -> (bthrid,frg) + auto layoutB_NK = right_inverse(layoutB_TV).with_shape(shape(ref_B)); + + // bthrid = (v,n,k) -> thr_idx + auto thrID_B = ThrLayoutVMNK{}(_,Int<0>{},_,_); + + return cute::make_tuple(layoutB_NK, thrID_B); + } + + CUTE_HOST_DEVICE constexpr static + auto + get_layoutB_TV() + { + // (N,K) -> (N,K) + auto ref_B = make_layout(make_shape(size<1>(TiledShape_MNK{}), size<2>(TiledShape_MNK{}))); + + return tidfrg_B(ref_B); + } +}; + +template +struct ThrMMA : TiledMMA +{ + // Use ThrVMNK and thrfrg rather than thr_idx and tidfrg + // to support swizzled threads partitioning dynamic layouts + ThrVMNK thr_vmnk_; + + CUTE_HOST_DEVICE constexpr + ThrMMA(ThrVMNK const& thr_vmnk) : thr_vmnk_(thr_vmnk) {} + + template + CUTE_HOST_DEVICE constexpr + auto + partition_C(CTensor&& ctensor) const + { + auto thr_tensor = make_tensor(std::forward(ctensor).data(), TiledMMA::thrfrg_C(ctensor.layout())); + + auto thr_vmn = make_coord(get<0>(thr_vmnk_), make_coord(get<1>(thr_vmnk_), get<2>(thr_vmnk_))); + return thr_tensor(thr_vmn, make_coord(_, repeat(thr_tensor)>(_))); + } + + template + CUTE_HOST_DEVICE constexpr + auto + partition_A(ATensor&& atensor) const + { + auto thr_tensor = make_tensor(std::forward(atensor).data(), TiledMMA::thrfrg_A(atensor.layout())); + + auto thr_vmk = make_coord(get<0>(thr_vmnk_), make_coord(get<1>(thr_vmnk_), get<3>(thr_vmnk_))); + return thr_tensor(thr_vmk, make_coord(_, repeat(thr_tensor)>(_))); + } + + template + CUTE_HOST_DEVICE constexpr + auto + partition_B(BTensor&& btensor) const + { + auto thr_tensor = make_tensor(std::forward(btensor).data(), TiledMMA::thrfrg_B(btensor.layout())); + + auto thr_vnk = make_coord(get<0>(thr_vmnk_), make_coord(get<2>(thr_vmnk_), get<3>(thr_vmnk_))); + return thr_tensor(thr_vnk, make_coord(_, repeat(thr_tensor)>(_))); + } + + template + CUTE_HOST_DEVICE constexpr + auto + partition_fragment_C(CTensor&& ctensor) const + { + return TiledMMA::make_fragment_C(partition_C(ctensor)); + } + + template + CUTE_HOST_DEVICE constexpr + auto + partition_fragment_A(ATensor&& atensor) const + { + return TiledMMA::make_fragment_A(partition_A(atensor)); + } + + template + CUTE_HOST_DEVICE constexpr + auto + partition_fragment_B(BTensor&& btensor) const + { + return TiledMMA::make_fragment_B(partition_B(btensor)); + } +}; + +// +// These tile the MMA_Atom as a whole +// + +template >, + class MMAValLayout = Layout>, + class Permutations = Tile> +CUTE_HOST_DEVICE constexpr +auto +make_tiled_mma(MMA_Atom const&, + MMAThrLayout const& thr_layout = {}, + MMAValLayout const& val_layout = {}, + Permutations const& permutations = {}) +{ + auto thr_layout_mnk = append<3>(thr_layout, Layout<_1,_0>{}); + auto val_layout_mnk = append<3>(val_layout, Layout<_1,_0>{}); + auto permutation_mnk = append<3>(permutations, _); + + return TiledMMA, + decltype(thr_layout_mnk), + decltype(val_layout_mnk), + decltype(permutation_mnk)>{}; +} + +template >, + class MMAValLayout = Layout>, + class Permutations = Tile> +CUTE_HOST_DEVICE constexpr +auto +make_tiled_mma(MMA_Op const&, + MMAThrLayout const& thr_layout = {}, + MMAValLayout const& val_layout = {}, + Permutations const& permutations = {}) +{ + // Attempt to wrap in an MMA_Atom<> and forward + return make_tiled_mma(MMA_Atom{}, thr_layout, val_layout, permutations); +} + +// +// partition_fragment_C -- static context +// + +template +CUTE_HOST_DEVICE constexpr +auto +partition_shape_C(TiledMMA const& mma, Shape_MN const& shape_MN) +{ + constexpr int R = rank_v; + static_assert(R >= 2, "Must have at least rank-2"); + auto atomMNK = typename TiledMMA::AtomShape_MNK{}; + auto thrVMNK = typename TiledMMA::ThrLayoutVMNK{}; + auto V = shape<1>(typename TiledMMA::AtomLayoutC_TV{}); + auto M = shape_div(size<0>(shape_MN), size<0>(atomMNK) * size<1>(thrVMNK)); + auto N = shape_div(size<1>(shape_MN), size<1>(atomMNK) * size<2>(thrVMNK)); + return cute::tuple_cat(make_shape(V,M,N), take<2,R>(shape_MN)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +partition_fragment_C(TiledMMA const& mma, Shape_MN const& shapeMN) +{ + return make_tensor::FrgTypeC>(partition_shape_C(mma, shapeMN)); +} + +// partition_fragment_A and partition_fragment_B often depend on the +// layout of A and B and/or the thread_idx that is requesting the partition. +// For these reasons, they should not be used in a static context. +// See TiledMMA::get_slice(thr_idx).partition_fragment_A(tensorA) instead. + +template +CUTE_HOST_DEVICE constexpr +auto +partition_shape_A(TiledMMA const& mma, Shape_MK const& shape_MK) +{ + constexpr int R = rank_v; + static_assert(R >= 2, "Must have at least rank-2"); + auto atomMNK = typename TiledMMA::AtomShape_MNK{}; + auto thrVMNK = typename TiledMMA::ThrLayoutVMNK{}; + auto V = shape<1>(typename TiledMMA::AtomLayoutA_TV{}); + auto M = shape_div(size<0>(shape_MK), size<0>(atomMNK) * size<1>(thrVMNK)); + auto K = shape_div(size<1>(shape_MK), size<2>(atomMNK) * size<3>(thrVMNK)); + return cute::tuple_cat(make_shape(V,M,K), take<2,R>(shape_MK)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +partition_shape_B(TiledMMA const& mma, Shape_NK const& shape_NK) +{ + constexpr int R = rank_v; + static_assert(R >= 2, "Must have at least rank-2"); + auto atomMNK = typename TiledMMA::AtomShape_MNK{}; + auto thrVMNK = typename TiledMMA::ThrLayoutVMNK{}; + auto V = shape<1>(typename TiledMMA::AtomLayoutB_TV{}); + auto N = shape_div(size<0>(shape_NK), size<1>(atomMNK) * size<2>(thrVMNK)); + auto K = shape_div(size<1>(shape_NK), size<2>(atomMNK) * size<3>(thrVMNK)); + return cute::tuple_cat(make_shape(V,N,K), take<2,R>(shape_NK)); +} + +// +// Size +// + +template +CUTE_HOST_DEVICE constexpr +auto +tile_size(TiledMMA const& mma) +{ + return size(typename TiledMMA::TiledShape_MNK{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tile_shape(TiledMMA const& mma) +{ + return shape(typename TiledMMA::TiledShape_MNK{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +size(TiledMMA const& mma) +{ + return size(typename TiledMMA::ThrLayoutVMNK{}); +} + +// +// Display utilities +// + +template +CUTE_HOST_DEVICE +void +print(MMA_Atom> const&) +{ + using Atom = MMA_Atom>; + print("MMA_Atom\n"); + print(" ThrID: "); print(typename Atom::ThrID{}); print("\n"); + print(" LayoutA_TV: "); print(typename Atom::LayoutA_TV{}); print("\n"); + print(" LayoutB_TV: "); print(typename Atom::LayoutB_TV{}); print("\n"); + print(" LayoutC_TV: "); print(typename Atom::LayoutC_TV{}); print("\n"); +} + +template +CUTE_HOST_DEVICE +void +print(TiledMMA const& mma) +{ + using MMA = TiledMMA; + print("TiledMMA\n"); + print(" TiledThr: "); print(TiledThr{}); print("\n"); + print(" TiledVal: "); print(TiledVal{}); print("\n"); + print(" TiledPerm: "); print(TiledPerm{}); print("\n"); + print(" TiledShape_MNK: "); print(typename MMA::TiledShape_MNK{}); print("\n"); + print(" ThrLayoutVMNK: "); print(typename MMA::ThrLayoutVMNK{}); print("\n"); + print(static_cast(mma)); +} + +template +CUTE_HOST_DEVICE +void +print(ThrMMA const&) +{ + print(TiledMMA{}); +} + +template +CUTE_HOST_DEVICE +auto +print_latex(TiledMMA const& mma) +{ + auto layout_and_thrid_C = mma.get_layoutC_MN(); + auto layoutC_MN = get<0>(layout_and_thrid_C); + auto thrID_C = get<1>(layout_and_thrid_C); + + auto layout_and_thrid_A = mma.get_layoutA_MK(); + auto layoutA_MK = get<0>(layout_and_thrid_A); + auto thrID_A = get<1>(layout_and_thrid_A); + + auto layout_and_thrid_B = mma.get_layoutB_NK(); + auto layoutB_NK = get<0>(layout_and_thrid_B); + auto thrID_B = get<1>(layout_and_thrid_B); + + print_latex_mma(layoutC_MN, thrID_C, + layoutA_MK, thrID_A, + layoutB_NK, thrID_B); +} + +// EXPERIMENTAL -- Doesn't work with Swizzled Thr TileMMAs... +template +CUTE_HOST_DEVICE +auto +print_latex_2(TiledMMA const& mma) +{ + print_latex_mma(typename TiledMMA::TiledShape_MNK{}, + mma.get_layoutC_TV(), + mma.get_layoutA_TV(), + mma.get_layoutB_TV()); +} + +// MNK MMA Layout to console printer -- 8-value color coded by thread +template +CUTE_HOST_DEVICE +void +print_layout_mma(LayoutC const& C, ThrIDC const& TC, // (m,n) -> (tid,vid) and tid -> thr_idx + LayoutA const& A, ThrIDA const& TA, // (m,k) -> (tid,vid) and tid -> thr_idx + LayoutB const& B, ThrIDB const& TB) // (n,k) -> (tid,vid) and tid -> thr_idx +{ + CUTE_STATIC_ASSERT_V(rank(C) == Int<2>{}); + CUTE_STATIC_ASSERT_V(rank(A) == Int<2>{}); + CUTE_STATIC_ASSERT_V(rank(B) == Int<2>{}); + + assert(size<0>(A) == size<0>(C)); + assert(size<0>(B) == size<1>(C)); + assert(size<1>(A) == size<1>(B)); + + int a_width = size<1>(A) * 6 + 4; + + // Print out B (white-shifted) k-by-n + for (int k = 0; k < size<1>(B); ++k) { + // Header + printf("%*s", a_width, ""); + for (int n = 0; n < size<0>(B); ++n) printf("+-----"); + printf("+\n"); + // Values + printf("%*s", a_width, ""); + for (int n = 0; n < size<0>(B); ++n) printf("|T%02dV%1d", int(TB(B(n,k) % size(TB))), int(B(n,k) / size(TB))); + printf("|\n"); + } + // Footer + printf("%*s", a_width, ""); + for (int n = 0; n < size<0>(B); ++n) printf("+-----"); + printf("+\n\n"); + + // Print out A m-by-k and C m-by-n + for (int m = 0; m < size<0>(A); ++m) { + // Header + for (int k = 0; k < size<1>(A); ++k) printf("+-----"); + printf("+ "); + for (int n = 0; n < size<1>(C); ++n) printf("+-----"); + printf("+\n"); + // Values + for (int k = 0; k < size<1>(A); ++k) printf("|T%02dV%1d", int(TA(A(m,k) % size(TA))), int(A(m,k) / size(TA))); + printf("| "); + for (int n = 0; n < size<1>(C); ++n) printf("|T%02dV%1d", int(TC(C(m,n) % size(TC))), int(C(m,n) / size(TC))); + printf("|\n"); + } + // Footer + for (int k = 0; k < size<1>(A); ++k) printf("+-----"); + printf("+ "); + for (int n = 0; n < size<1>(C); ++n) printf("+-----"); + printf("+\n"); +} + +// MNK MMA Layout to Latex TIKZ -- 8-value color coded by thread +template +CUTE_HOST_DEVICE +void +print_latex_mma(LayoutC const& C, ThrIDC const& TC, // (m,n) -> (tid,vid) and tid -> thr_idx + LayoutA const& A, ThrIDA const& TA, // (m,k) -> (tid,vid) and tid -> thr_idx + LayoutB const& B, ThrIDB const& TB) // (n,k) -> (tid,vid) and tid -> thr_idx +{ + CUTE_STATIC_ASSERT_V(rank(C) == Int<2>{}); + CUTE_STATIC_ASSERT_V(rank(A) == Int<2>{}); + CUTE_STATIC_ASSERT_V(rank(B) == Int<2>{}); + + assert(size<0>(A) == size<0>(C)); + assert(size<0>(B) == size<1>(C)); + assert(size<1>(A) == size<1>(B)); + + char const* latex_header = + "\\documentclass{standalone}\n" + "\\usepackage{tikz}\n" + "\\usetikzlibrary{external}\n" + "\\tikzexternalize\n" + "\\begin{document}\n" + "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},box/.style={rectangle,draw=black,thick,minimum size=1cm,anchor=center}]\n\n"; + char const* latex_footer = + "\\end{tikzpicture}\n" + "\\end{document}\n"; + + char const* color_map[8] = {"{rgb,255:red,175;green,175;blue,255}", + "{rgb,255:red,175;green,255;blue,175}", + "{rgb,255:red,255;green,255;blue,175}", + "{rgb,255:red,255;green,175;blue,175}", + "{rgb,255:red,210;green,210;blue,255}", + "{rgb,255:red,210;green,255;blue,210}", + "{rgb,255:red,255;green,255;blue,210}", + "{rgb,255:red,255;green,210;blue,210}"}; + + // Header + printf("%% LayoutC: "); print(C); printf("\n"); + printf("%% ThrIDC : "); print(TC); printf("\n"); + printf("%% LayoutA: "); print(A); printf("\n"); + printf("%% ThrIDA : "); print(TA); printf("\n"); + printf("%% LayoutB: "); print(B); printf("\n"); + printf("%% ThrIDB : "); print(TB); printf("\n\n"); + + printf(latex_header); + + // C starting at 0,0 + for (int m = 0; m < size<0>(C); ++m) { + for (int n = 0; n < size<1>(C); ++n) { + int thrid = C(m,n) % size(TC); + int val_idx = C(m,n) / size(TC); + int thr_idx = TC(thrid); + + printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n", + color_map[thr_idx % 8], + m, n, + thr_idx, val_idx); + } + } + + // A starting at 0,-size<1>(A)-1 + for (int m = 0; m < size<0>(A); ++m) { + for (int k = 0; k < size<1>(A); ++k) { + int thrid = A(m,k) % size(TA); + int val_idx = A(m,k) / size(TA); + int thr_idx = TA(thrid); + + printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n", + color_map[thr_idx % 8], + m, k-1-size<1>(A), + thr_idx, val_idx); + } + } + + // B starting at -size<1>(B)-1,0 + for (int n = 0; n < size<0>(B); ++n) { + for (int k = 0; k < size<1>(B); ++k) { + int thrid = B(n,k) % size(TB); + int val_idx = B(n,k) / size(TB); + int thr_idx = TB(thrid); + + printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n", + color_map[thr_idx % 8], + k-1-size<1>(B), n, + thr_idx, val_idx); + } + } + + // A labels + for (int m = 0, k = -1; m < size<0>(A); ++m) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k-1-size<1>(A), m); + } + for (int k = 0, m = -1; k < size<1>(A); ++k) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k-1-size<1>(A), k); + } + // B labels + for (int n = 0, k = -1; n < size<0>(B); ++n) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k-1-size<1>(B), n, n); + } + for (int k = 0, n = -1; k < size<1>(B); ++k) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k-1-size<1>(B), n, k); + } + + // Footer + printf(latex_footer); +} + +// ThrVal MMA Layout to Latex TIKZ -- 8-value color coded by thread +template +CUTE_HOST_DEVICE +void +print_latex_mma(Shape_MNK const& shape_mnk, + LayoutC const& C, // (thr_idx,vid) -> (m,n) + LayoutA const& A, // (thr_idx,vid) -> (m,k) + LayoutB const& B) // (thr_idx,vid) -> (n,k) +{ + CUTE_STATIC_ASSERT_V(rank(C) == Int<2>{}); + CUTE_STATIC_ASSERT_V(rank(A) == Int<2>{}); + CUTE_STATIC_ASSERT_V(rank(B) == Int<2>{}); + + char const* latex_header = + "\\documentclass{standalone}\n" + "\\usepackage{tikz}\n" + "\\usetikzlibrary{external}\n" + "\\tikzexternalize\n" + "\\begin{document}\n" + "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},box/.style={rectangle,draw=black,thick,minimum size=1cm,anchor=center}]\n\n"; + char const* latex_footer = + "\\end{tikzpicture}\n" + "\\end{document}\n"; + + char const* color_map[8] = {"{rgb,255:red,175;green,175;blue,255}", + "{rgb,255:red,175;green,255;blue,175}", + "{rgb,255:red,255;green,255;blue,175}", + "{rgb,255:red,255;green,175;blue,175}", + "{rgb,255:red,210;green,210;blue,255}", + "{rgb,255:red,210;green,255;blue,210}", + "{rgb,255:red,255;green,255;blue,210}", + "{rgb,255:red,255;green,210;blue,210}"}; + + // Header + printf("%% Shape_MNK: "); print(shape_mnk); printf("\n"); + printf("%% LayoutC : "); print(C); printf("\n"); + printf("%% LayoutA : "); print(A); printf("\n"); + printf("%% LayoutB : "); print(B); printf("\n\n"); + + printf(latex_header); + + constexpr int M = size<0>(shape_mnk); + constexpr int N = size<1>(shape_mnk); + constexpr int K = size<2>(shape_mnk); + + // C starting at 0,0 + bool c_filled[M][N] = {}; + for (int t = 0; t < size<0>(C); ++t) { + for (int v = 0; v < size<1>(C); ++v) { + int m = C(t,v) % M; + int n = C(t,v) / M; + + if (not c_filled[m][n]) { + printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n", + color_map[t % 8], + m, n, + t, v); + c_filled[m][n] = true; + } + } + } + + // A starting at 0,-size<1>(A)-1 + bool a_filled[M][K] = {}; + for (int t = 0; t < size<0>(A); ++t) { + for (int v = 0; v < size<1>(A); ++v) { + int m = A(t,v) % M; + int k = A(t,v) / M; + + if (not a_filled[m][k]) { + printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n", + color_map[t % 8], + m, k - 1 - K, + t, v); + a_filled[m][k] = true; + } + } + } + + // B starting at -size<1>(B)-1,0 + bool b_filled[N][K] = {}; + for (int t = 0; t < size<0>(B); ++t) { + for (int v = 0; v < size<1>(B); ++v) { + int n = B(t,v) % N; + int k = B(t,v) / N; + + if (not b_filled[n][k]) { + printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n", + color_map[t % 8], + k - 1 - K, n, + t, v); + b_filled[n][k] = true; + } + } + } + + // A labels + for (int m = 0, k = -1; m < M; ++m) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k - 1 - K, m); + } + for (int k = 0, m = -1; k < K; ++k) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k - 1 - K, k); + } + // B labels + for (int n = 0, k = -1; n < N; ++n) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k - 1 - K, n, n); + } + for (int k = 0, n = -1; k < K; ++k) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k - 1 - K, n, k); + } + + // Footer + printf(latex_footer); +} + +} // namespace cute + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#include +#include +#include +#include +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/cute/atom/mma_traits.hpp b/csrc/mctlass/include/cute/atom/mma_traits.hpp new file mode 100644 index 0000000..7242e2d --- /dev/null +++ b/csrc/mctlass/include/cute/atom/mma_traits.hpp @@ -0,0 +1,208 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +namespace cute +{ + +namespace detail { + +template +struct supports_output_scaling { static constexpr bool value = false; }; + +template +struct supports_output_scaling().accumulate_)>> { static constexpr bool value = true; }; + +} // end namespace detail + +/** + * concept MMA_Traits + * { + * using ElementDVal = // Logical A-value type + * using ElementAVal = // Logical B-value type + * using ElementBVal = // Logical C-value type + * using ElementCVal = // Logical D-value type (NOTE: Not used? Assumed == ElementDVal) + * + * using ElementAFrg = // A-type consumed by MMA (if ommitted, same as ElementAVal) + * using ElementBFrg = // B_type consumed by MMA (if ommitted, same as ElementBVal) + * using ElementCFrg = // C_type consumed by MMA (if ommitted, same as ElementCVal) + * + * using Shape_MNK = // Logical MxNxK shape of the MMA + * + * using ThrID = // Logical thread id (tid) -> tidx + * + * using ALayout = // (Logical thread id (tid), Logical value id (vid)) -> Flat MK-coord + * using BLayout = // (Logical thread id (tid), Logical value id (vid)) -> Flat NK-coord + * using CLayout = // (Logical thread id (tid), Logical value id (vid)) -> Flat MN-coord + * }; + */ + +template +struct MMA_Traits +{ + static_assert(sizeof(MMAOperation) == 0, "MMA_Traits not implemented for this MMA_Operation."); +}; + +template +struct MMA_Traits> +{ + using ElementDVal = D; + using ElementAVal = A; + using ElementBVal = B; + using ElementCVal = C; + + // Logical shape of the MMA + using Shape_MNK = Shape<_1,_1,_1>; + + // Logical thread id (tid) -> tidx + using ThrID = Layout<_1>; + + // (Logical thread id (tid), Logical value id (vid)) -> coord + + // (tid,vid) -> (m,k) + using ALayout = Layout>; + // (tid,vid) -> (n,k) + using BLayout = Layout>; + // (tid,vid) -> (m,n) + using CLayout = Layout>; +}; + +// +// Generic mma_unpack for any MMA_Traits +// +template +CUTE_HOST_DEVICE constexpr +void +mma_unpack(MMA_Traits const& traits, + Tensor & D, + Tensor const& A, + Tensor const& B, + Tensor const& C) +{ + static_assert(is_rmem::value, "Expected registers in MMA_Atom::call"); + static_assert(is_rmem::value, "Expected registers in MMA_Atom::call"); + static_assert(is_rmem::value, "Expected registers in MMA_Atom::call"); + static_assert(is_rmem::value, "Expected registers in MMA_Atom::call"); + + // Register value types from the MMA_Operation register arrays + using RegTypeD = typename remove_extent::type; + using RegTypeA = typename remove_extent::type; + using RegTypeB = typename remove_extent::type; + using RegTypeC = typename remove_extent::type; + using MMATraits = MMA_Traits; + + constexpr int RegNumD = extent::value; + constexpr int RegNumA = extent::value; + constexpr int RegNumB = extent::value; + constexpr int RegNumC = extent::value; + + Tensor rA = recast(A); + Tensor rB = recast(B); + + CUTE_STATIC_ASSERT_V(size(rA) == Int{}); + CUTE_STATIC_ASSERT_V(size(rB) == Int{}); + + if constexpr (is_same::value) + { + static_assert(is_same::value, "GMMA C and D value_type must match."); + static_assert(is_same::value, "GMMA C and D layouts must match."); + // assert((void*)&C == (void*)&D); + + Tensor rC = recast(D); // NOTE: D and C are same, so use mutable D + + //CUTE_STATIC_ASSERT_V(size(rC) == Int{}); + + if constexpr (detail::supports_output_scaling::value) { + detail::explode_with_d_scaling(MMA_Op::fma, + rA, make_int_sequence{}, + rB, make_int_sequence{}, + rC, make_int_sequence{}, + traits.accumulate_); + } + else { + detail::explode(MMA_Op::fma, + rA, make_int_sequence{}, + rB, make_int_sequence{}, + rC, make_int_sequence{}); + } + } + else { + Tensor rD = recast(D); + Tensor rC = recast(C); + + CUTE_STATIC_ASSERT_V(size(rD) == Int{}); + CUTE_STATIC_ASSERT_V(size(rC) == Int{}); + if constexpr (detail::supports_output_scaling::value) { + detail::explode_with_d_scaling(MMA_Op::fma, + rD, make_int_sequence{}, + rA, make_int_sequence{}, + rB, make_int_sequence{}, + rC, make_int_sequence{}, + traits.accumulate_); + } + else { + detail::explode(MMA_Op::fma, + rD, make_int_sequence{}, + rA, make_int_sequence{}, + rB, make_int_sequence{}, + rC, make_int_sequence{}); + } + } +} + +namespace detail { + +template +struct FrgTypeA_or_Default { using type = typename X::ElementAVal; }; +template +struct FrgTypeA_or_Default> { using type = typename X::ElementAFrg; }; + +template +struct FrgTypeB_or_Default { using type = typename X::ElementBVal; }; +template +struct FrgTypeB_or_Default> { using type = typename X::ElementBFrg; }; + +template +struct FrgTypeC_or_Default { using type = typename X::ElementCVal; }; +template +struct FrgTypeC_or_Default> { using type = typename X::ElementCFrg; }; + +} // end namespace detail + +} // namespace cute diff --git a/csrc/mctlass/include/cute/atom/mma_traits_sm75.hpp b/csrc/mctlass/include/cute/atom/mma_traits_sm75.hpp new file mode 100644 index 0000000..405e871 --- /dev/null +++ b/csrc/mctlass/include/cute/atom/mma_traits_sm75.hpp @@ -0,0 +1,81 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include + +namespace cute +{ + +template <> +struct MMA_Traits +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using Shape_MNK = Shape<_16,_8,_8>; + using ThrID = Layout<_32>; + using ALayout = Layout,Shape < _2,_2>>, + Stride,Stride<_16,_1>>>; + using BLayout = Layout,_2>, + Stride,_8>>; + using CLayout = Layout,Shape < _2,_2>>, + Stride,Stride<_16,_1>>>; +}; + +/////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using Shape_MNK = Shape<_8,_8,_16>; + using ThrID = Layout<_32>; + using ALayout = Layout,_4>, + Stride,_8>>; + using BLayout = Layout,_4>, + Stride,_8>>; + using CLayout = Layout,_2>, + Stride,_8>>; +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace cute diff --git a/csrc/mctlass/include/cute/atom/mma_traits_sm80.hpp b/csrc/mctlass/include/cute/atom/mma_traits_sm80.hpp new file mode 100644 index 0000000..0e17e06 --- /dev/null +++ b/csrc/mctlass/include/cute/atom/mma_traits_sm80.hpp @@ -0,0 +1,604 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +#include + +#include + +#include + +#include + +namespace cute +{ + +namespace { + +// (T32,V1) -> (M8,N8) +using SM80_8x4 = Layout,_1>, + Stride,_0>>; +// (T32,V2) -> (M8,N8) +using SM80_8x8_Row = Layout,_2>, + Stride,_8>>; +// (T32,V4) -> (M8,N16) +using SM80_8x16_Row = Layout,_4>, + Stride,_8>>; +// (T32,V4) -> (M16,N8) +using SM80_16x8_Row = Layout,Shape < _2,_2>>, + Stride,Stride<_16,_8>>>; + +} + +/////////////////////////////////////////////////////////////////////////////// +//////////////////////// fp16 = fp16 * fp16 + fp16 //////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using Shape_MNK = Shape<_16,_8,_8>; + using ThrID = Layout<_32>; + using ALayout = SM80_16x8_Row; + using BLayout = SM80_8x8_Row; + using CLayout = SM80_16x8_Row; +}; + +template <> +struct MMA_Traits +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using Shape_MNK = Shape<_16,_8,_16>; + using ThrID = Layout<_32>; + using ALayout = Layout,Shape < _2,_2, _2>>, + Stride,Stride<_16,_8,_128>>>; + using BLayout = Layout,Shape <_2, _2>>, + Stride,Stride<_8,_64>>>; + using CLayout = SM80_16x8_Row; +}; + +/////////////////////////////////////////////////////////////////////////////// +//////////////////////// fp32 = fp16 * fp16 + fp32 //////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; +}; + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; +}; + +template <> +struct MMA_Traits +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using Shape_MNK = Shape<_16, _16, _16>; + using ThrID = Layout<_64>; + using ALayout = Layout, _4>, + Stride, _16>>; + using BLayout = Layout, _4>, + Stride, _16>>; + using CLayout = Layout, _4>, + Stride, _16>>; + +}; + +template <> +struct MMA_Traits +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using Shape_MNK = Shape<_16, _16, _32>; + using ThrID = Layout<_64>; + using ALayout = Layout, _8>, + Stride, _16>>; + using BLayout = Layout, _8>, + Stride, _16>>; + using CLayout = Layout, _4>, + Stride, _16>>; + +}; + +// use for lds4x4 + perm4x4 +template <> +struct MMA_Traits +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using Shape_MNK = Shape<_16, _64, _16>; + using ThrID = Layout<_64>; + using ALayout = Layout, _4>, + Stride, _16>>; + using BLayout = Layout, Shape<_4, _4>>, + Stride, Stride<_1, _64>>>; + using CLayout = Layout, _16>, + Stride, _16>>; + +}; + +/////////////////////////////////////////////////////////////////////////////// +//////////////////////// int32 = int8 * int8 + int32 ////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using Shape_MNK = Shape<_16, _16, _16>; + using ThrID = Layout<_64>; + using ALayout = Layout, _4>, + Stride, _16>>; + using BLayout = Layout, _4>, + Stride, _16>>; + using CLayout = Layout, _4>, + Stride, _16>>; + +}; + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using Shape_MNK = Shape<_16, _16, _32>; + using ThrID = Layout<_64>; + using ALayout = Layout, Shape<_8>>, + Stride, Stride<_16>>>; + using BLayout = Layout, Shape<_8>>, + Stride, Stride<_16>>>; + using CLayout = Layout, _4>, + Stride, _16>>; + +}; + +/////////////////////////////////////////////////////////////////////////////// +//////////////////////// fp32 = bf16 * bf16 + fp32 //////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; +}; + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; +}; + +template <> +struct MMA_Traits +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using Shape_MNK = Shape<_16, _16, _16>; + using ThrID = Layout<_64>; + using ALayout = Layout, _4>, + Stride, _16>>; + using BLayout = Layout, _4>, + Stride, _16>>; + using CLayout = Layout, _4>, + Stride, _16>>; + +}; + +template <> +struct MMA_Traits +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using Shape_MNK = Shape<_16, _16, _32>; + using ThrID = Layout<_64>; + using ALayout = Layout, _8>, + Stride, _16>>; + using BLayout = Layout, _8>, + Stride, _16>>; + using CLayout = Layout, _4>, + Stride, _16>>; + +}; + +// use for lds4x4 + perm4x4 +template <> +struct MMA_Traits +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using Shape_MNK = Shape<_16, _64, _16>; + using ThrID = Layout<_64>; + using ALayout = Layout, _4>, + Stride, _16>>; + using BLayout = Layout, Shape<_4, _4>>, + Stride, Stride<_1, _64>>>; + using CLayout = Layout, _16>, + Stride, _16>>; + +}; + +/////////////////////////////////////////////////////////////////////////////// +//////////////////////// fp32 = tf32 * tf32 + fp32 //////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = float; + using ElementAVal = mctlass::tfloat32_t; + using ElementBVal = mctlass::tfloat32_t; + using ElementCVal = float; + + using Shape_MNK = Shape<_16,_8,_4>; + using ThrID = Layout<_32>; + using ALayout = Layout,_2>, + Stride,_8>>; + using BLayout = SM80_8x4; + using CLayout = SM80_16x8_Row; +}; + +template <> +struct MMA_Traits +{ + using ElementDVal = float; + using ElementAVal = mctlass::tfloat32_t; + using ElementBVal = mctlass::tfloat32_t; + using ElementCVal = float; + + using Shape_MNK = Shape<_16,_8,_8>; + using ThrID = Layout<_32>; + using ALayout = Layout,Shape <_2, _2>>, + Stride,Stride<_8,_64>>>; + using BLayout = Layout, _2>, + Stride,_32>>; + using CLayout = SM80_16x8_Row; +}; + +/////////////////////////////////////////////////////////////////////////////// +//////////////////////// fp64 = fp64 * fp64 + fp64 //////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = double; + using ElementAVal = double; + using ElementBVal = double; + using ElementCVal = double; + + using Shape_MNK = Shape<_8,_8,_4>; + using ThrID = Layout<_32>; + using ALayout = SM80_8x4; + using BLayout = SM80_8x4; + using CLayout = SM80_8x8_Row; +}; + +// Custom complex fp64 MMA composed of 4 fp64 MMAs -- same layouts +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = complex; + using ElementAVal = complex; + using ElementBVal = complex; + using ElementCVal = complex; +}; + +// Custom complex fp64 MMA composed of 3 fp64 MMAs -- same layouts +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = typename SM80_8x8x4_GC64C64C64GC64_TN::GaussComplex; + using ElementAVal = complex; + using ElementBVal = complex; + using ElementCVal = typename SM80_8x8x4_GC64C64C64GC64_TN::GaussComplex; +}; + +/////////////////////////////////////////////////////////////////////////////// +/////////////////////////// s32 = s8 * s8 + s32 /////////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using Shape_MNK = Shape<_8,_8,_16>; + using ThrID = Layout<_32>; + using ALayout = SM80_8x16_Row; + using BLayout = SM80_8x16_Row; + using CLayout = SM80_8x8_Row; +}; + +template <> +struct MMA_Traits + : MMA_Traits {}; + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using Shape_MNK = Shape<_16,_8,_16>; + using ThrID = Layout<_32>; + using ALayout = Layout,Shape < _4,_2>>, + Stride,Stride<_16,_8>>>; + using BLayout = SM80_8x16_Row; + using CLayout = SM80_16x8_Row; +}; + +template <> +struct MMA_Traits + : MMA_Traits {}; + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using Shape_MNK = Shape<_16,_8,_32>; + using ThrID = Layout<_32>; + using ALayout = Layout,Shape < _4,_2, _2>>, + Stride,Stride<_16,_8,_256>>>; + using BLayout = Layout, Shape <_4, _2>>, + Stride, Stride<_8,_128>>>; + using CLayout = SM80_16x8_Row; +}; + +template <> +struct MMA_Traits + : MMA_Traits {}; + +/////////////////////////////////////////////////////////////////////////////// +/////////////////////////// s32 = s8 * u8 + s32 /////////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; +}; + +template <> +struct MMA_Traits + : MMA_Traits {}; + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; +}; + +template <> +struct MMA_Traits + : MMA_Traits {}; + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; +}; + +template <> +struct MMA_Traits + : MMA_Traits {}; + +/////////////////////////////////////////////////////////////////////////////// +/////////////////////////// s32 = u8 * s8 + s32 /////////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; +}; + +template <> +struct MMA_Traits + : MMA_Traits {}; + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; +}; + +template <> +struct MMA_Traits + : MMA_Traits {}; + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; +}; + +template <> +struct MMA_Traits + : MMA_Traits {}; + +/////////////////////////////////////////////////////////////////////////////// +/////////////////////////// s32 = u8 * u8 + s32 /////////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; +}; + +template <> +struct MMA_Traits + : MMA_Traits {}; + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; +}; + +template <> +struct MMA_Traits + : MMA_Traits {}; + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; +}; + +template <> +struct MMA_Traits + : MMA_Traits {}; + +/////////////////////////////////////////////////////////////////////////////// +/////////////////////////// s32 = b1 ^ b1 + s32 /////////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = cute::uint1b_t; + using ElementBVal = cute::uint1b_t; + using ElementCVal = int32_t; + + using Shape_MNK = Shape<_16,_8,_256>; + using ThrID = Layout<_32>; + using ALayout = Layout>, + Stride<_64,Stride<_64,_16,_8,_2048>>>; + using BLayout = Layout>, + Stride<_32,Stride< _1,_1024>>>; + using CLayout = SM80_16x8_Row; +}; +} // end namespace cute diff --git a/csrc/mctlass/include/cute/atom/mma_traits_sm90.hpp b/csrc/mctlass/include/cute/atom/mma_traits_sm90.hpp new file mode 100644 index 0000000..b7a12b9 --- /dev/null +++ b/csrc/mctlass/include/cute/atom/mma_traits_sm90.hpp @@ -0,0 +1,132 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +#include + +#include + +namespace cute { + +/////////////////////////////////////////////////////////////////////////////// +//////////////////////// fp64 = fp64 * fp64 + fp64 //////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = double; + using ElementAVal = double; + using ElementBVal = double; + using ElementCVal = double; + + using Shape_MNK = Shape<_16,_8,_4>; + using ThrID = Layout<_32>; + using ALayout = Layout,_2>, + Stride,_8>>; + using BLayout = Layout,_1>, + Stride,_0>>; + using CLayout = Layout,Shape < _2,_2>>, + Stride,Stride<_16,_8>>>; +}; + +template <> +struct MMA_Traits +{ + using ElementDVal = double; + using ElementAVal = double; + using ElementBVal = double; + using ElementCVal = double; + + using Shape_MNK = Shape<_16,_8,_8>; + using ThrID = Layout<_32>; + using ALayout = Layout,Shape <_2, _2>>, + Stride,Stride<_8,_64>>>; + using BLayout = Layout, _2>, + Stride,_32>>; + using CLayout = Layout,Shape < _2,_2>>, + Stride,Stride<_16,_8>>>; +}; + +template <> +struct MMA_Traits +{ + using ElementDVal = double; + using ElementAVal = double; + using ElementBVal = double; + using ElementCVal = double; + + using Shape_MNK = Shape<_16,_8,_16>; + using ThrID = Layout<_32>; + using ALayout = Layout,Shape <_2, _4>>, + Stride,Stride<_8,_64>>>; + using BLayout = Layout, _4>, + Stride,_32>>; + using CLayout = Layout,Shape < _2,_2>>, + Stride,Stride<_16,_8>>>; +}; + +/////////////////////////////////////////////////////////////////////////////////// +//////////////////////// cfp64 = cfp64 * cfp64 + cfp64 //////////////////////////// +/////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = complex; + using ElementAVal = complex; + using ElementBVal = complex; + using ElementCVal = complex; +}; + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = complex; + using ElementAVal = complex; + using ElementBVal = complex; + using ElementCVal = complex; +}; + +template <> +struct MMA_Traits + : MMA_Traits +{ + using ElementDVal = complex; + using ElementAVal = complex; + using ElementBVal = complex; + using ElementCVal = complex; +}; + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/atom/mma_traits_sm90_gmma.hpp b/csrc/mctlass/include/cute/atom/mma_traits_sm90_gmma.hpp new file mode 100644 index 0000000..752023c --- /dev/null +++ b/csrc/mctlass/include/cute/atom/mma_traits_sm90_gmma.hpp @@ -0,0 +1,3185 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +#include + +#include + +namespace cute { + +// Fence between the async destination accumulators of GMMA & source for their dependent use +template +CUTE_HOST_DEVICE +void +warpgroup_fence_operand(Tensor& frg) { + CUTE_STATIC_ASSERT(is_static::value); + if constexpr (is_same_v) { + auto f32_frg = recast(frg); + CUTE_UNROLL + for (int i = 0; i < size(f32_frg); ++i) { + warpgroup_fence_operand(f32_frg(i)); + } + } + else { + CUTE_STATIC_ASSERT(is_rmem::value); + auto u32_frg = recast(frg); + CUTE_UNROLL + for (int i = 0; i < size(u32_frg); ++i) { + warpgroup_fence_operand(u32_frg(i)); + } + } +} + +namespace GMMA { + +/////////////////////////////////////////// +// Common layouts for GMMA Shared Memory // +/////////////////////////////////////////// + +// M|N-major GMMA layouts in units of bits +using Layout_MN_INTER_Atom_Bits = Layout,Stride<_1,_128>>; +using Layout_MN_SW32_Atom_Bits = ComposedLayout, smem_ptr_flag, Layout,Stride<_1, _256>>>; +using Layout_MN_SW64_Atom_Bits = ComposedLayout, smem_ptr_flag, Layout,Stride<_1, _512>>>; +using Layout_MN_SW128_Atom_Bits = ComposedLayout, smem_ptr_flag, Layout,Stride<_1,_1024>>>; + +// K-major GMMA layouts in units of bits +using Layout_K_INTER_Atom_Bits = Layout,Stride<_128,_1>>; +using Layout_K_SW32_Atom_Bits = ComposedLayout, smem_ptr_flag, Layout,Stride< _256,_1>>>; +using Layout_K_SW64_Atom_Bits = ComposedLayout, smem_ptr_flag, Layout,Stride< _512,_1>>>; +using Layout_K_SW128_Atom_Bits = ComposedLayout, smem_ptr_flag, Layout,Stride<_1024,_1>>>; + +// M|N-major layouts in units of Type +template +using Layout_MN_INTER_Atom = decltype(upcast::value>(Layout_MN_INTER_Atom_Bits{})); +template +using Layout_MN_SW32_Atom = decltype(upcast::value>(Layout_MN_SW32_Atom_Bits{})); +template +using Layout_MN_SW64_Atom = decltype(upcast::value>(Layout_MN_SW64_Atom_Bits{})); +template +using Layout_MN_SW128_Atom = decltype(upcast::value>(Layout_MN_SW128_Atom_Bits{})); + +// K-major layouts in units of Type +template +using Layout_K_INTER_Atom = decltype(upcast::value>(Layout_K_INTER_Atom_Bits{})); +template +using Layout_K_SW32_Atom = decltype(upcast::value>(Layout_K_SW32_Atom_Bits{})); +template +using Layout_K_SW64_Atom = decltype(upcast::value>(Layout_K_SW64_Atom_Bits{})); +template +using Layout_K_SW128_Atom = decltype(upcast::value>(Layout_K_SW128_Atom_Bits{})); + +// With GMMA::Major param +template +using Layout_INTER_Atom = typename conditional, + Layout_K_INTER_Atom>::type; +template +using Layout_SW32_Atom = typename conditional, + Layout_K_SW32_Atom>::type; +template +using Layout_SW64_Atom = typename conditional, + Layout_K_SW64_Atom>::type; +template +using Layout_SW128_Atom = typename conditional, + Layout_K_SW128_Atom>::type; + +// +// Tensor to LayoutType utility +// + +// smem_ptr_swizzle LayoutType +template +CUTE_HOST_DEVICE constexpr +LayoutType +layout_type(Tensor>>, + Layout> const&) +{ + static_assert(M == 4, "Unsupported layout swizzle"); + static_assert(0 <= B && B <= 3, "Unsupported layout swizzle"); + static_assert(S == 3, "Unsupported layout swizzle"); + + switch (B) { + case 0: return LayoutType::INTERLEAVE; + case 1: return LayoutType::B32; + case 2: return LayoutType::B64; + case 3: return LayoutType::B128; + } + return LayoutType::INTERLEAVE; // ERROR +} + +// smem_ptr non-swizzled LayoutType +template +CUTE_HOST_DEVICE constexpr +LayoutType +layout_type(Tensor>, + Layout> const&) +{ + return LayoutType::INTERLEAVE; +} + +/////////////////////////////////////////////////////////////////////////////// +// Construction method for GMMA Descriptors +/////////////////////////////////////////////////////////////////////////////// + +/** +* /////////////////////////////// +* // make_gmma_desc // +* /////////////////////////////// +* Each GmmaDescriptor Major-MN describes a canonical layout of the form +* +* LayoutType::INTERLEAVE : Swizzle<0,4,3> o smem_ptr o ((T,1,m),(8,k)):((1,T,SBO),(1T,LBO)) +* LayoutType::B32 : Swizzle<1,4,3> o smem_ptr o ((T,2,m),(8,k)):((1,T,LBO),(2T,SBO)) +* LayoutType::B64 : Swizzle<2,4,3> o smem_ptr o ((T,4,m),(8,k)):((1,T,LBO),(4T,SBO)) +* LayoutType::B128 : Swizzle<3,4,3> o smem_ptr o ((T,8,m),(8,k)):((1,T,LBO),(8T,SBO)) +* +* where +* T : sizeof(uint128_t) / sizeof(value_type) +* m : integer in [1,16] corresponding to GMMA shape +* k : integer in [1,32] corresponding to GMMA shape +* SBO: stride byte offset +* LBO: leading byte offset +* +* See GMMA::Layout_MN_XXX_Atom for building canonical GmmaDescriptor Major-MN layouts. +* For example, +* auto smem_layout = tile_to_shape(Layout_MN_SW128_Atom{}, Shape<_128,_64>{}); +* is guaranteed to be accepted by make_gmma_desc for appropriate value_type. +* +* ////////////////////////////// +* // make_gmma_desc // +* ////////////////////////////// +* Each GmmaDescriptor Major-K describes a canonical layout of the form +* +* LayoutType::INTERLEAVE : Swizzle<0,4,3> o smem_ptr o ((8,m),(T,2)):((1T,SBO),(1,LBO)) +* LayoutType::B32 : Swizzle<1,4,3> o smem_ptr o ((8,m),(T,2)):((2T,SBO),(1, T )) +* LayoutType::B64 : Swizzle<2,4,3> o smem_ptr o ((8,m),(T,2)):((4T,SBO),(1, T )) +* LayoutType::B128 : Swizzle<3,4,3> o smem_ptr o ((8,m),(T,2)):((8T,SBO),(1, T )) +* +* See GMMA::Layout_K_XXX_Atom for building canonical GmmaDescriptor Major-K layouts. +* For example, +* auto smem_layout = tile_to_shape(Layout_K_SW128_Atom{}, Shape<_128,_64>{}); +* is guaranteed to be accepted by make_gmma_desc for appropriate value_type. +*/ +template +CUTE_HOST_DEVICE constexpr +GmmaDescriptor +make_gmma_desc(Tensor const& tensor) +{ + static_assert(is_smem::value, "GMMA Descriptors can only be constructed on smem."); + static_assert(TLayout::rank == 2, "GMMA Descriptors can only be constructed on rank-2 tensors."); + using value_type = typename TEngine::value_type; + + Tensor u128_tensor = recast(tensor); + + // Result + GmmaDescriptor desc; + + // Layout type + constexpr GMMA::LayoutType LAYOUT_TYPE = GMMA::layout_type(u128_tensor); + desc.layout_type_ = uint8_t(LAYOUT_TYPE); + + // Start address (4LSB not included) + uint32_t start_address = cast_smem_ptr_to_uint(u128_tensor.data().get()); + desc.start_address_ = start_address >> 4; + + constexpr uint8_t base_offset = 0; + desc.base_offset_ = base_offset; + + // LayoutType meta + constexpr int W = LAYOUT_TYPE == GMMA::LayoutType::INTERLEAVE ? 1 : + LAYOUT_TYPE == GMMA::LayoutType::B32 ? 2 : + LAYOUT_TYPE == GMMA::LayoutType::B64 ? 4 : + LAYOUT_TYPE == GMMA::LayoutType::B128 ? 8 : -1; + + if constexpr (MajorMode == GMMA::Major::MN) + { + /* In units of uint128_t, each GmmaDescriptor Major-MN describes a canonical layout of the form + * + * LayoutType::INTERLEAVE : Swizzle<0,4,3> o smem_ptr o ((1,n),(8,k)):((X,SBO),(1,LBO)) + * LayoutType::B32 : Swizzle<1,4,3> o smem_ptr o ((2,n),(8,k)):((1,LBO),(2,SBO)) + * LayoutType::B64 : Swizzle<2,4,3> o smem_ptr o ((4,n),(8,k)):((1,LBO),(4,SBO)) + * LayoutType::B128 : Swizzle<3,4,3> o smem_ptr o ((8,n),(8,k)):((1,LBO),(8,SBO)) + */ + static_assert(size<1>(u128_tensor) == Int<(256 / cute::sizeof_bits::value)>{}, // K size + "Not a canonical GMMA_MN Layout: Expected K-size 256/sizeof_bits."); + + // Construct the canonical GMMA T Layout with shape ((W,n),(8,2)) + Layout canonical_layout = logical_divide(layout(u128_tensor), make_tile(Layout,_1>{}, Layout,_1>{})); + + // Check ranks of canonical + CUTE_STATIC_ASSERT_V(rank<0>(canonical_layout) == Int<2>{}, "Not a canonical GMMA_MN Layout: No flat offset mode"); + CUTE_STATIC_ASSERT_V(rank<1>(canonical_layout) == Int<2>{}, "Not a canonical GMMA_MN Layout: No flat offset mode"); + // Check canonical mode strides + constexpr uint32_t stride_00 = stride<0,0>(canonical_layout); + constexpr uint32_t expected_stride_00 = LAYOUT_TYPE == GMMA::LayoutType::INTERLEAVE ? stride<0,0>(canonical_layout) : 1; + static_assert(stride_00 == expected_stride_00, "Not a canonical GMMA_MN Layout: Expected stride failure."); + constexpr uint32_t stride_10 = stride<1,0>(canonical_layout); + constexpr uint32_t expected_stride_10 = W; + static_assert(stride_10 == expected_stride_10, "Not a canonical GMMA_MN Layout: Expected stride failure."); + + // stride dimension byte offset and leading dimension byte offset (4LSB not included == uint128_t units) + constexpr uint32_t stride_01 = stride<0,1>(canonical_layout); + constexpr uint32_t stride_11 = stride<1,1>(canonical_layout); + + desc.stride_byte_offset_ = (LAYOUT_TYPE == GMMA::LayoutType::INTERLEAVE) ? stride_01 : stride_11; + desc.leading_byte_offset_ = (LAYOUT_TYPE == GMMA::LayoutType::INTERLEAVE) ? stride_11 : stride_01; + } + else if constexpr (MajorMode == GMMA::Major::K) + { + /* In units of uint128_t, each GmmaDescriptor Major-K describes a canonical layout of the form + * + * LayoutType::INTERLEAVE : Swizzle<0,4,3> o smem_ptr o ((8,n),2):((1,SBO),LBO) + * LayoutType::B32 : Swizzle<1,4,3> o smem_ptr o ((8,n),2):((2,SBO),1) + * LayoutType::B64 : Swizzle<2,4,3> o smem_ptr o ((8,n),2):((4,SBO),1) + * LayoutType::B128 : Swizzle<3,4,3> o smem_ptr o ((8,n),2):((8,SBO),1) + */ + CUTE_STATIC_ASSERT_V(size<0>(u128_tensor) % Int<8>{} == Int<0>{}, // N|M size + "Not a canonical GMMA_K Layout: Expected MN-size multiple of 8."); + CUTE_STATIC_ASSERT_V(size<1>(u128_tensor) == Int<2>{}, // K size + "Not a canonical GMMA_K Layout: Expected K-size 2 (in units of uint128_t)."); + + // Construct the canonical GMMA N Layout with shape ((8,n),(2,1)) + Layout canonical_layout = logical_divide(layout(u128_tensor), make_tile(Layout<_8,_1>{}, Layout<_2,_1>{})); + + // Check ranks of canonical + CUTE_STATIC_ASSERT_V(rank<0>(canonical_layout) == Int<2>{}, "Not a canonical GMMA_K Layout: No flat offset mode"); + CUTE_STATIC_ASSERT_V(rank<1>(canonical_layout) == Int<2>{}, "Not a canonical GMMA_K Layout: No flat offset mode"); + // Check canonical mode strides + constexpr uint32_t stride_00 = stride<0,0>(canonical_layout); + constexpr uint32_t expected_stride_00 = W; + static_assert(stride_00 == expected_stride_00, "Not a canonical GMMA_K Layout: Expected stride failure."); + constexpr uint32_t stride_10 = stride<1,0>(canonical_layout); + constexpr uint32_t expected_stride_10 = (LAYOUT_TYPE == GMMA::LayoutType::INTERLEAVE) ? stride<1,0>(canonical_layout) : 1; + static_assert(stride_10 == expected_stride_10, "Not a canonical GMMA_K Layout: Expected stride failure."); + + // stride dimension byte offset and leading dimension byte offset (4LSB not included == uint128_t units) + constexpr uint32_t stride_01 = stride<0,1>(canonical_layout); + + desc.stride_byte_offset_ = stride_01; + desc.leading_byte_offset_ = stride_10; + } else { + static_assert(MajorMode != GMMA::Major::MN && MajorMode != GMMA::Major::K, "Unrecognized MajorMode!"); + } + +#if 0 + // DEBUG and SANITY + assert((start_address & 0b0000001111) == 0); // Must be 16B aligned (4LSB are 0) no negotiation + assert((start_address & 0b1110000000) == 0); // Assert base_offset is 0, generalize later + if (thread0()) { + print("smem_desc input tensor: "); print(tensor.data()); print(" o "); print(tensor.layout()); print("\n"); + print("smem_desc uint128_t tensor: "); print(u128_tensor.data()); print(" o "); print(u128_tensor.layout()); print("\n"); + //print(" desc canonical layout: "); print(canonical_layout); print("\n"); + print(desc); + } +#endif + + return desc; +} + +/////////////////////////////////////////////////////////////////////////////// +// Higher level GMMA Descriptor utilities +/////////////////////////////////////////////////////////////////////////////// + +struct DescriptorIterator +{ + GmmaDescriptor desc_; + + // Dereference returns the GmmaDescriptor + CUTE_HOST_DEVICE constexpr + GmmaDescriptor const& operator*() const { return desc_; } + + // Advance and return a new GmmaDescriptor + template + CUTE_HOST_DEVICE constexpr + GmmaDescriptor operator[](Index const& i) const { return *(*this + i); } + + // Return an advanced iterator + template + CUTE_HOST_DEVICE constexpr + DescriptorIterator operator+(Index const& offset) const + { + // offset is in the units of uint128_t (4LSB of start_address not included) + + //GmmaDescriptor desc = desc_; + //desc.start_address_ += uint16_t(offset); + //desc.reg32_[0] += uint16_t(offset); // Generates better asm than adding to the bitfield + + // May need to update base_offset if swizzle alignment isn't guaranteed + //desc.base_offset_ = 0; + //assert((desc.start_address_ & 0b111000) == 0); // Assert base_offset is 0, generalize later + + //return {desc}; + + // The above seems to not work for some reason... + return { GmmaDescriptor {desc_ + uint64_t(offset)} }; + } + + CUTE_HOST_DEVICE friend void + print(DescriptorIterator const&) { printf("GMMA::DescriptorIterator"); } +}; + +// The GMMA Traits below have custom fragment type flags for their smem desc tensors. +// These flags specialize a MakeTensor customization point to correctly make the fragment that is desired. +template +struct smem_desc : DescriptorIterator {}; + +// Recast a DescriptorIterator Tensor to uint64_t, it's RegType +template +CUTE_HOST_DEVICE constexpr +auto +recast(Tensor,TLayout> const& tensor, type_list) +{ + static_assert(is_same::value, "Can only cast descriptors to uint64_t."); + return make_tensor(tensor.data(), Layout<_1,_0>{}); +} + +} // end namespace GMMA + +// Customization point for creating a GMMA::smem_desc Tensor +template +struct MakeTensor> +{ + template + CUTE_HOST_DEVICE constexpr auto + operator()(Tensor const& smem_tensor) + { + static_assert(is_smem::value, "Expected SMEM Tensor to construct a GMMA Desc Tensor"); + return make_tensor(GMMA::DescriptorIterator{GMMA::make_gmma_desc(tensor<0>(smem_tensor))}, + recast(smem_tensor).layout()); + } +}; + +/////////////////////////////////////////////////////////////////////////////// +//////////////////////////// MMA_TRAITS /////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +namespace GMMA { + +// Accumulator layouts +using CLayout_64x8 = Layout,Shape < _2,_2>>, + Stride,Stride<_64,_8>>>; + +using CLayout_64x16 = Layout,Shape < _2,_2, _2>>, + Stride,Stride<_64,_8,_512>>>; + +using CLayout_64x32 = Layout,Shape < _2,_2, _4>>, + Stride,Stride<_64,_8,_512>>>; + +using CLayout_64x64 = Layout,Shape < _2,_2, _8>>, + Stride,Stride<_64,_8,_512>>>; + +using CLayout_64x96 = Layout,Shape < _2,_2, _12>>, + Stride,Stride<_64,_8,_512>>>; + +using CLayout_64x128 = Layout,Shape < _2,_2, _16>>, + Stride,Stride<_64,_8,_512>>>; + +using CLayout_64x192 = Layout,Shape < _2,_2, _24>>, + Stride,Stride<_64,_8,_512>>>; + +using CLayout_64x256 = Layout,Shape < _2,_2, _32>>, + Stride,Stride<_64,_8,_512>>>; + +// Register source layout for 32-bit value types +using ALayout_64x8 = Layout,Shape < _2, _2>>, + Stride,Stride< _8,_256>>>; + +// Register source layout for 16-bit value types +using ALayout_64x16 = CLayout_64x16; + +// Register source layout for 8-bit value types +using ALayout_64x32 = Layout,Shape < _4,_2, _2>>, + Stride,Stride<_64,_8,_1024>>>; + +// Shared memory source layouts for any value type +template +using ABLayout = Layout,Int>>, + Stride< _0,Stride< _1,Int>>>; + +} // namespace GMMA + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 8, 16>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 8, 16>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 16, 16>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 16, 16>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 32, 16>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 32, 16>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 64, 16>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 64, 16>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 96, 16>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 96, 16>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout<128, 16>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout<128, 16>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout<192, 16>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout<192, 16>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout<256, 16>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = half_t; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = half_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout<256, 16>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 8, 16>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 8, 16>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 16, 16>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 16, 16>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 32, 16>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 32, 16>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 64, 16>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 64, 16>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 96, 16>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 96, 16>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout<128, 16>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout<128, 16>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout<192, 16>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout<192, 16>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout<256, 16>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = half_t; + using ElementBVal = half_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout<256, 16>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 8, 16>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 8, 16>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 16, 16>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 16, 16>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 32, 16>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 32, 16>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 64, 16>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 64, 16>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout< 96, 16>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout< 96, 16>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout<128, 16>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout<128, 16>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout<192, 16>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout<192, 16>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 16>; + using BLayout = GMMA::ABLayout<256, 16>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = bfloat16_t; + using ElementBVal = bfloat16_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_16>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x16; + using BLayout = GMMA::ABLayout<256, 16>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 8>; + using BLayout = GMMA::ABLayout< 8, 8>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x8; + using BLayout = GMMA::ABLayout< 8, 8>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 8>; + using BLayout = GMMA::ABLayout< 16, 8>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x8; + using BLayout = GMMA::ABLayout< 16, 8>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 8>; + using BLayout = GMMA::ABLayout< 32, 8>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x8; + using BLayout = GMMA::ABLayout< 32, 8>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 8>; + using BLayout = GMMA::ABLayout< 64, 8>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x8; + using BLayout = GMMA::ABLayout< 64, 8>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 8>; + using BLayout = GMMA::ABLayout< 96, 8>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x8; + using BLayout = GMMA::ABLayout< 96, 8>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 8>; + using BLayout = GMMA::ABLayout<128, 8>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x8; + using BLayout = GMMA::ABLayout<128, 8>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 8>; + using BLayout = GMMA::ABLayout<192, 8>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x8; + using BLayout = GMMA::ABLayout<192, 8>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 8>; + using BLayout = GMMA::ABLayout<256, 8>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct MMA_Traits> +{ + using ElementDVal = float; + using ElementAVal = tfloat32_t; + using ElementBVal = tfloat32_t; + using ElementCVal = float; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_8>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x8; + using BLayout = GMMA::ABLayout<256, 8>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 8, 32>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 16, 32>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 32, 32>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 64, 32>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 96, 32>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout<128, 32>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout<192, 32>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout<256, 32>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 8, 32>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 16, 32>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 32, 32>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 64, 32>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 96, 32>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout<128, 32>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout<192, 32>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout<256, 32>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 8, 32>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 16, 32>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 32, 32>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 64, 32>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 96, 32>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout<128, 32>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout<192, 32>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout<256, 32>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 8, 32>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 16, 32>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 32, 32>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 64, 32>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 96, 32>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout<128, 32>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout<192, 32>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = int8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout<256, 32>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 8, 32>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 16, 32>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 32, 32>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 64, 32>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 96, 32>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout<128, 32>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout<192, 32>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout<256, 32>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 8, 32>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 16, 32>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 32, 32>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 64, 32>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 96, 32>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout<128, 32>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout<192, 32>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = int8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout<256, 32>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 8, 32>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 16, 32>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 32, 32>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 64, 32>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout< 96, 32>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout<128, 32>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout<192, 32>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementAFrg = GMMA::smem_desc; + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ABLayout< 64, 32>; + using BLayout = GMMA::ABLayout<256, 32>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_8,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 8, 32>; + using CLayout = GMMA::CLayout_64x8; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_16,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 16, 32>; + using CLayout = GMMA::CLayout_64x16; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_32,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 32, 32>; + using CLayout = GMMA::CLayout_64x32; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_64,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 64, 32>; + using CLayout = GMMA::CLayout_64x64; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_96,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout< 96, 32>; + using CLayout = GMMA::CLayout_64x96; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_128,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout<128, 32>; + using CLayout = GMMA::CLayout_64x128; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_192,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout<192, 32>; + using CLayout = GMMA::CLayout_64x192; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +struct MMA_Traits +{ + using ElementDVal = int32_t; + using ElementAVal = uint8_t; + using ElementBVal = uint8_t; + using ElementCVal = int32_t; + + using ElementBFrg = GMMA::smem_desc; + + using Shape_MNK = Shape<_64,_256,_32>; + using ThrID = Layout<_128>; + using ALayout = GMMA::ALayout_64x32; + using BLayout = GMMA::ABLayout<256, 32>; + using CLayout = GMMA::CLayout_64x256; + + GMMA::ScaleOut accumulate_ = GMMA::ScaleOut::One; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/config.hpp b/csrc/mctlass/include/cute/config.hpp new file mode 100644 index 0000000..b7da1fe --- /dev/null +++ b/csrc/mctlass/include/cute/config.hpp @@ -0,0 +1,169 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#if defined(__MACA_ARCH__) || defined(__clang__) +# define CUTE_HOST_DEVICE __forceinline__ __host__ __device__ +# define CUTE_DEVICE __forceinline__ __device__ +# define CUTE_HOST __forceinline__ __host__ +#else +# define CUTE_HOST_DEVICE inline +# define CUTE_DEVICE inline +# define CUTE_HOST inline +#endif // CUTE_HOST_DEVICE, CUTE_DEVICE + +#if !defined(__MACACC_RTC__) && (defined(__MACA_ARCH__)) +# define CUTE_UNROLL _Pragma("unroll") +# define CUTE_NO_UNROLL _Pragma("unroll 1") +#elif defined(__MACACC_RTC__) +# define CUTE_UNROLL _Pragma("unroll") +# define CUTE_NO_UNROLL _Pragma("unroll 1") +#else +# define CUTE_UNROLL +# define CUTE_NO_UNROLL +#endif // CUTE_UNROLL + +#if defined(__MACA_ARCH__) +# define CUTE_INLINE_CONSTANT static const __device__ +#else +# define CUTE_INLINE_CONSTANT static constexpr +#endif + +#if defined(__MACA_ARCH__) && (__MACA_ARCH__ == 1000) +# define CUTE_MACA_XCORE1000_ENABLED +#endif + +#if defined(__MACA_ARCH__) && (__MACA_ARCH__ == 1500) +# define CUTE_MACA_XCORE1500_ENABLED +#endif + +// __grid_constant__ was introduced in CUDA 11.7. +// #if ((__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 7))) +#if 0 +# define CUTE_GRID_CONSTANT_SUPPORTED +#endif + +// __grid_constant__ can be enabled only on SM70+. +#if defined(__MACA_ARCH__) +# define CUTE_GRID_CONSTANT_ENABLED +#endif + +#if ! defined(CUTE_GRID_CONSTANT) +# if defined(CUTE_GRID_CONSTANT_SUPPORTED) && defined(CUTE_GRID_CONSTANT_ENABLED) +# define CUTE_GRID_CONSTANT __grid_constant__ +# else +# define CUTE_GRID_CONSTANT +# endif +#endif + +// Some versions of GCC < 11 have trouble deducing that a +// function with "auto" return type and all of its returns in an "if +// constexpr ... else" statement must actually return. Thus, GCC +// emits spurious "missing return statement" build warnings. +// Developers can suppress these warnings by using the +// CUTE_GCC_UNREACHABLE macro, which must be followed by a semicolon. +// It's harmless to use the macro for other GCC versions or other +// compilers, but it has no effect. +#if ! defined(CUTE_GCC_UNREACHABLE) +# if defined(__GNUC__) && __GNUC__ < 11 + // GCC 10, but not 7.5, 9.4.0, or 11, issues "missing return + // statement" warnings without this little bit of help. +# define CUTE_GCC_UNREACHABLE __builtin_unreachable() +# else +# define CUTE_GCC_UNREACHABLE +# endif +#endif + +#ifdef _MSC_VER +// Provides support for alternative operators 'and', 'or', and 'not' +#include +#endif // _MSC_VER + +#if defined(__MACACC_RTC__) +#define CUTE_STL_NAMESPACE cuda::std +#define CUTE_STL_NAMESPACE_IS_CUDA_STD +#else +#define CUTE_STL_NAMESPACE std +#endif + +// +// Assertion helpers +// + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#define CUTE_STATIC_ASSERT static_assert +#define CUTE_STATIC_ASSERT_V(x,...) static_assert(decltype(x)::value, ##__VA_ARGS__) + +#if defined(__MACA_ARCH__) +# define CUTE_RUNTIME_ASSERT(x) assert(0 && x);__brkpt() +#else +# define CUTE_RUNTIME_ASSERT(x) assert(0 && x) +#endif + +// +// IO +// + +#if !defined(__MACACC_RTC__) +#include +#include +#include +#endif + +// +// Support +// + +#include + +// +// Basic types +// + +#include +#include +#include +#include +#include +#include +#include + +// +// Debugging utilities +// + +#include +#include diff --git a/csrc/mctlass/include/cute/container/alignment.hpp b/csrc/mctlass/include/cute/container/alignment.hpp new file mode 100644 index 0000000..dc73f21 --- /dev/null +++ b/csrc/mctlass/include/cute/container/alignment.hpp @@ -0,0 +1,70 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include + +namespace cute +{ + +// Test if a pointer is aligned to N bytes +template +CUTE_HOST_DEVICE constexpr +bool +is_byte_aligned(void const* const ptr) +{ + static_assert(N > 0 && (N & (N - 1)) == 0, "N must be a power of 2 in alignment check"); + return (reinterpret_cast(ptr) & (N-1)) == 0; +} + +#if defined(__MACACC__) +# define CUTE_ALIGNAS(n) __align__(n) +#else +# define CUTE_ALIGNAS(n) alignas(n) +#endif + +template +struct aligned_struct {}; + +template <> struct CUTE_ALIGNAS( 1) aligned_struct< 1> {}; +template <> struct CUTE_ALIGNAS( 2) aligned_struct< 2> {}; +template <> struct CUTE_ALIGNAS( 4) aligned_struct< 4> {}; +template <> struct CUTE_ALIGNAS( 8) aligned_struct< 8> {}; +template <> struct CUTE_ALIGNAS( 16) aligned_struct< 16> {}; +template <> struct CUTE_ALIGNAS( 32) aligned_struct< 32> {}; +template <> struct CUTE_ALIGNAS( 64) aligned_struct< 64> {}; +template <> struct CUTE_ALIGNAS(128) aligned_struct<128> {}; +template <> struct CUTE_ALIGNAS(256) aligned_struct<256> {}; + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/container/array.hpp b/csrc/mctlass/include/cute/container/array.hpp new file mode 100644 index 0000000..ff9cbcc --- /dev/null +++ b/csrc/mctlass/include/cute/container/array.hpp @@ -0,0 +1,334 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include + +namespace cute +{ + +template +struct array +{ + using value_type = T; + using size_type = size_t; + using difference_type = ptrdiff_t; + using reference = value_type&; + using const_reference = const value_type&; + using pointer = value_type*; + using const_pointer = const value_type*; + using iterator = pointer; + using const_iterator = const_pointer; + + CUTE_HOST_DEVICE constexpr + reference operator[](size_type pos) + { + return begin()[pos]; + } + + CUTE_HOST_DEVICE constexpr + const_reference operator[](size_type pos) const + { + return begin()[pos]; + } + + CUTE_HOST_DEVICE constexpr + reference front() + { + return *begin(); + } + + CUTE_HOST_DEVICE constexpr + const_reference front() const + { + return *begin(); + } + + CUTE_HOST_DEVICE constexpr + reference back() + { + // return *rbegin(); + return operator[](N-1); + } + + CUTE_HOST_DEVICE constexpr + const_reference back() const + { + // return *rbegin(); + return operator[](N-1); + } + + CUTE_HOST_DEVICE constexpr + T* data() + { + return __elems_; + } + + CUTE_HOST_DEVICE constexpr + T const* data() const + { + return __elems_; + } + + CUTE_HOST_DEVICE constexpr + iterator begin() + { + return data(); + } + + CUTE_HOST_DEVICE constexpr + const_iterator begin() const + { + return data(); + } + + CUTE_HOST_DEVICE constexpr + const_iterator cbegin() + { + return begin(); + } + + CUTE_HOST_DEVICE constexpr + const_iterator cbegin() const + { + return begin(); + } + + CUTE_HOST_DEVICE constexpr + iterator end() + { + return data() + size(); + } + + CUTE_HOST_DEVICE constexpr + const_iterator end() const + { + return data() + size(); + } + + CUTE_HOST_DEVICE constexpr + const_iterator cend() + { + return end(); + } + + CUTE_HOST_DEVICE constexpr + const_iterator cend() const + { + return end(); + } + + CUTE_HOST_DEVICE constexpr + bool empty() const + { + return size() == 0; + } + + CUTE_HOST_DEVICE constexpr + size_type size() const + { + return N; + } + + CUTE_HOST_DEVICE constexpr + size_type max_size() const + { + return size(); + } + + CUTE_HOST_DEVICE constexpr + void fill(const T& value) + { + for (auto& e : *this) { + e = value; + } + } + + CUTE_HOST_DEVICE constexpr + void clear() + { + fill(T(0)); + } + + CUTE_HOST_DEVICE constexpr + void swap(array& other) + { + using CUTE_STL_NAMESPACE::swap; + for (size_type i = 0; i < size(); ++i) { + swap((*this)[i], other[i]); + } + } + + value_type __elems_[N > 0 ? N : 1]; +}; + + +template +CUTE_HOST_DEVICE constexpr +bool operator==(array const& lhs, array const& rhs) +{ + for (size_t i = 0; i < N; ++i) { + if (lhs[i] != rhs[i]) { + return false; + } + } + return true; +} + +template +CUTE_HOST_DEVICE constexpr +void clear(array& a) +{ + a.fill(T(0)); +} + +template +CUTE_HOST_DEVICE constexpr +void fill(array& a, T const& value) +{ + a.fill(value); +} + +template +CUTE_HOST_DEVICE constexpr +void swap(array& a, array& b) +{ + a.swap(b); +} + +} // end cute + + +// +// Specialize tuple-related functionality for cute::array +// + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +namespace cute +{ + +template +CUTE_HOST_DEVICE constexpr +T& get(array& a) +{ + static_assert(I < N, "Index out of range"); + return a[I]; +} + +template +CUTE_HOST_DEVICE constexpr +T const& get(array const& a) +{ + static_assert(I < N, "Index out of range"); + return a[I]; +} + +template +CUTE_HOST_DEVICE constexpr +T&& get(array&& a) +{ + static_assert(I < N, "Index out of range"); + return std::move(a[I]); +} + +} // end namespace cute + +namespace CUTE_STL_NAMESPACE +{ + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> +{ + using type = T; +}; + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> +{ + using type = T; +}; + +} // end namespace CUTE_STL_NAMESPACE + +#ifdef CUTE_STL_NAMESPACE_IS_CUDA_STD +namespace std +{ + +#if defined(__MACACC_RTC__) +template +struct tuple_size; + +template +struct tuple_element; +#endif + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> +{ + using type = T; +}; + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> +{ + using type = T; +}; + +} // end namepsace std +#endif // CUTE_STL_NAMESPACE_IS_CUDA_STD diff --git a/csrc/mctlass/include/cute/container/array_aligned.hpp b/csrc/mctlass/include/cute/container/array_aligned.hpp new file mode 100644 index 0000000..6bf9da3 --- /dev/null +++ b/csrc/mctlass/include/cute/container/array_aligned.hpp @@ -0,0 +1,42 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +#include + +namespace cute +{ + +template +struct CUTE_ALIGNAS(Alignment) array_aligned : cute::array {}; + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/container/array_subbyte.hpp b/csrc/mctlass/include/cute/container/array_subbyte.hpp new file mode 100644 index 0000000..a636196 --- /dev/null +++ b/csrc/mctlass/include/cute/container/array_subbyte.hpp @@ -0,0 +1,633 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Statically sized array of elements that accommodates subbyte trivial types + in a packed storage. +*/ + +#pragma once + +#include + +#include // sizeof_bits +#include + +namespace cute +{ + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Statically sized array for any data type +template +class array_subbyte +{ + public: + + /// Number of total bits in the array + static constexpr int kSizeBits = sizeof_bits::value * N; + + /// Storage type + using Storage = conditional_t<(kSizeBits % 32) == 0, uint32_t, + conditional_t<(kSizeBits % 16) == 0, uint16_t, + uint8_t>>; + + /// Number of logical elements per stored object + static constexpr int kElementsPerStoredItem = sizeof_bits::value / sizeof_bits::value; + + /// Number of storage elements + static constexpr size_t kStorageElements = (N + kElementsPerStoredItem - 1) / kElementsPerStoredItem; + + /// Bitmask for covering one item + static constexpr Storage bit_mask_ = ((Storage(1) << sizeof_bits::value) - 1); + + // + // C++ standard members with reference and iterator types omitted + // + + using value_type = T; + using pointer = value_type*; + using const_pointer = value_type const*; + + using size_type = size_t; + using difference_type = ptrdiff_t; + + // + // References + // + + /// Reference object inserts or extracts sub-byte items + class reference { + /// Pointer to storage element + Storage* ptr_; + + /// Index into elements packed into Storage object + int idx_; + + public: + + /// Default ctor + CUTE_HOST_DEVICE constexpr + reference() : ptr_(nullptr), idx_(0) {} + + /// Ctor + CUTE_HOST_DEVICE constexpr + reference(Storage* ptr, int idx = 0) : ptr_(ptr), idx_(idx) {} + + /// Assignment + CUTE_HOST_DEVICE constexpr + reference& operator=(T x) { + Storage item = (x & bit_mask_); + Storage kUpdateMask = Storage(~(bit_mask_ << (idx_ * sizeof_bits::value))); + *ptr_ = Storage((*ptr_ & kUpdateMask) | (item << (idx_ * sizeof_bits::value))); + return *this; + } + + CUTE_HOST_DEVICE constexpr + T get() const { + if constexpr (is_same::value) { + // Extract to bool -- potentially faster impl + return bool((*ptr_) & (bit_mask_ << (idx_ * sizeof_bits::value))); + } else { + // Extract to T + Storage item = Storage((*ptr_ >> (idx_ * sizeof_bits::value)) & bit_mask_); + return reinterpret_cast(item); + } + } + + /// Extract to type T + CUTE_HOST_DEVICE constexpr + operator T() const { + return get(); + } + }; + + /// Reference object extracts sub-byte items + class const_reference { + + /// Pointer to storage element + Storage const* ptr_; + + /// Index into elements packed into Storage object + int idx_; + + public: + + /// Default ctor + CUTE_HOST_DEVICE constexpr + const_reference(): ptr_(nullptr), idx_(0) { } + + /// Ctor + CUTE_HOST_DEVICE constexpr + const_reference(Storage const* ptr, int idx = 0): ptr_(ptr), idx_(idx) { } + + CUTE_HOST_DEVICE constexpr + const T get() const { + if constexpr (is_same::value) { + // Extract to bool -- potentially faster impl + return bool((*ptr_) & (bit_mask_ << (idx_ * sizeof_bits::value))); + } else { + // Extract to T + Storage item = Storage((*ptr_ >> (idx_ * sizeof_bits::value)) & bit_mask_); + return reinterpret_cast(item); + } + } + + /// Extract to type T + CUTE_HOST_DEVICE constexpr + operator T() const { + return get(); + } + }; + + // + // Iterators + // + + /// Bidirectional iterator over elements + class iterator { + + /// Pointer to storage element + Storage* ptr_; + + /// Index into elements packed into Storage object + int idx_; + + public: + + CUTE_HOST_DEVICE constexpr + iterator(): ptr_(nullptr), idx_(0) { } + + CUTE_HOST_DEVICE constexpr + iterator(Storage* ptr, int idx = 0): ptr_(ptr), idx_(idx) { } + + CUTE_HOST_DEVICE constexpr + iterator& operator++() { + ++idx_; + if (idx_ == kElementsPerStoredItem) { + ++ptr_; + idx_ = 0; + } + return *this; + } + + CUTE_HOST_DEVICE constexpr + iterator& operator--() { + if (idx_) { + --idx_; + } else { + --ptr_; + idx_ = kElementsPerStoredItem - 1; + } + return *this; + } + + CUTE_HOST_DEVICE constexpr + iterator operator++(int) { + iterator ret(*this); + ++(*this); + return ret; + } + + CUTE_HOST_DEVICE constexpr + iterator operator--(int) { + iterator ret(*this); + --(*this); + return ret; + } + + CUTE_HOST_DEVICE constexpr + iterator& operator+=(int k) { + idx_ += k; + ptr_ += idx_ / kElementsPerStoredItem; + idx_ = idx_ % kElementsPerStoredItem; + return *this; + } + + CUTE_HOST_DEVICE constexpr + iterator operator+(int k) const { + return iterator(ptr_,idx_) += k; + } + + CUTE_HOST_DEVICE constexpr + reference operator*() const { + return reference(ptr_, idx_); + } + + CUTE_HOST_DEVICE constexpr + reference operator[](int k) const { + return *(*this + k); + } + + CUTE_HOST_DEVICE constexpr + bool operator==(iterator const& other) const { + return ptr_ == other.ptr_ && idx_ == other.idx_; + } + + CUTE_HOST_DEVICE constexpr + bool operator!=(iterator const& other) const { + return !(*this == other); + } + }; + + /// Bidirectional constant iterator over elements + class const_iterator { + + /// Pointer to storage element + Storage const* ptr_; + + /// Index into elements packed into Storage object + int idx_; + + public: + + CUTE_HOST_DEVICE constexpr + const_iterator(): ptr_(nullptr), idx_(0) { } + + CUTE_HOST_DEVICE constexpr + const_iterator(Storage const* ptr, int idx = 0): ptr_(ptr), idx_(idx) { } + + CUTE_HOST_DEVICE constexpr + const_iterator& operator++() { + ++idx_; + if (idx_ == kElementsPerStoredItem) { + ++ptr_; + idx_ = 0; + } + return *this; + } + + CUTE_HOST_DEVICE constexpr + const_iterator& operator--() { + if (idx_) { + --idx_; + } else { + --ptr_; + idx_ = kElementsPerStoredItem - 1; + } + return *this; + } + + CUTE_HOST_DEVICE constexpr + const_iterator operator++(int) { + iterator ret(*this); + ++idx_; + if (idx_ == kElementsPerStoredItem) { + ++ptr_; + idx_ = 0; + } + return ret; + } + + CUTE_HOST_DEVICE constexpr + const_iterator operator--(int) { + iterator ret(*this); + if (idx_) { + --idx_; + } else { + --ptr_; + idx_ = kElementsPerStoredItem - 1; + } + return ret; + } + + CUTE_HOST_DEVICE constexpr + const_iterator& operator+=(int k) { + idx_ += k; + ptr_ += idx_ / kElementsPerStoredItem; + idx_ = idx_ % kElementsPerStoredItem; + return *this; + } + + CUTE_HOST_DEVICE constexpr + const_iterator operator+(int k) const { + return const_iterator(ptr_,idx_) += k; + } + + CUTE_HOST_DEVICE constexpr + const_reference operator*() const { + return const_reference(ptr_, idx_); + } + + CUTE_HOST_DEVICE constexpr + const_reference operator[](int k) const { + return *(*this + k); + } + + CUTE_HOST_DEVICE constexpr + bool operator==(iterator const& other) const { + return ptr_ == other.ptr_ && idx_ == other.idx_; + } + + CUTE_HOST_DEVICE constexpr + bool operator!=(iterator const& other) const { + return !(*this == other); + } + }; + +private: + + /// Internal storage + Storage storage[kStorageElements]; + +public: + + CUTE_HOST_DEVICE constexpr + array_subbyte() { } + + CUTE_HOST_DEVICE constexpr + array_subbyte(array_subbyte const& x) { + CUTE_UNROLL + for (unsigned i = 0; i < kStorageElements; ++i) { + storage[i] = x.storage[i]; + } + } + + CUTE_HOST_DEVICE constexpr + size_type size() const { + return N; + } + + CUTE_HOST_DEVICE constexpr + size_type max_size() const { + return N; + } + + CUTE_HOST_DEVICE constexpr + bool empty() const { + return !N; + } + + /// Efficient clear method + CUTE_HOST_DEVICE constexpr + void clear() { + CUTE_UNROLL + for (unsigned i = 0; i < kStorageElements; ++i) { + storage[i] = Storage(0); + } + } + + // Efficient fill method + CUTE_HOST_DEVICE constexpr + void fill(T const& value) { + Storage item = (reinterpret_cast(value) & bit_mask_); + + // Reproduce the value over the bits of the storage item + CUTE_UNROLL + for (unsigned s = sizeof_bits::value; s < sizeof_bits::value; s *= 2) { + item |= item << s; + } + + CUTE_UNROLL + for (unsigned i = 0; i < kStorageElements; ++i) { + storage[i] = item; + } + } + + CUTE_HOST_DEVICE constexpr + reference at(size_type pos) { + return reference(storage + pos / kElementsPerStoredItem, pos % kElementsPerStoredItem); + } + + CUTE_HOST_DEVICE constexpr + const_reference at(size_type pos) const { + return const_reference(storage + pos / kElementsPerStoredItem, pos % kElementsPerStoredItem); + } + + CUTE_HOST_DEVICE constexpr + reference operator[](size_type pos) { + return at(pos); + } + + CUTE_HOST_DEVICE constexpr + const_reference operator[](size_type pos) const { + return at(pos); + } + + CUTE_HOST_DEVICE constexpr + reference front() { + return at(0); + } + + CUTE_HOST_DEVICE constexpr + const_reference front() const { + return at(0); + } + + CUTE_HOST_DEVICE constexpr + reference back() { + return reference(storage + kStorageElements - 1, kElementsPerStoredItem - 1); + } + + CUTE_HOST_DEVICE constexpr + const_reference back() const { + return const_reference(storage + kStorageElements - 1, kElementsPerStoredItem - 1); + } + + CUTE_HOST_DEVICE constexpr + pointer data() { + return reinterpret_cast(storage); + } + + CUTE_HOST_DEVICE constexpr + const_pointer data() const { + return reinterpret_cast(storage); + } + + CUTE_HOST_DEVICE constexpr + Storage* raw_data() { + return storage; + } + + CUTE_HOST_DEVICE constexpr + Storage const* raw_data() const { + return storage; + } + + CUTE_HOST_DEVICE constexpr + iterator begin() { + return iterator(storage); + } + + CUTE_HOST_DEVICE constexpr + const_iterator begin() const { + return const_iterator(storage); + } + + CUTE_HOST_DEVICE constexpr + const_iterator cbegin() const { + return begin(); + } + + CUTE_HOST_DEVICE constexpr + iterator end() { + return iterator(storage + N / kElementsPerStoredItem, N % kElementsPerStoredItem); + } + + CUTE_HOST_DEVICE constexpr + const_iterator end() const { + return const_iterator(storage + N / kElementsPerStoredItem, N % kElementsPerStoredItem); + } + + CUTE_HOST_DEVICE constexpr + const_iterator cend() const { + return end(); + } + + // + // Comparison operators + // + +}; + +// +// Operators +// + +template +CUTE_HOST_DEVICE constexpr +void clear(array_subbyte& a) +{ + a.clear(); +} + +template +CUTE_HOST_DEVICE constexpr +void fill(array_subbyte& a, T const& value) +{ + a.fill(value); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cute + +// +// Specialize tuple-related functionality for cute::array_subbyte +// + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +namespace cute +{ + +template +CUTE_HOST_DEVICE constexpr +T& get(array_subbyte& a) +{ + static_assert(I < N, "Index out of range"); + return a[I]; +} + +template +CUTE_HOST_DEVICE constexpr +T const& get(array_subbyte const& a) +{ + static_assert(I < N, "Index out of range"); + return a[I]; +} + +template +CUTE_HOST_DEVICE constexpr +T&& get(array_subbyte&& a) +{ + static_assert(I < N, "Index out of range"); + return std::move(a[I]); +} + +} // end namespace cute + +namespace CUTE_STL_NAMESPACE +{ + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> +{ + using type = T; +}; + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> +{ + using type = T; +}; + +} // end namespace CUTE_STL_NAMESPACE + +#ifdef CUTE_STL_NAMESPACE_IS_CUDA_STD +namespace std +{ + +#if defined(__MACACC_RTC__) +template +struct tuple_size; + +template +struct tuple_element; +#endif + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> +{ + using type = T; +}; + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> +{ + using type = T; +}; + +} // end namespace std +#endif // CUTE_STL_NAMESPACE_IS_CUDA_STD diff --git a/csrc/mctlass/include/cute/container/bit_field.hpp b/csrc/mctlass/include/cute/container/bit_field.hpp new file mode 100644 index 0000000..5398e32 --- /dev/null +++ b/csrc/mctlass/include/cute/container/bit_field.hpp @@ -0,0 +1,131 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Portable bit field that supports byte and word straddling that can + be used in unions to bit-wise define parameters. +*/ + +#pragma once + +#include + +#include // uint_bit_t + +namespace cute +{ + +class dummy_type {}; + +template +struct bit_field +{ + static_assert(0 < NumBits && NumBits <= 64, "bit_fields with more than 64 bits are not supported."); + + // value_type: Use the smallest value type that fits NumBits + static constexpr uint32_t value_type_bits = (NumBits <= 8) ? 8 : + (NumBits <= 16) ? 16 : + (NumBits <= 32) ? 32 : 64; + using value_type = cute::uint_bit_t; + // storage_type: Use the smallest storage_type that avoids boundary crossing + static constexpr uint32_t storage_type_bits = (BitStart / 8 == (BitStart + NumBits - 1) / 8) ? 8 : + (BitStart / 16 == (BitStart + NumBits - 1) / 16) ? 16 : + (BitStart / 32 == (BitStart + NumBits - 1) / 32) ? 32 : 64; + using storage_type = cute::uint_bit_t; + + static_assert(sizeof(OtherValueType) == sizeof(value_type) || is_same::value, + "sizeof(OtherValueType) must be same as sizeof(value_type)."); + + // Number of storage values needed: ceil_div(BitStart + NumBits, storage_type_bits) + static constexpr uint32_t N = (BitStart + NumBits + storage_type_bits - 1) / storage_type_bits; + // Index of storage value for BitStart + static constexpr uint32_t idx = BitStart / storage_type_bits; + // Bit of data_[idx] for BitStart + static constexpr uint32_t bit_lo = BitStart % storage_type_bits; + // Number of bits in data_[idx] used for NumBits if straddling, else 0 + static constexpr uint32_t bit_hi = (idx + 1 < N) ? (storage_type_bits - bit_lo) : 0; + + // NumBits mask + static constexpr value_type mask = (NumBits < 64) ? ((uint64_t(1) << NumBits) - 1) : uint64_t(-1); + // NumBits mask for BitStart + static constexpr storage_type mask_lo = storage_type(mask) << bit_lo; + // NumBits mask for leftover bits in data_[idx+1] if straddling, else 0 + static constexpr storage_type mask_hi = (idx + 1 < N) ? (storage_type(mask) >> bit_hi) : 0; + + storage_type data_[N]; + + // Get value + CUTE_HOST_DEVICE constexpr + value_type get() const { + storage_type result = (data_[idx] & mask_lo) >> bit_lo; + if constexpr (bit_hi) { + result |= (data_[idx+1] & mask_hi) << bit_hi; + } + return static_cast(result); + } + + // Set value + CUTE_HOST_DEVICE constexpr + void set(value_type x) { + storage_type item = static_cast(x & mask); + data_[idx] = static_cast((data_[idx] & ~mask_lo) | (item << bit_lo)); + if constexpr (bit_hi) { + data_[idx+1] = static_cast((data_[idx+1] & ~mask_hi) | (item >> bit_hi)); + } + } + + // Assign value + CUTE_HOST_DEVICE constexpr + bit_field& operator=(value_type x) { + set(x); + return *this; + } + + // Cast to value + CUTE_HOST_DEVICE constexpr + operator value_type () const { + return get(); + } + + // Assign OtherValueType + CUTE_HOST_DEVICE constexpr + bit_field& operator=(OtherValueType x) { + return *this = *reinterpret_cast(&x); + } + + // Cast to OtherValueType + CUTE_HOST_DEVICE constexpr + operator OtherValueType () const { + value_type x = get(); + return *reinterpret_cast(&x); + } +}; + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/container/cuda_types.hpp b/csrc/mctlass/include/cute/container/cuda_types.hpp new file mode 100644 index 0000000..b436f83 --- /dev/null +++ b/csrc/mctlass/include/cute/container/cuda_types.hpp @@ -0,0 +1,186 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + + +#include + +#include +#include + +namespace cute +{ + +// +// dim3 +// + +using dim3 = ::dim3; + +// MSVC doesn't define its C++ version macro to match +// its C++ language version. This means that when +// building with MSVC, dim3 isn't constexpr-friendly. +template +CUTE_HOST_DEVICE +#if ! defined(_MSC_VER) +constexpr +#endif +uint32_t& get(dim3& a) +{ + static_assert(I < 3, "Index out of range"); + if constexpr (I == 0) { + return a.x; + } else if constexpr (I == 1) { + return a.y; + } else if constexpr (I == 2) { + return a.z; + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE +#if ! defined(_MSC_VER) +constexpr +#endif +uint32_t const& get(dim3 const& a) +{ + static_assert(I < 3, "Index out of range"); + if constexpr (I == 0) { + return a.x; + } else if constexpr (I == 1) { + return a.y; + } else if constexpr (I == 2) { + return a.z; + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE +#if ! defined(_MSC_VER) +constexpr +#endif +uint32_t&& get(dim3&& a) +{ + static_assert(I < 3, "Index out of range"); + if constexpr (I == 0) { + return std::move(a.x); + } else if constexpr (I == 1) { + return std::move(a.y); + } else if constexpr (I == 2) { + return std::move(a.z); + } + + CUTE_GCC_UNREACHABLE; +} + +// Specialize cute::tuple-traits for external types +template <> +struct tuple_size + : integral_constant +{}; + +template +struct tuple_element +{ + using type = uint32_t; +}; + +// +// uint3 +// + +using uint3 = ::uint3; + +template +CUTE_HOST_DEVICE constexpr +uint32_t& get(uint3& a) +{ + static_assert(I < 3, "Index out of range"); + if constexpr (I == 0) { + return a.x; + } else if constexpr (I == 1) { + return a.y; + } else if constexpr (I == 2) { + return a.z; + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +uint32_t const& get(uint3 const& a) +{ + static_assert(I < 3, "Index out of range"); + if constexpr (I == 0) { + return a.x; + } else if constexpr (I == 1) { + return a.y; + } else if constexpr (I == 2) { + return a.z; + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +uint32_t&& get(uint3&& a) +{ + static_assert(I < 3, "Index out of range"); + if constexpr (I == 0) { + return std::move(a.x); + } else if constexpr (I == 1) { + return std::move(a.y); + } else if constexpr (I == 2) { + return std::move(a.z); + } + + CUTE_GCC_UNREACHABLE; +} + +// Specialize cute::tuple-traits for external types +template <> +struct tuple_size + : integral_constant +{}; + +template +struct tuple_element +{ + using type = uint32_t; +}; + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/container/tuple.hpp b/csrc/mctlass/include/cute/container/tuple.hpp new file mode 100644 index 0000000..3713d73 --- /dev/null +++ b/csrc/mctlass/include/cute/container/tuple.hpp @@ -0,0 +1,702 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +#include +#include // cute::true_type, cute::false_type +#include + +#include + +//#include // Advanced optimizations + +// +// cute::tuple is like std::tuple, with two differences. +// +// 1. It works on both host and device. +// 2. Its template arguments must be semiregular types. +// +// Semiregular types are default constructible and copyable. +// They include "value types" like int or float, +// but do _not_ include references like int& or float&. +// (See std::tie for an example of a tuple of references.) +// +// This is simplified over the implementations in std::, cuda::std::, and thrust:: by ignoring much of +// the conversion SFINAE, special overloading, and avoiding cvref template types. +// Furthermore, the empty base optimization (EBO) is MORE aggressive by avoiding +// construction calls, and ignoring any need for unique element addresses. +// +// Over standard-conforming tuple implementations, this appears to accelerate compilation times by over 3x. + +namespace cute +{ + +namespace detail +{ + +// EBO stands for "empty base optimization." +// We use this technique to ensure that cute::tuple +// doesn't need to waste space storing any template arguments +// of cute::tuple that have no data (like integral_constant). +// Otherwise, cute::tuple would need to spend at least 1 byte +// for each of its template arguments. +// +// EBO always "holds" a single value of type T. +// N is like an array index that TupleBase uses +// to access the desired tuple element. +template ::value> +struct EBO; + +// Specialization for types T that have no data; +// the "static tuple leaf." Valid T here include +// integral_constant, Int, +// and any other semiregular type +// for which std::is_empty_v is true. +template +struct EBO +{ + CUTE_HOST_DEVICE constexpr + EBO() {} + + CUTE_HOST_DEVICE constexpr + EBO(T const&) {} +}; + +template +CUTE_HOST_DEVICE constexpr T getv(EBO const&) +{ return {}; } + +// Specialization for types T that are not empty; +// the "dynamic tuple leaf." Valid T here include int, +// any other integral or floating-point type, +// or any semiregular type for which std::is_empty_v is false. +template +struct EBO +{ + CUTE_HOST_DEVICE constexpr + EBO() : t_{} {} + + template + CUTE_HOST_DEVICE constexpr + EBO(U const& u) : t_{u} {} + + T t_; +}; + +template +CUTE_HOST_DEVICE constexpr T const& getv(EBO const& x) +{ return x.t_; } + +template +CUTE_HOST_DEVICE constexpr T& getv(EBO& x) +{ return x.t_; } + +template +CUTE_HOST_DEVICE constexpr T&& getv(EBO&& x) +{ return static_cast(x.t_); } + +template +struct TupleBase; + +// Base class of cute::tuple. +// It inherits from EBO for each (i, t) in (I..., T...). +// The actual storage (for nonempty t) lives in the base classes. +// index_sequence is a way to wrap up a sequence of zero or more +// compile-time integer values in a single type. +// We only ever use index_sequence<0, 1, ..., sizeof...(T)> in practice, +// as the type alias TupleBase below indicates. +template +struct TupleBase, T...> + : EBO... +{ + CUTE_HOST_DEVICE constexpr + TupleBase() {} + + template + CUTE_HOST_DEVICE constexpr explicit + TupleBase(U const&... u) + : EBO(u)... {} + + template + CUTE_HOST_DEVICE constexpr + TupleBase(TupleBase, U...> const& u) + : EBO(getv(static_cast const&>(u)))... {} +}; + +} // end namespace detail + +// Attempting to use the following commented-out alias +// in the declaration of `struct tuple` causes MSVC 2022 build errors. +// +//template +//using TupleBase = detail::TupleBase, T...>; + +// This is the actual cute::tuple class. +// The storage (if any) lives in TupleBase's EBO base classes. +// +// Inheriting from the above alias TupleBase +// causes MSVC 2022 build errors when assigning one tuple to another: +// +// illegal member initialization: +// 'TupleBase< /* template arguments */ >' is not a base or member +// +// Not using the alias or any kind of alias fixed the errors. +// In summary: this is verbose as a work-around for MSVC build errors. +template +struct tuple : detail::TupleBase, T...> +{ + CUTE_HOST_DEVICE constexpr + tuple() {} + + template + CUTE_HOST_DEVICE constexpr + tuple(U const&... u) : detail::TupleBase, T...>(u...) {} + + template + CUTE_HOST_DEVICE constexpr + tuple(tuple const& u) + : detail::TupleBase, T...>(static_cast, U...> const&>(u)) {} +}; + +// +// get for cute::tuple (just like std::get for std::tuple) +// + +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +get(tuple const& t) noexcept +{ + static_assert(I < sizeof...(T), "Index out of range"); + return detail::getv(t); +} + +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +get(tuple& t) noexcept +{ + static_assert(I < sizeof...(T), "Index out of range"); + return detail::getv(t); +} + +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +get(tuple&& t) noexcept +{ + static_assert(I < sizeof...(T), "Index out of range"); + return detail::getv(static_cast&&>(t)); +} + +// +// Custom is_tuple trait simply checks the existence of tuple_size +// and assumes std::get(.), std::tuple_element +// +namespace detail { + +template +auto has_tuple_size( T*) -> integral_constant::value>; +auto has_tuple_size(...) -> false_type; + +} // end namespace detail + +template +struct is_tuple : decltype(detail::has_tuple_size((T*)0)) {}; + +// +// make_tuple (value-based implementation) +// + +template +CUTE_HOST_DEVICE constexpr +tuple +make_tuple(T const&... t) +{ + return {t...}; +} + +// +// tuple_cat concatenates multiple cute::tuple into a single cute::tuple, +// just like std::tuple_cat for std::tuple. +// + +#if 0 +// Original implementation + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1, + index_sequence, index_sequence) +{ + return cute::make_tuple(get(t0)..., get(t1)...); +} + +} // end namespace detail + +CUTE_HOST_DEVICE constexpr +tuple<> +tuple_cat() +{ + return {}; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +Tuple const& +tuple_cat(Tuple const& t) +{ + return t; +} + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1) +{ + return detail::tuple_cat(t0, t1, + make_index_sequence::value>{}, + make_index_sequence::value>{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, Ts const&... ts) +{ + return cute::tuple_cat(cute::tuple_cat(t0,t1),t2,ts...); +} +#endif + +#if 1 +// Extended implementation + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1, + index_sequence, index_sequence) +{ + return cute::make_tuple(get(t0)..., get(t1)...); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, + index_sequence, index_sequence, index_sequence) +{ + return cute::make_tuple(get(t0)..., get(t1)..., get(t2)...); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, T3 const& t3, + index_sequence, index_sequence, index_sequence, index_sequence) +{ + return cute::make_tuple(get(t0)..., get(t1)..., get(t2)..., get(t3)...); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, T3 const& t3, T4 const& t4, + index_sequence, index_sequence, index_sequence, index_sequence, index_sequence) +{ + return cute::make_tuple(get(t0)..., get(t1)..., get(t2)..., get(t3)..., get(t4)...); +} + +} // end namespace detail + +CUTE_HOST_DEVICE constexpr +tuple<> +tuple_cat() +{ + return {}; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +Tuple const& +tuple_cat(Tuple const& t) +{ + return t; +} + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1) +{ + return detail::tuple_cat(t0, t1, + make_index_sequence::value>{}, + make_index_sequence::value>{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2) +{ + return detail::tuple_cat(t0, t1, t2, + make_index_sequence::value>{}, + make_index_sequence::value>{}, + make_index_sequence::value>{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, T3 const& t3) +{ + return detail::tuple_cat(t0, t1, t2, t3, + make_index_sequence::value>{}, + make_index_sequence::value>{}, + make_index_sequence::value>{}, + make_index_sequence::value>{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, T3 const& t3, T4 const& t4) +{ + return detail::tuple_cat(t0, t1, t2, t3, t4, + make_index_sequence::value>{}, + make_index_sequence::value>{}, + make_index_sequence::value>{}, + make_index_sequence::value>{}, + make_index_sequence::value>{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, T3 const& t3, T4 const& t4, T5 const& t5, Ts const&... ts) +{ + return cute::tuple_cat(cute::tuple_cat(t0,t1,t2,t3,t4), t5, ts...); +} +#endif + +#if 0 +// Outer-Inner indexing trick to concat all tuples at once + +namespace detail { + +template +struct tuple_cat_helper +{ + static constexpr cute::array ns = {Ns...}; + + static constexpr size_t total_size() { + size_t sum = 0; + for (size_t n : ns) sum += n; + return sum; + } + static constexpr size_t total_size_ = total_size(); + + static constexpr auto values() { + cute::array outer_inner = {}; + + size_t idx = 0; + for (size_t i = 0; i < ns.size(); ++i) { + for (size_t j = 0; j < ns[i]; ++j, ++idx) { + outer_inner[idx][0] = i; + outer_inner[idx][1] = j; + } + } + return outer_inner; + } + static constexpr auto outer_inner_ = values(); + + using total_sequence = make_index_sequence; +}; + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(Tuple const& t, index_sequence) +{ + return cute::make_tuple(get(get(t))...); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1, + index_sequence, index_sequence) +{ + return cute::make_tuple(get(t0)..., get(t1)...); +} + +} // end namespace detail + +CUTE_HOST_DEVICE constexpr +tuple<> +tuple_cat() +{ + return {}; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +Tuple const& +tuple_cat(Tuple const& t) +{ + return t; +} + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(T0 const& t0, T1 const& t1) +{ + return detail::tuple_cat(t0, t1, + make_index_sequence::value>{}, + make_index_sequence::value>{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tuple_cat(Tuples const&... ts) +{ + using Helper = detail::tuple_cat_helper::value...>; + return detail::tuple_cat(cute::make_tuple(ts...), typename Helper::total_sequence{}); +} +#endif + +// +// Equality operators +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +equal_impl(TupleA const& a, TupleB const& b) +{ + if constexpr (I == tuple_size::value) { + return cute::true_type{}; // Terminal: TupleA is exhausted + } else if constexpr (I == tuple_size::value) { + return cute::false_type{}; // Terminal: TupleA is not exhausted, TupleB is exhausted + } else { + return (get(a) == get(b)) && equal_impl(a,b); + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace detail + +template ::value && is_tuple::value)> +CUTE_HOST_DEVICE constexpr +auto +operator==(TupleT const& t, TupleU const& u) +{ + return detail::equal_impl<0>(t, u); +} + +template ::value ^ is_tuple::value)> +CUTE_HOST_DEVICE constexpr +auto +operator==(TupleT const& t, TupleU const& u) +{ + return cute::false_type{}; +} + +template ::value && is_tuple::value)> +CUTE_HOST_DEVICE constexpr +auto +operator!=(TupleT const& t, TupleU const& u) +{ + return !(t == u); +} + +template ::value ^ is_tuple::value)> +CUTE_HOST_DEVICE constexpr +auto +operator!=(TupleT const& t, TupleU const& u) +{ + return cute::true_type{}; +} + +// +// Comparison operators +// + +// +// There are many ways to compare tuple of elements and because CuTe is built +// on parameterizing layouts of coordinates, some comparisons are appropriate +// only in certain cases. +// -- lexicographical comparison [reverse, reflected, revref] +// -- colexicographical comparison [reverse, reflected, revref] +// -- element-wise comparison [any,all] +// This can be very confusing. To avoid errors in selecting the appropriate +// comparison, op<|op<=|op>|op>= are *not* implemented for cute::tuple. +// +// That said, see int_tuple for more explicitly named common comparison ops. +// + +// +// Display utilities +// + +namespace detail { + +template +CUTE_HOST_DEVICE void print_tuple(Tuple const& t, + index_sequence, char s = '(', char e = ')') +{ + using eat = int[]; + using cute::print; + (void) eat {(print(s), 0), + (print(Is == 0 ? "" : ","), print(get(t)), 0)..., + (print(e), 0)}; +} + +#if !defined(__MACACC_RTC__) +template +CUTE_HOST std::ostream& print_tuple_os(std::ostream& os, Tuple const& t, + index_sequence, char s = '(', char e = ')') +{ + using eat = int[]; + (void) eat {(void(os << s), 0), + (void(os << (Is == 0 ? "" : ",") << get(t)), 0)..., + (void(os << e), 0)}; + return os; +} +#endif // !defined(__MACACC_RTC__) + +} // end namespace detail + +template ::value)> +CUTE_HOST_DEVICE void print(Tuple const& t) +{ + return detail::print_tuple(t, make_index_sequence::value>{}); +} + +#if !defined(__MACACC_RTC__) +template ::value)> +CUTE_HOST std::ostream& operator<<(std::ostream& os, Tuple const& t) +{ + return detail::print_tuple_os(os, t, make_index_sequence::value>{}); +} +#endif // !defined(__MACACC_RTC__) + +} // end namespace cute + +namespace CUTE_STL_NAMESPACE +{ + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> + : CUTE_STL_NAMESPACE::tuple_element> +{}; + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> + : CUTE_STL_NAMESPACE::tuple_element> +{}; + +} // end namespace CUTE_STL_NAMESPACE + +// +// std compatibility +// + +#ifdef CUTE_STL_NAMESPACE_IS_CUDA_STD +namespace std +{ + +#if defined(__MACACC_RTC__) +template +struct tuple_size; + +template +struct tuple_element; +#endif + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> + : CUTE_STL_NAMESPACE::tuple_element> +{}; + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> + : CUTE_STL_NAMESPACE::tuple_element> +{}; + +} // end namepsace std +#endif // CUTE_STL_NAMESPACE_IS_CUDA_STD diff --git a/csrc/mctlass/include/cute/container/type_list.hpp b/csrc/mctlass/include/cute/container/type_list.hpp new file mode 100644 index 0000000..a9c28ca --- /dev/null +++ b/csrc/mctlass/include/cute/container/type_list.hpp @@ -0,0 +1,136 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +namespace cute +{ + +template +struct type_c { + using type = T; +}; + +template +struct type_list {}; + +} // end namespace cute + +// +// Specialize tuple-related functionality for cute::type_list +// + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include + +namespace cute +{ + +template +CUTE_HOST_DEVICE constexpr +CUTE_STL_NAMESPACE::tuple_element_t> +get(type_list&) noexcept { + return {}; +} +template +CUTE_HOST_DEVICE constexpr +CUTE_STL_NAMESPACE::tuple_element_t> +get(type_list const& t) noexcept { + return {}; +} + +} // end namespace cute + +namespace CUTE_STL_NAMESPACE +{ + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> + : cute::type_c>::type> +{}; + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> + : cute::type_c>::type> +{}; + +} // end namespace std + +#ifdef CUTE_STL_NAMESPACE_IS_CUDA_STD +namespace std +{ + +#if defined(__MACACC_RTC__) +template +struct tuple_size; + +template +struct tuple_element; +#endif + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> + : cute::type_c>::type> +{}; + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> + : cute::type_c>::type> +{}; + +} // end namespace std +#endif // CUTE_STL_NAMESPACE_IS_CUDA_STD diff --git a/csrc/mctlass/include/cute/int_tuple.hpp b/csrc/mctlass/include/cute/int_tuple.hpp new file mode 100644 index 0000000..b73e2ec --- /dev/null +++ b/csrc/mctlass/include/cute/int_tuple.hpp @@ -0,0 +1,875 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include +#include +#include + +namespace cute +{ + +template +using IntTuple = cute::tuple; + +// Construct an IntTuple with all value-elements +template +CUTE_HOST_DEVICE constexpr +IntTuple +make_int_tuple(Ts const&... t) +{ + return {t...}; +} + +/** if rank(int) == 1, then get<0>(int) should work too + */ +template >::value)> +CUTE_HOST_DEVICE constexpr +decltype(auto) +get(T&& t) noexcept +{ + static_assert(I == 0, "Index out of range"); + return static_cast(t); +} + +/** Custom recursive get for anything that implements get(.) + */ +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +get(Tuple&& t) noexcept +{ + return get(get(static_cast(t))); +} + +// +// rank +// + +template +CUTE_HOST_DEVICE constexpr +auto +rank(IntTuple const& t) +{ + if constexpr (sizeof...(Is) == 0) { + if constexpr (is_tuple::value) { + return Int::value>{}; + } else { + return Int<1>{}; + } + } else { + return rank(get(t)); + } + + CUTE_GCC_UNREACHABLE; +} + +template +using rank_t = decltype(rank(declval())); + +template +static constexpr int rank_v = rank_t::value; + +// +// shape +// + +template +CUTE_HOST_DEVICE constexpr +auto +shape(IntTuple const& s) +{ + if constexpr (is_tuple::value) { + return transform(s, [](auto const& a) { return shape(a); }); + } else { + return s; + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +shape(IntTuple const& s) +{ + if constexpr (is_tuple::value) { + return shape(get(s)); + } else { + return get(shape(s)); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// max +// + +template +CUTE_HOST_DEVICE constexpr +auto +max(T0 const& t0, Ts const&... ts) +{ + if constexpr (is_tuple::value) { + return cute::max(cute::apply(t0, [](auto const&... a){ return cute::max(a...); }), ts...); + } else if constexpr (sizeof...(Ts) == 0) { + return t0; + } else { + return cute::max(t0, cute::max(ts...)); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// min +// + +template +CUTE_HOST_DEVICE constexpr +auto +min(T0 const& t0, Ts const&... ts) +{ + if constexpr (is_tuple::value) { + return cute::min(cute::apply(t0, [](auto const&... a){ return cute::min(a...); }), ts...); + } else if constexpr (sizeof...(Ts) == 0) { + return t0; + } else { + return cute::min(t0, cute::min(ts...)); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// depth +// + +template +CUTE_HOST_DEVICE constexpr +auto +depth(IntTuple const& t) +{ + if constexpr (sizeof...(Is) == 0) { + if constexpr (is_tuple::value) { + return Int<1>{} + cute::apply(t, [](auto const&... v){ return cute::max(depth(v)...); }); + } else { + return Int<0>{}; + } + } else { + return depth(get(t)); + } + + CUTE_GCC_UNREACHABLE; +} + +template +using depth_t = decltype(depth(declval())); + +template +static constexpr int depth_v = depth_t::value; + +// +// product +// + +template +CUTE_HOST_DEVICE constexpr +auto +product(IntTuple const& a) +{ + if constexpr (is_tuple::value) { + return cute::apply(a, [](auto const&... v){ return (Int<1>{} * ... * product(v)); }); + } else { + return a; + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +product_each(Tuple const& t) +{ + return transform(t, [](auto const& x) { return product(x); }); +} + +// Return the product of elements in a mode +template +CUTE_HOST_DEVICE constexpr +auto +size(IntTuple const& a) +{ + if constexpr (sizeof...(Is) == 0) { + return product(a); + } else { + return product(get(a)); + } + + CUTE_GCC_UNREACHABLE; +} + +template +static constexpr int size_v = decltype(size(declval()))::value; + +// +// sum +// + +template +CUTE_HOST_DEVICE constexpr +auto +sum(IntTuple const& a) +{ + if constexpr (is_tuple::value) { + return cute::apply(a, [](auto const&... v){ return (Int<0>{} + ... + sum(v)); }); + } else { + return a; + } + + CUTE_GCC_UNREACHABLE; +} + +// +// inner_product +// + +template +CUTE_HOST_DEVICE constexpr +auto +inner_product(IntTupleA const& a, IntTupleB const& b) +{ + if constexpr (is_tuple::value && is_tuple::value) { + static_assert(tuple_size::value == tuple_size::value, "Mismatched ranks"); + return transform_apply(a, b, [](auto const& x, auto const& y) { return inner_product(x,y); }, + [](auto const&... v) { return (Int<0>{} + ... + v); }); + } else { + return a * b; + } + + CUTE_GCC_UNREACHABLE; +} + +// +// ceil_div +// + +template +CUTE_HOST_DEVICE constexpr +auto +ceil_div(IntTupleA const& a, IntTupleB const& b) +{ + if constexpr (is_tuple::value && is_tuple::value) { + static_assert(tuple_size::value >= tuple_size::value, "Mismatched ranks"); + constexpr int R = tuple_size::value; // Missing ranks in TupleB are implictly 1 + return transform(a, append(b,Int<1>{}), [](auto const& x, auto const& y) { return ceil_div(x,y); }); + } else { + return (a + b - Int<1>{}) / b; + } + + CUTE_GCC_UNREACHABLE; +} + +/** Division for Shapes + */ +template +CUTE_HOST_DEVICE constexpr +auto +shape_div(IntTupleA const& a, IntTupleB const& b) +{ + if constexpr (is_tuple::value) { + if constexpr (is_tuple::value) { // tuple tuple + static_assert(tuple_size::value == tuple_size::value, "Mismatched ranks"); + return transform(a, b, [](auto const& x, auto const& y) { return shape_div(x,y); }); + } else { // tuple int + auto const [result, rest] = fold(a, cute::make_tuple(cute::make_tuple(), b), + [] (auto const& init, auto const& ai) { + return cute::make_tuple(append(get<0>(init), shape_div(ai, get<1>(init))), shape_div(get<1>(init), ai)); + }); + return result; + } + } else { + if constexpr (is_tuple::value) { // int tuple + return shape_div(a, product(b)); + } else { // int int + //assert(a % b == 0 || b % a == 0); + return a / b != 0 ? a / b : signum(a) * signum(b); // divide with rounding away from zero + } + } + + CUTE_GCC_UNREACHABLE; +} + +/** Division for Shapes that are static constants + * @pre t % u == 0 || u % t == 0 + * @result if t % u == 0, then t / u + * if u % t == 0, then signum(t) * signum(u) + */ +template +CUTE_HOST_DEVICE constexpr +constant +shape_div(constant const&, constant const&) +{ + static_assert(t % u == 0 || u % t == 0, "Static shape_div failure"); + return {}; +} + +/** Return a tuple the same profile as A scaled by corresponding elements in B + */ +template +CUTE_HOST_DEVICE constexpr +auto +elem_scale(A const& a, B const& b) +{ + if constexpr (is_tuple::value) { + return transform(a, b, [](auto const& x, auto const& y) { return elem_scale(x,y); }); + } else { + return a * product(b); + } + + CUTE_GCC_UNREACHABLE; +} + +/** Test if two IntTuple have the same profile (hierarchical rank division) + */ +template +CUTE_HOST_DEVICE constexpr +auto +congruent(IntTupleA const& a, IntTupleB const& b) +{ + return bool_constant::value>{}; +} + +template +using is_congruent = decltype(congruent(declval(), declval())); + +/** Test if two IntTuple have the similar profiles up to Shape A (hierarchical rank division) + */ +template +CUTE_HOST_DEVICE constexpr +auto +weakly_congruent(IntTupleA const& a, IntTupleB const& b) +{ + if constexpr (is_tuple::value && is_tuple::value) { + if constexpr (tuple_size::value != tuple_size::value) { + return false_type{}; + } else { + return transform_apply(a, b, [](auto const& x, auto const& y) { return weakly_congruent(x,y); }, + [](auto const&... z) { return (true_type{} && ... && z); }); + } + } else if constexpr (is_integral::value) { + return true_type{}; + } else if constexpr (is_integral::value) { + return false_type{}; + } else { + return weakly_congruent(shape(a), shape(b)); + } + + CUTE_GCC_UNREACHABLE; +} + +template +using is_weakly_congruent = decltype(weakly_congruent(declval(), declval())); + +/** Test if Shape B is compatible with Shape A: + * Any coordinate into A can also be used as a coordinate into B + * A <= B is a partially ordered set of factored shapes + */ +template +CUTE_HOST_DEVICE constexpr +auto +compatible(IntTupleA const& a, IntTupleB const& b) +{ + if constexpr (is_tuple::value && is_tuple::value) { + if constexpr (tuple_size::value != tuple_size::value) { + return false_type{}; + } else { + return transform_apply(a, b, [](auto const& x, auto const& y) { return compatible(x,y); }, + [](auto const&... z) { return (true_type{} && ... && z); }); + } + } else if constexpr (is_integral::value) { + return a == size(b); + } else if constexpr (is_integral::value) { + return false_type{}; + } else { + return compatible(shape(a), shape(b)); + } + + CUTE_GCC_UNREACHABLE; +} + +template +using is_compatible = decltype(compatible(declval(), declval())); + +/** Test if Shape B is weakly compatible with Shape A: + * Shape B divides Shape A at some level of refinement + */ +template +CUTE_HOST_DEVICE constexpr +auto +weakly_compatible(IntTupleA const& a, IntTupleB const& b) +{ + if constexpr (is_tuple::value && is_tuple::value) { + if constexpr (tuple_size::value != tuple_size::value) { + return false_type{}; + } else { + return transform_apply(a, b, [](auto const& x, auto const& y) { return weakly_compatible(x,y); }, + [](auto const&... z) { return (true_type{} && ... && z); }); + } + } else if constexpr (is_integral::value) { + return a % size(b) == Int<0>{}; + } else if constexpr (is_integral::value) { + return false_type{}; + } else { + return weakly_compatible(shape(a), shape(b)); + } + + CUTE_GCC_UNREACHABLE; +} + +template +using is_weakly_compatible = decltype(weakly_compatible(declval(), declval())); + +/** Replace the elements of Tuple B that are paired with an Int<0> with an Int<1> + */ +template +CUTE_HOST_DEVICE constexpr +auto +filter_zeros(IntTupleA const& a, IntTupleB const& b) +{ + if constexpr (is_tuple::value) { + return transform(a, b, [](auto const& x, auto const& y) { return filter_zeros(x,y); }); + } else if constexpr (is_constant<0, IntTupleA>::value) { + return Int<1>{}; + } else { + return b; + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +filter_zeros(Tuple const& t) +{ + return filter_zeros(t, t); +} + +// +// Converters and constructors with arrays and params +// + +/** Make an IntTuple of rank N from an Indexable array. + * Access elements up to a dynamic index n, then use init (requires compatible types) + * Consider cute::take if all indexing is known to be valid + * \code + * std::vector a = {6,3,4}; + * auto tup = make_int_tuple<5>(a, a.size(), 0) // (6,3,4,0,0) + * \endcode + */ +template +CUTE_HOST_DEVICE constexpr +auto +make_int_tuple(Indexable const& t, int n, T const& init) +{ + static_assert(N > 0); + if constexpr (N == 1) { + return 0 < n ? t[0] : init; + } else { + return transform(make_seq{}, [&](auto i) { return i < n ? t[i] : init; }); + } + + CUTE_GCC_UNREACHABLE; +} + +/** Fill the dynamic values of a Tuple with values from another Tuple + * \code + * auto params = make_int_tuple(6,3,4); + * cute::tuple, cute::tuple>, int, Int<2>> result; + * fill_int_tuple_from(result, params); // (_1,(6,3,_3),4,_2) + * \endcode + */ +template +CUTE_HOST_DEVICE constexpr +auto +fill_int_tuple_from(Tuple& result, TupleV const& vals) +{ + return fold(result, vals, [](auto const& init, auto&& r) { + if constexpr (is_static>::value) { // Skip static elements of result + return init; + } else if constexpr (is_tuple>::value) { // Recurse into tuples + return fill_int_tuple_from(r, init); + } else { // Assign and consume arg + static_assert(tuple_size>::value > 0, "Not enough values to fill with!"); + r = get<0>(init); + return remove<0>(init); + } + + CUTE_GCC_UNREACHABLE; + }); +} + +/** Make a "Tuple" by filling in the dynamic values in order from the arguments + * \code + * using result_t = cute::tuple, cute::tuple>, int, Int<2>>; + * auto result = make_int_tuple_from(6,3,4); // (_1,(6,3,_3),4,_2) + * \endcode + */ +template +CUTE_HOST_DEVICE constexpr +Tuple +make_int_tuple_from(Ts const&... ts) +{ + Tuple result = Tuple{}; + fill_int_tuple_from(result, cute::make_tuple(ts...)); + return result; +} + +/** Convert a tuple to a flat homogeneous array of type T + * \code + * auto tup = cute::make_tuple(Int<1>{}, cute::make_tuple(6,3,Int<3>{}),4,Int<2>{}); + * cute::array result = to_array(tup); // [1,6,3,3,4,2] + * \endcode + */ +template +CUTE_HOST_DEVICE constexpr +auto +to_array(IntTuple const& t) +{ + auto flat_t = flatten_to_tuple(t); + constexpr int N = tuple_size::value; + cute::array result; + for_each(make_seq{}, [&] (auto i) { result[i] = get(flat_t); }); + return result; +} + +// +// Comparison operators +// + +// +// There are many ways to compare tuple of elements and because CuTe is built +// on parameterizing layouts of coordinates, some comparisons are appropriate +// only in certain cases. +// -- lexicographical comparison [reverse, reflected, revref] : Correct for coords in RowMajor Layout +// -- colexicographical comparison [reverse, reflected, revref] : Correct for coords in ColMajor Layout +// -- element-wise comparison [any,all] : +// This can be very confusing. To avoid errors in selecting the appropriate +// comparison, op<|op<=|op>|op>= are *not* implemented for cute::tuple. +// +// When actually desiring to order coordinates, the user should map them to +// their indices within the Layout they came from: +// e.g. layoutX(coordA) < layoutX(coordB) +// That said, we implement the three most common ways to compare tuples below. +// These are implemented with slighly more explicit names than op<. +// + +template +CUTE_HOST_DEVICE constexpr +auto +lex_less(IntTupleA const& a, IntTupleB const& b); + +template +CUTE_HOST_DEVICE constexpr +auto +colex_less(IntTupleA const& a, IntTupleB const& b); + +template +CUTE_HOST_DEVICE constexpr +auto +elem_less(IntTupleA const& a, IntTupleB const& b); + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +lex_less_impl(TupleA const& a, TupleB const& b) +{ + if constexpr (I == tuple_size::value) { + return cute::false_type{}; // Terminal: TupleB is exhausted + } else if constexpr (I == tuple_size::value) { + return cute::true_type{}; // Terminal: TupleA is exhausted, TupleB is not exhausted + } else { + return lex_less(get(a), get(b)) || (get(a) == get(b) && lex_less_impl(a,b)); + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +colex_less_impl(TupleA const& a, TupleB const& b) +{ + if constexpr (I == tuple_size::value) { + return cute::false_type{}; // Terminal: TupleB is exhausted + } else if constexpr (I == tuple_size::value) { + return cute::true_type{}; // Terminal: TupleA is exhausted, TupleB is not exhausted + } else { + constexpr size_t A = tuple_size::value - 1 - I; + constexpr size_t B = tuple_size::value - 1 - I; + return colex_less(get(a), get(b)) || (get(a) == get(b) && colex_less_impl(a,b)); + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +elem_less_impl(TupleA const& a, TupleB const& b) +{ + if constexpr (I == tuple_size::value) { + return cute::true_type{}; // Terminal: TupleA is exhausted + } else if constexpr (I == tuple_size::value) { + return cute::false_type{}; // Terminal: TupleA is not exhausted, TupleB is exhausted + } else { + return elem_less(get(a), get(b)) && elem_less_impl(a,b); + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace detail + +// Lexicographical comparison + +template +CUTE_HOST_DEVICE constexpr +auto +lex_less(IntTupleA const& a, IntTupleB const& b) +{ + if constexpr (is_tuple::value && is_tuple::value) { + return detail::lex_less_impl<0>(a, b); + } else { + return a < b; + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +lex_leq(T const& t, U const& u) { + return !lex_less(u, t); +} + +template +CUTE_HOST_DEVICE constexpr +auto +lex_gtr(T const& t, U const& u) { + return lex_less(u, t); +} + +template +CUTE_HOST_DEVICE constexpr +auto +lex_geq(T const& t, U const& u) { + return !lex_less(t, u); +} + +// Colexicographical comparison + +template +CUTE_HOST_DEVICE constexpr +auto +colex_less(IntTupleA const& a, IntTupleB const& b) +{ + if constexpr (is_tuple::value && is_tuple::value) { + return detail::colex_less_impl<0>(a, b); + } else { + return a < b; + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +colex_leq(T const& t, U const& u) { + return !colex_less(u, t); +} + +template +CUTE_HOST_DEVICE constexpr +auto +colex_gtr(T const& t, U const& u) { + return colex_less(u, t); +} + +template +CUTE_HOST_DEVICE constexpr +auto +colex_geq(T const& t, U const& u) { + return !colex_less(t, u); +} + +// Elementwise [all] comparison + +template +CUTE_HOST_DEVICE constexpr +auto +elem_less(IntTupleA const& a, IntTupleB const& b) +{ + if constexpr (is_tuple::value && is_tuple::value) { + return detail::elem_less_impl<0>(a, b); + } else { + return a < b; + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +elem_leq(T const& t, U const& u) { + return !elem_less(u, t); +} + +template +CUTE_HOST_DEVICE constexpr +auto +elem_gtr(T const& t, U const& u) { + return elem_less(u, t); +} + +template +CUTE_HOST_DEVICE constexpr +auto +elem_geq(T const& t, U const& u) { + return !elem_less(t, u); +} + +/** Increment a (dynamic) coord lexicographically within a shape + * \code + * auto shape = make_shape(1,2,make_shape(2,3),3); + * + * int i = 0; + * for (auto coord = repeat_like(shape, 0); back(coord) != back(shape); increment(coord, shape)) { + * std::cout << i++ << ": " << coord << std::endl; + * } + * assert(i == size(shape)); + * \endcode + */ +template +CUTE_HOST_DEVICE constexpr +void +increment(Coord& coord, Shape const& shape); + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +void +increment(Coord& coord, Shape const& shape, seq) +{ + cute::increment(get(coord), get(shape)); + if constexpr (sizeof...(Is) != 0) { + if (back(get(coord)) == back(get(shape))) { + back(get(coord)) = 0; + increment(coord, shape, seq{}); + } + } +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +void +increment(Coord& coord, Shape const& shape) +{ + if constexpr (is_integral::value && is_integral::value) { + ++coord; + } else if constexpr (is_tuple::value && is_tuple::value) { + static_assert(tuple_size::value == tuple_size::value, "Mismatched ranks"); + detail::increment(coord, shape, tuple_seq{}); + } else { + static_assert(sizeof(Coord) == 0, "Invalid parameters"); + } +} + +struct ForwardCoordIteratorSentinal +{}; + +// A forward iterator for a coordinate that starts from zero and goes to shape +template +struct ForwardCoordIterator +{ + static_assert(is_congruent::value); + + CUTE_HOST_DEVICE constexpr + Coord const& operator*() const { return coord; } + + CUTE_HOST_DEVICE constexpr + ForwardCoordIterator& operator++() { increment(coord, shape); return *this; } + + // Sentinal for the end of the implied range + CUTE_HOST_DEVICE constexpr + bool operator< (ForwardCoordIteratorSentinal const&) const { return back(coord) < back(shape); } + CUTE_HOST_DEVICE constexpr + bool operator==(ForwardCoordIteratorSentinal const&) const { return back(coord) == back(shape); } + CUTE_HOST_DEVICE constexpr + bool operator!=(ForwardCoordIteratorSentinal const&) const { return back(coord) != back(shape); } + // NOTE: These are expensive, avoid use + CUTE_HOST_DEVICE constexpr + bool operator< (ForwardCoordIterator const& other) const { return colex_less(coord, other.coord); } + CUTE_HOST_DEVICE constexpr + bool operator==(ForwardCoordIterator const& other) const { return coord == other.coord; } + CUTE_HOST_DEVICE constexpr + bool operator!=(ForwardCoordIterator const& other) const { return coord != other.coord; } + + Coord coord; + Shape const& shape; +}; + +// A forward iterator for a coordinate that starts from zero +template +CUTE_HOST_DEVICE constexpr +auto +make_coord_iterator(Shape const& shape) +{ + auto coord = repeat_like(shape, int(0)); + return ForwardCoordIterator{coord,shape}; +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/layout.hpp b/csrc/mctlass/include/cute/layout.hpp new file mode 100644 index 0000000..1ee99c5 --- /dev/null +++ b/csrc/mctlass/include/cute/layout.hpp @@ -0,0 +1,1687 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include +#include +#include + +namespace cute +{ + +// Aliases + +template +using Shape = IntTuple; + +template +using Stride = IntTuple; + +template +using Step = IntTuple; + +template +using Coord = IntTuple; + +template +CUTE_HOST_DEVICE constexpr +Shape +make_shape(Ts const&... t) { + return {t...}; +} +template +CUTE_HOST_DEVICE constexpr +Stride +make_stride(Ts const&... t) { + return {t...}; +} +template +CUTE_HOST_DEVICE constexpr +Step +make_step(Ts const&... t) { + return {t...}; +} +template +CUTE_HOST_DEVICE constexpr +Coord +make_coord(Ts const&... t) { + return {t...}; +} + + +template > +struct Layout + : private cute::tuple // EBO for static layouts +{ + // Expensive in compilation time... + //static_assert(is_congruent::value, "Shape and Stride must be congruent"); + + // NOTE: This defaults static Shapes/Strides correctly, but not dynamic + CUTE_HOST_DEVICE constexpr + Layout(Shape const& shape = {}, Stride const& stride = {}) + : cute::tuple(shape, stride) + {} + + // + // Accessors + // + + static constexpr int rank = rank_v; + + CUTE_HOST_DEVICE constexpr + decltype(auto) + layout() { + return *this; + } + + CUTE_HOST_DEVICE constexpr + decltype(auto) + layout() const { + return *this; + } + + template + CUTE_HOST_DEVICE constexpr + decltype(auto) + shape() { + return get<0,I...>(static_cast&>(*this)); + } + + template + CUTE_HOST_DEVICE constexpr + decltype(auto) + shape() const { + return get<0,I...>(static_cast const&>(*this)); + } + + template + CUTE_HOST_DEVICE constexpr + decltype(auto) + stride() { + return get<1,I...>(static_cast&>(*this)); + } + + template + CUTE_HOST_DEVICE constexpr + decltype(auto) + stride() const { + return get<1,I...>(static_cast const&>(*this)); + } + + // + // Mappings + // + + // Map a logical coordinate to a linear index (Coord has no Underscore slice operators) + // OR + // Slice the layout and return the sublayout (Coord has an Underscore slice op) + template + CUTE_HOST_DEVICE constexpr + auto + operator()(Coord const& coord) const { + if constexpr (has_underscore::value) { + return slice(coord, *this); + } else { + return crd2idx(coord, shape(), stride()); + } + + CUTE_GCC_UNREACHABLE; + } + + // Convenience function for multi-dimensional coordinates + template + CUTE_HOST_DEVICE constexpr + auto + operator()(Coord0 const& c0, Coord1 const& c1, Coords const&... cs) const { + return operator()(make_coord(c0,c1,cs...)); + } + + // Map a linear index to a hier ND logical coordinate + // NOTE: Dangerous and error-prone + template + CUTE_HOST_DEVICE constexpr + auto + operator[](Int const& linear_idx) const { + static_assert(is_integral::value); + return get_hier_coord(linear_idx); + } + + // + // Compose + // + + template + CUTE_HOST_DEVICE constexpr + auto + compose(OtherLayout const& other) const { + return composition(*this, other); + } + + template + CUTE_HOST_DEVICE constexpr + auto + compose(Layouts const&... layouts) const { + return composition(*this, make_tile(layouts...)); + } + + template + CUTE_HOST_DEVICE constexpr + auto + with_shape(OtherShape const& shape) const { + return composition(*this, make_layout(shape)); + } + + template + CUTE_HOST_DEVICE constexpr + auto + with_shape(Shapes const&... shapes) const { + return composition(*this, make_layout(make_shape(shapes...))); + } + + // + // Tile + // + + template + CUTE_HOST_DEVICE constexpr + auto + tile(OtherLayout const& other) const { + return tiled_divide(*this, other); + } + + template + CUTE_HOST_DEVICE constexpr + auto + tile(Layouts const&... layouts) const { + return tiled_divide(*this, make_tile(layouts...)); + } + + // + // Utility + // + + // + // Index to Coordinate + // + + // NOTE: Only valid for compact layouts + + // Return the (hierarchical) ND logical coordinate corresponding to the linear index + // @post crd2idx(@a result, shape(), stride()) == idx + // @post congruent(@a result, shape()) + template ::value)> + CUTE_HOST_DEVICE constexpr + auto + get_hier_coord(IInt const& idx) const { + return cute::idx2crd(idx, shape(), stride()); + } + + // Return the (flat) ND logical coordinate corresponding to the linear index + // @post crd2idx(@a result, shape(), stride()) == idx + // @post rank(@a result) == rank(shape()) && depth(@a result) == 1 + template ::value)> + CUTE_HOST_DEVICE constexpr + auto + get_flat_coord(IInt const& idx) const { + return cute::crd2crd(this->get_hier_coord(idx), shape(), repeat(Int<1>{})); + } + + // Return the generalized column-major 1D logical coordinate corresponding to the linear index + // @post crd2idx(@a result, shape(), stride()) == idx + // @post is_integral::value + template ::value)> + CUTE_HOST_DEVICE constexpr + auto + get_1d_coord(IInt const& idx) const { + return cute::crd2idx(this->get_hier_coord(idx), shape()); + } + + // + // Coordinate to Coordinate + // + +#if 0 + // Return the (hierarchical) ND logical coordinate corresponding to the linear index + // @post congruent(@a result, shape()) + template + CUTE_HOST_DEVICE constexpr + auto + crd_2_hier_coord(Coord const& crd) const { + return cute::crd2crd(crd, shape(), shape()); + } + + // Return the (flat) ND logical coordinate corresponding to the linear index + // @post rank(@a result) == rank(shape()) && depth(@a result) == 1 + template + CUTE_HOST_DEVICE constexpr + auto + crd_2_flat_coord(Coord const& crd) const { + return cute::crd2crd(crd, shape(), product_each(shape())); + } + + // Return the generalized column-major 1D logical coordinate corresponding to the linear index + // @post is_integral::value + template + CUTE_HOST_DEVICE constexpr + auto + crd_2_1d_coord(Coord const& crd) const { + //return cute::crd2crd(crd, shape(), product(shape())); + return cute::crd2idx(crd, shape()); + } +#endif +}; + +template +struct is_layout : false_type {}; +template +struct is_layout> : true_type {}; + + +template ::value || is_integral::value) && + (is_tuple::value || is_integral::value))> +CUTE_HOST_DEVICE constexpr +auto +make_layout(Shape const& shape, Stride const& stride) +{ + return Layout(shape, stride); +} + +template ::value || is_integral::value)> +CUTE_HOST_DEVICE constexpr +auto +make_layout(Shape const& shape) +{ + return make_layout(shape, compact_col_major(shape)); +} + +// Construct a layout from multiple layouts by +// concatenating each layout as an independent mode +template +CUTE_HOST_DEVICE constexpr +auto +make_layout(Layout const&... layouts) +{ + return make_layout(make_shape (layouts.shape()...), + make_stride(layouts.stride()...)); +} + +// +// Convenience tags for common layouts +// + +template +CUTE_HOST_DEVICE constexpr +auto +make_layout(Shape const& shape, GenColMajor) +{ + return make_layout(shape, compact_col_major(shape)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +make_layout(Shape const& shape, GenRowMajor) +{ + return make_layout(shape, compact_row_major(shape)); +} + +// Follow the same ordering induced by the strides, but make the layout compact +template +CUTE_HOST_DEVICE constexpr +auto +make_ordered_layout(Shape const& shape, Order const& order) +{ + static_assert(is_static::value && is_static::value); + return make_layout(shape, compact_order(shape, order)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +make_ordered_layout(Layout const& layout) +{ + return make_ordered_layout(layout.shape(), layout.stride()); +} + +// Make a layout of the same shape that is either ordered or colmajor depending on staticness +template +CUTE_HOST_DEVICE constexpr +auto +make_layout_like(Layout const& layout) +{ + auto any_zero = any_of(layout.stride(), [](auto d) { return is_constant<0, decltype(d)>{}; }); + if constexpr (any_zero) { + // If there are static-0 strides, then make a col-major layout that keeps those 0s + return make_layout(layout.shape(), + compact_col_major(filter_zeros(layout.stride(), layout.shape()))); + } else + if constexpr (is_static::value && is_static::value) { + // If the layout is fully static, then make a layout that follows the same order as the strides + // Assumes the strides are unique + return make_ordered_layout(layout.shape(), layout.stride()); + } else { + return make_layout(layout.shape()); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Make a layout of the same shape, +// with mode-0 being colmajor then following the mode order in layout +// +template +CUTE_HOST_DEVICE constexpr +auto +make_fragment_like(Layout const& layout) +{ + constexpr int R = Layout::rank; + if constexpr (R > 1 && is_static::value && is_static::value) { + return tiled_product(make_layout(shape<0>(layout)), make_ordered_layout(take<1,R>(layout))); + } else { + return make_layout(layout.shape()); + } + + CUTE_GCC_UNREACHABLE; +} + +template ::value || is_integral::value)> +CUTE_HOST_DEVICE constexpr +auto +make_fragment_like(Shape const& shape) +{ + return make_layout(shape); +} + +// +// Make an identity layout that maps a coordinate to itself +// + +template +CUTE_HOST_DEVICE constexpr +auto +make_identity_layout(Shape const& shape) +{ + return make_layout(shape, make_basis_like(shape)); +} + +// +// Operations to manipulate Layouts like a tuple of pairs +// + +template +CUTE_HOST_DEVICE constexpr +auto +get(Layout const& layout) +{ + // Let the static_asserts in get(shape|stride) catch problems + return make_layout(get(layout.shape()), get(layout.stride())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +take(Layout const& layout) +{ + // Let the static_asserts in take(shape|stride) catch problems + return make_layout(take(layout.shape()), take(layout.stride())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +flatten(Layout const& layout) +{ + return make_layout(flatten(layout.shape()), flatten(layout.stride())); +} + +// +// Utilities +// + +// Return the layout of a mode +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +layout(Layout const& layout) +{ + if constexpr (sizeof...(Is) == 0) { + return layout; + } else { + return get(layout); + } + + CUTE_GCC_UNREACHABLE; +} + +// Return the shape of a mode +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +shape(Layout& layout) +{ + return layout.template shape(); +} + +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +shape(Layout const& layout) +{ + return layout.template shape(); +} + +// Return the stride of a mode +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +stride(Layout& layout) +{ + return layout.template stride(); +} + +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +stride(Layout const& layout) +{ + return layout.template stride(); +} + +// Return the number of elements in a mode +template +CUTE_HOST_DEVICE constexpr +auto +size(Layout const& layout) +{ + return size(shape(layout)); +} + +// Return the number of modes +template +CUTE_HOST_DEVICE constexpr +auto +rank(Layout const& layout) +{ + return rank(shape(layout)); +} + +// Return the depth of the layout +template +CUTE_HOST_DEVICE constexpr +auto +depth(Layout const& layout) +{ + return depth(shape(layout)); +} + +// Return the codomain size of a mode +// @return M smallest integer such that @a sub_layout(c) < M for all c < size(@a sub_layout) +// where sub_layout = get(layout). +template +CUTE_HOST_DEVICE constexpr +auto +cosize(Layout const& layout) +{ + // Protect against negative strides + auto abs_sub_layout = make_layout(shape(layout), + transform_leaf(stride(layout), abs_fn{})); + return abs_sub_layout(size(abs_sub_layout) - Int<1>{}) + Int<1>{}; +} + +template +using cosize_t = decltype(cosize(declval())); + +template +static constexpr int cosize_v = cosize_t::value; + +// Equality +// Return a static or dynamic boolean +template +CUTE_HOST_DEVICE constexpr +auto +operator==(Layout const& layoutA, Layout const& layoutB) +{ + return layoutA.shape() == layoutB.shape() && layoutA.stride() == layoutB.stride(); +} + +// With crd2idx(coord, shape), makes sense to have crd2idx(coord, Layout) as well +template +CUTE_HOST_DEVICE constexpr +auto +crd2idx(Coord const& c, Layout const& layout) +{ + return crd2idx(c, layout.shape(), layout.stride()); +} + +// +// Slice and Dice a layout +// + +template +CUTE_HOST_DEVICE constexpr +auto +slice(Coord const& c, Layout const& layout) +{ + return make_layout(slice(c, layout.shape()), + slice(c, layout.stride())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +slice_and_offset(Coord const& c, Layout const& layout) +{ + return cute::make_tuple(slice(c, layout), crd2idx(c, layout)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +dice(Coord const& c, Layout const& layout) +{ + return make_layout(dice(c, layout.shape()), + dice(c, layout.stride())); +} + +// +// Transform the modes of a layout +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +transform_layout(Tuple const& t, F&& f, seq) +{ + return make_layout(f(get(t))...); +} + +template +CUTE_HOST_DEVICE constexpr +auto +transform_layout(Tuple0 const& t0, Tuple1 const& t1, F&& f, seq, seq, seq) +{ + return make_layout(f(get(t0),get(t1))..., get(t0)..., get(t1)...); +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +transform_layout(Tuple const& t, F&& f) +{ + return detail::transform_layout(t, f, make_seq{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +transform_layout(Tuple0 const& t0, Tuple1 const& t1, F&& f) +{ + constexpr int R0 = decltype(rank(t0))::value; + constexpr int R1 = decltype(rank(t1))::value; + constexpr int R = (R0 < R1) ? R0 : R1; + return detail::transform_layout(t0, t1, f, make_seq{}, make_range{}, make_range{}); +} + +// +// Coalesce and Filter +// + +namespace detail { + +// Look at each element and the front of the stack (in order of priority) +// front(NewLayout) get(Layout) +// s0:d0 _1:d1 => continue +// _1:d0 s1:d1 => replace_front s1:d1 +// s0:s1*d1 s1:d1 => replace_front s0*s1:d1 +// s0:d0 s1:d1 => prepend s1:d1 +// +// @pre OldShape and OldStride are flat +template +CUTE_HOST_DEVICE constexpr +auto +bw_coalesce(OldShape const& old_shape, OldStride const& old_stride, + NewShape const& new_shape, NewStride const& new_stride) +{ + if constexpr (I == -1) { + // Base case, we're done + if constexpr (is_constant<1, NewShape>::value) { + return Layout<_1,_0>{}; + } else { + return Layout{new_shape,new_stride}; + } + } else if constexpr (is_constant<1, decltype(get(old_shape))>::value) { + // shape(layout) == _1, skip it and continue + return bw_coalesce(old_shape, old_stride, new_shape, new_stride); + } else if constexpr (is_constant<1, NewShape>::value) { + // Replace our shape-1 with anything (Can only happen on input new_shape/new_stride) + return bw_coalesce(old_shape, old_stride, get(old_shape), get(old_stride)); + } else if constexpr (is_constant(old_shape) * get(old_stride) == get<0>(new_stride))>::value) { + // Merge modes because the shapes and strides match + return bw_coalesce(old_shape, old_stride, + replace_front(new_shape, get(old_shape) * get<0>(new_shape)), + replace_front(new_stride, get(old_stride))); + } else { + // Can't replace or merge, so prepend a new mode + return bw_coalesce(old_shape, old_stride, + prepend(new_shape, get(old_shape)), + prepend(new_stride, get(old_stride))); + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace detail + +// Combine all the modes that are possible to combine +// Does not respect the profile of the layout, but does preserve total size +template +CUTE_HOST_DEVICE constexpr +auto +coalesce(Layout const& layout) +{ + auto flat_shape = flatten(layout.shape()); + auto flat_stride = flatten(layout.stride()); + + constexpr int R = decltype(rank(flat_shape))::value; + return detail::bw_coalesce(flat_shape, flat_stride, get(flat_shape), get(flat_stride)); +} + +// Apply coalesce at the terminals of trg_profile +template +CUTE_HOST_DEVICE constexpr +auto +coalesce(Layout const& layout, IntTuple const& trg_profile) +{ + if constexpr (is_tuple::value) { + static_assert(tuple_size::value <= Layout::rank); + return transform_layout(layout, trg_profile, [](auto const& l, auto const& t) { return coalesce(l,t); }); + } else { + return coalesce(layout); + } + + CUTE_GCC_UNREACHABLE; +} + +// Replace the modes in layout that have a 0-stride with a 1-size +template +CUTE_HOST_DEVICE constexpr +auto +filter_zeros(Layout const& layout) +{ + return make_layout(filter_zeros(layout.stride(), layout.shape()), layout.stride()); +} + +// Remove all of the 0-strides and 1-sizes +// Return 1-shape if empty +template +CUTE_HOST_DEVICE constexpr +auto +filter(Layout const& layout) +{ + return coalesce(filter_zeros(layout)); +} + +// Apply filter at the terminals of trg_profile +template +CUTE_HOST_DEVICE constexpr +auto +filter(Layout const& layout, IntTuple const& trg_profile) +{ + if constexpr (is_tuple::value) { + static_assert(tuple_size::value <= Layout::rank); + return transform_layout(layout, trg_profile, [](auto const& l, auto const& t) { return filter(l,t); }); + } else { + return filter(layout); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Append, Prepend, Replace +// + +template +CUTE_HOST_DEVICE constexpr +auto +append(Layout const& layout, + Layout const& x = {}) +{ + return make_layout(append(layout.shape(), x.shape()), + append(layout.stride(), x.stride())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +prepend(Layout const& layout, + Layout const& x = {}) +{ + return make_layout(prepend(layout.shape(), x.shape()), + prepend(layout.stride(), x.stride())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +replace(Layout const& layout, + Layout const& x) +{ + return make_layout(replace(layout.shape(), x.shape()), + replace(layout.stride(), x.stride())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +group(Layout const& layout) +{ + return make_layout(group(layout.shape()), + group(layout.stride())); +} + +// +// Composition of two layouts: lhs o rhs +// @post compatible(rhs, result) +// @post result(c) = lhs(rhs(c)) +// for all c in the domain of result +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +composition(Layout const& lhs, + RShape const& rhs_shape, RStride const& rhs_stride) +{ + if constexpr (is_tuple::value) { + // Apply the right-distributivity of Layout composition + return transform_layout(rhs_shape, rhs_stride, [&](auto const& s, auto const& d) { return composition(lhs, s, d); }); + } else + if constexpr (is_scaled_basis::value) { + // Special case for a ScaledBasis stride + return composition(get(lhs), rhs_shape, rhs_stride.value()); + } else + if constexpr (is_integral::value) { + // Integral Rstride (and RShape) + + // NOTE: Should only flatten once for efficiency + auto flat_shape = flatten(lhs.shape()); + [[maybe_unused]] auto flat_stride = flatten(lhs.stride()); + [[maybe_unused]] constexpr int R = rank(flat_shape); + + if constexpr (is_constant<0, RStride>::value) { + // Special case shortcut for any static stride-0 + return Layout{rhs_shape, rhs_stride}; + } else + if constexpr (is_integral::value) { + // Special case shortcut for any integral LShape + auto result_stride = rhs_stride * flat_stride; + return Layout{rhs_shape, result_stride}; + } else + if constexpr (is_constant<1, RStride>::value) { + // Special case shortcut for any static stride-1 + auto result_shape_0 = take<0,R-1>(flat_shape); + + // Mod out the rhs_shape from the lhs.shape() + auto const [result_shape_1, rest_shape] = fold(result_shape_0, cute::make_tuple(cute::make_tuple(), rhs_shape), + [] (auto const& init, auto const& si) { + return cute::make_tuple(append(get<0>(init), cute::min(abs(si), get<1>(init))), shape_div(get<1>(init), abs(si))); + }); + + // Jump into coalesce and append (rest_shape, get(lhs.stride()) + return detail::bw_coalesce(result_shape_1, flat_stride, rest_shape, get(flat_stride)); + } else + { + // General case + auto result_shape_0 = take<0,R-1>(flat_shape); + auto result_stride_0 = take<0,R-1>(flat_stride); + + // Divide out the rhs_stride from the lhs.shape() + auto const [result_shape_1, rest_stride] = fold(result_shape_0, cute::make_tuple(cute::make_tuple(), rhs_stride), + [] (auto const& init, auto const& di) { + return cute::make_tuple(append(get<0>(init), shape_div(di, get<1>(init))), shape_div(get<1>(init), di)); + }); + + // Apply any lhs.shape() changes to the stride + auto result_stride_1 = elem_scale(result_stride_0, shape_div(result_shape_0, result_shape_1)); + + // Mod out the rhs_shape from the lhs.shape() + auto const [result_shape_2, rest_shape] = fold(result_shape_1, cute::make_tuple(cute::make_tuple(), rhs_shape), + [] (auto const& init, auto const& si) { + return cute::make_tuple(append(get<0>(init), cute::min(abs(si), get<1>(init))), shape_div(get<1>(init), abs(si))); + }); + + // Jump into coalesce and append (rest_shape, rest_stride * get(lhs.stride()) + return detail::bw_coalesce(result_shape_2, result_stride_1, rest_shape, rest_stride * get(flat_stride)); + } + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +composition(Layout const& lhs, + Layout const& rhs) +{ + //return detail::composition(flatten(lhs), rhs.shape(), rhs.stride()); + return detail::composition(lhs, rhs.shape(), rhs.stride()); +} + +template +CUTE_HOST_DEVICE constexpr +auto +composition(Layout const& lhs, + IntTuple const& rhs) +{ + if constexpr (is_tuple::value) { + static_assert(tuple_size::value <= Layout::rank); + // Drop any modes of lhs that aren't hit by rhs + return detail::transform_layout(lhs, rhs, [](auto const& l, auto const& r) { return composition(l,r); }, make_seq::value>{}, seq<>{}, seq<>{}); + } else if constexpr (is_underscore::value) { + return lhs; + } else { + return composition(lhs, make_layout(rhs)); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Complement +// +// Build the complement of a layout. +// @post size(@a result) >= @a cosize_hi / size(filter(@a layout))); +// @post For all i in [1,size(@a result)), +// @a result(i) < @a result(i-1) +// For all j in [0, size(@a layout)), +// @a result(i) != @a layout(j) +// + +namespace detail { + +// @pre @a layout has been filtered (flattened and no stride-0 or size-1 modes). +template +CUTE_HOST_DEVICE constexpr +auto +complement(Shape const& shape, Stride const& stride, CoSizeHi const& cosize_hi) +{ + if constexpr (is_constant<0, Stride>::value) { + // Special case for irreducible rank-1 stride-0 layout + return make_layout(cosize_hi); + } else { + // General case + constexpr int R = rank_v; + static_assert(R == 1 || is_static::value, + "Dynamic-stride complement only for rank-1 layouts"); + + // Should just be a sort and a fold... + // Then we could even handle dynamic strides (but they would destroy all static strides) + auto result = fold(make_seq{}, + cute::make_tuple(shape, + stride, + cute::make_tuple(), + cute::make_tuple(Int<1>{})), + [](auto const& init, auto i) + { + auto curr_stride = cute::min(get<1>(init)); + auto curr_idx = find(get<1>(init), curr_stride); + auto curr_shape = get(get<0>(init)); + + return cute::make_tuple(remove(get<0>(init)), // Remove the curr shape + remove(get<1>(init)), // Remove the curr stride + append(get<2>(init), curr_stride / get<3,i>(init)), // new shape = curr_stride / last_stride + append(get<3>(init), curr_shape * curr_stride)); // new stride = curr_shape * curr_stride + }); + + // Append the last shape mode + auto result_stride = get<3>(result); + auto result_shape = append(get<2>(result), get<1,0>(result) / back(result_stride)); // new shape = curr_stride / last_stride + + // Compute the rest_stride + auto rest_stride = get<0,0>(result) * get<1,0>(result); + //return make_layout(append(result_shape, ceil_div(cosize_hi, rest_stride)), append(result_stride, rest_stride)); + // Jump into coalesce and append (ceil_div(cosize_hi, rest_stride), rest_stride) + return detail::bw_coalesce(result_shape, result_stride, ceil_div(cosize_hi, rest_stride), rest_stride); + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +complement(Layout const& layout, CoSizeHi const& cosize_hi) +{ + static_assert(cute::is_integral::value, "Expected integral codomain size in complement."); + auto filter_layout = filter(layout); + return detail::complement(filter_layout.shape(), filter_layout.stride(), cosize_hi); +} + +template +CUTE_HOST_DEVICE constexpr +auto +complement(Layout const& layout) +{ + auto filter_layout = filter(layout); + return detail::complement(filter_layout.shape(), filter_layout.stride(), cosize(filter_layout)); +} + +// +// Right-Inverse and Left-Inverse +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +inverse_seq(Shape const& shape, Stride const& stride, seq) +{ + if constexpr (I == decltype(rank(stride))::value) { + return seq{}; + } else { + //auto next_stride = get(shape) * get(stride); + using next_stride = decltype(get(shape) * get(stride)); // NOTE: WAR for g++-7 + + if constexpr (is_static::value) { + auto next_idx = find_if(stride, [](auto a) { return is_constant{}; }); + return inverse_seq(shape, stride, seq{}); + } else { + return seq{}; + } + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace detail + +// +// Build the right-inverse of a layout +// @pre is_static +// @result A layout @a result such that +// @a layout(@a result(i)) == i for all i < size(@a result) +// @result A layout @a result such that +// composition(@a layout, @a result) is identical to make_layout(shape(result)) +// + +template +CUTE_HOST_DEVICE constexpr +auto +right_inverse(Layout const& layout) +{ + auto flat_layout = coalesce(layout); + auto astride = transform_leaf(flat_layout.stride(), abs_fn{}); + + // Find Int<1>{}, the starting idx, and follow the strides to gen inverse_seq + auto next_I = find_if(astride, [](auto a) { return is_constant<1, decltype(a)>{}; }); + [[maybe_unused]] auto iseq = detail::inverse_seq(flat_layout.shape(), astride, seq<>{}); + + if constexpr (tuple_size::value == 0) { + return Layout<_1,_0>{}; // Empty case, nothing found + } else { + // Generate the corresponding new strides and construct + auto rstride = compact_col_major(flat_layout.shape()); + return make_layout(unwrap(transform(iseq, [&](auto i) { return shape(flat_layout); })), + unwrap(transform(iseq, [&](auto i) { return signum(stride(flat_layout)) * get(rstride); }))); + } + + CUTE_GCC_UNREACHABLE; +} + +CUTE_HOST_DEVICE constexpr +auto +right_inverse(Underscore const& _) +{ + return _; +} + +// +// Build the left-inverse of a layout +// @pre is_static +// @pre @a layout is an injective function +// @result A layout @a result such that +// @a result(@a layout(i)) == i for all i < size(@a layout) +// @result A layout @a result such that +// composition(@a result, @a layout) is identical to make_layout(shape(layout)) +// + +template +CUTE_HOST_DEVICE constexpr +auto +left_inverse(Layout const& layout) +{ + return right_inverse(make_layout(layout, complement(layout))); +} + +CUTE_HOST_DEVICE constexpr +auto +left_inverse(Underscore const& _) +{ + return _; +} + +// +// Max Common Layout +// + +/* Return a layout that points to the maximum number of contiguous elements + * that logically correspond in the layouts of @a a and @a b. This is, + * the elements that could reasonably be "vectorized" in the layouts. + * + * @returns Layout R + * @post For all 0 <= i < size(R), a(R(i)) == i && b(R(i)) == i + */ +template +CUTE_HOST_DEVICE constexpr +auto +max_common_layout(Layout const& a, + Layout const& b) +{ + if constexpr (is_static::value && is_static::value && + is_static::value && is_static::value) + { + Layout inv_b = right_inverse(b); + Layout common = coalesce(composition(a, inv_b)); + + if constexpr (is_constant<1, decltype(stride<0>(common))>::value) { + // Truncate to the size of the contiguous vector (static stride-1 mode) + return composition(inv_b, layout<0>(common)); + } else { + return Layout<_1,_0>{}; + } + } else { + // CASE: One of the layouts is dynamic, can't prove alignment+vectorization is valid + // NOTE: Could weaken if we assume dynamic shapes/strides obey alignment requirements + // (i.e. are large and multiples of the vector) + return Layout<_1,_0>{}; + } + + CUTE_GCC_UNREACHABLE; +} + +/* Return Int such that N is the maximum number of contiguous elements + * that logically correspond in the layouts of @a a and @a b. This is, + * the number of elements that could reasonably be "vectorized" in the layouts. + * + * @returns Int with N >= 1 + * @post For all 0 <= n < N, a(b[n]) == n (NOTE: Problems with negative strides/coords in this post-condition) + */ +template +CUTE_HOST_DEVICE constexpr +auto +max_common_vector(Layout const& a, + Layout const& b) +{ + return size(max_common_layout(a, b)); +} + +// +// Zip +// + +template +CUTE_HOST_DEVICE constexpr +auto +zip(Layout const& layout) +{ + return make_layout(zip(layout.shape()), + zip(layout.stride())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +zip(Layout const& layoutA, + Layout const& layoutB) +{ + return make_layout(zip(layoutA.shape(), layoutB.shape()), + zip(layoutA.stride(), layoutB.stride())); +} + +// +// Tile unzip +// Logical product and logical divide (on layouts) produce rank-2 results by design. +// Follow the profile of @a tile and zip the rank-2 modes located at the terminals into +// their own mode. +// + +template +CUTE_HOST_DEVICE constexpr +auto +tile_unzip(Layout const& layout, + IntTuple const& tile) +{ + return make_layout(zip2_by(layout.shape(), tile), + zip2_by(layout.stride(), tile)); +} + +// +// Logical divide +// + +template +CUTE_HOST_DEVICE constexpr +auto +logical_divide(Layout const& layout, + Layout const& tile) +{ + //CUTE_STATIC_ASSERT_V(size(layout) % size(tile) == Int<0>{}, + // "Tiling does not evenly divide the block"); + // NOTE: With tiles that have stride-0, this doesn't have to be true + + return composition(layout, make_layout(tile, complement(tile, size(layout)))); +} + +template +CUTE_HOST_DEVICE constexpr +auto +logical_divide(Layout const& layout, + IntTuple const& tile) +{ + if constexpr (is_tuple::value) { + static_assert(tuple_size::value <= Layout::rank, "logical_divide: Too many modes in tile."); + return transform_layout(layout, tile, [](auto const& l, auto const& t) { return logical_divide(l,t); }); + } else if constexpr (is_underscore::value) { + return layout; + } else if constexpr (is_integral::value) { + return logical_divide(layout, make_layout(tile)); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Convenience operator +// that produces layouts like ((BLK_A,BLK_B,...),(a,b,...,x,y)) +// by gathering the tile modes and residuals into a rank-2 result. +// + +template +CUTE_HOST_DEVICE constexpr +auto +zipped_divide(Layout const& layout, + Tile const& tile) +{ + return tile_unzip(logical_divide(layout, tile), tile); +} + +// Same as zipped_divide, but unpacks the second mode: ((BLK_A,BLK_B,...),a,b,...,x,y) +template +CUTE_HOST_DEVICE constexpr +auto +tiled_divide(Layout const& layout, + Tile const& tile) +{ + auto div = zipped_divide(layout, tile); + + auto R = rank<1>(div); + return div(_, repeat(_)); +} + +// +// Logical product +// + +template +CUTE_HOST_DEVICE constexpr +auto +logical_product(Layout const& layout, + Layout const& tile) +{ + return make_layout(layout, composition(complement(layout, size(layout)*cosize(tile)), tile)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +logical_product(Layout const& layout, + IntTuple const& tile) +{ + if constexpr (is_tuple::value) { + static_assert(tuple_size::value <= Layout::rank); + return transform_layout(layout, tile, [](auto const& l, auto const& t) { return logical_product(l,t); }); + } else if constexpr (is_underscore::value) { + return layout; + } else if constexpr (is_integral::value) { + return logical_product(layout, make_layout(tile)); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Convenience operator +// that produces layouts like ((BLK_A,BLK_B,...),(a,b,...,x,y)) +// by gathering the block modes and products into a rank-2 result. +// + +template +CUTE_HOST_DEVICE constexpr +auto +zipped_product(Layout const& layout, + Tile const& tile) +{ + return tile_unzip(logical_product(layout, tile), tile); +} + +// Same as zipped_product, but unpacks the second mode: ((BLK_A,BLK_B,...),a,b,...,x,y) +template +CUTE_HOST_DEVICE constexpr +auto +tiled_product(Layout const& layout, + Tile const& tile) +{ + auto div = zipped_product(layout, tile); + + auto R = rank(tile); + return div(_, repeat(_)); +} + +// Attempts to reproduce layout "block" over layout "layout" +// That is, think of every element of "layout" as a "block" +// and return the layout of the resulting structure +template +CUTE_HOST_DEVICE constexpr +auto +blocked_product(Layout const& block, + Layout const& layout) +{ + constexpr int R = cute::max(rank_v, rank_v); + auto padded_block = append(block); + auto padded_layout = append(layout); + + auto result = logical_product(padded_block, padded_layout); + + return coalesce(zip(get<0>(result), get<1>(result)), repeat(Int<1>{})); +} + +template +CUTE_HOST_DEVICE constexpr +auto +raked_product(Layout const& block, + Layout const& layout) +{ + constexpr int R = cute::max(rank_v, rank_v); + auto padded_block = append(block); + auto padded_layout = append(layout); + + auto result = logical_product(padded_block, padded_layout); + + return coalesce(zip(get<1>(result), get<0>(result)), repeat(Int<1>{})); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tile_to_shape(Layout const& layout, + TrgShape const& trg_shape, + ModeOrder const& ord_shape = {}) +{ + CUTE_STATIC_ASSERT_V(rank(layout) <= rank(trg_shape), "Rank of layout must be <= rank of target shape."); + constexpr int R = rank_v; + + auto padded_layout = append(layout); + + auto layout_shape = product_each(padded_layout.shape()); + auto target_shape = product_each(trg_shape); + + // Assert proper division + CUTE_STATIC_ASSERT_V(sum(transform(target_shape, layout_shape, modulus{})) == Int<0>{}, + "Layout shape does not divide the target shape."); + + auto product_shape = shape_div(target_shape, layout_shape); + + return coalesce(blocked_product(padded_layout, make_ordered_layout(product_shape, ord_shape)), product_shape); +} + +// +// Upcast +// For stride-1 mode, divide size by N. Divide all other strides by N. +// + +template +CUTE_HOST_DEVICE constexpr +auto +upcast(Shape const& shape, Stride const& stride) +{ + if constexpr (is_tuple::value) { // tuple stride + return transform_layout(shape, stride, [](auto const& s, auto const& d) { return upcast(s,d); }); + } else if constexpr (is_constant<0, Stride>::value) { // static-0 stride + return Layout{shape,stride}; + } else if constexpr (is_static::value) { // static stride + return make_layout(shape_div(shape, shape_div(Int{}, abs(stride))), + shape_div(stride, Int{})); + } else { // dynamic stride + // assume dynamic strides are larger than N and divisible + // assert(stride % N == 0); + return make_layout(shape, safe_div(stride, Int{})); + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +upcast(Layout const& layout) +{ + return upcast(layout.shape(), layout.stride()); +} + +// +// Downcast +// For stride-1 mode, multiply size by N. Multiply all other strides by N. +// + +template +CUTE_HOST_DEVICE constexpr +auto +downcast(Shape const& shape, Stride const& stride) +{ + if constexpr (is_tuple::value) { + return transform_layout(shape, stride, [](auto const& s, auto const& d) { return downcast(s,d); }); + } else if constexpr (is_constant<1, Stride>::value || is_constant<-1, Stride>::value) { + return make_layout(shape * Int{}, stride); + } else { + return make_layout(shape, stride * Int{}); + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +downcast(Layout const& layout) +{ + CUTE_STATIC_ASSERT(has_int1::value, "Downcast requires adjacent elements"); + return downcast(layout.shape(), layout.stride()); +} + +// +// Recast +// + +template +CUTE_HOST_DEVICE constexpr +auto +recast(Layout const& layout) +{ + if constexpr (sizeof(NewType) == sizeof(OldType)) { + return layout; + } else if constexpr (sizeof(NewType) > sizeof(OldType)) { + static_assert(sizeof(NewType) % sizeof(OldType) == 0, "NewType must be a multiple of OldType"); + return upcast(layout); + } else if constexpr (sizeof(NewType) < sizeof(OldType)) { + static_assert(sizeof(OldType) % sizeof(NewType) == 0, "NewType must be a divisor of OldType"); + return downcast(layout); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Display utilities +// + +template +CUTE_HOST_DEVICE void print(Layout const& layout) +{ + print(layout.shape()); print(":"); print(layout.stride()); +} + +#if !defined(__MACACC_RTC__) +template +CUTE_HOST std::ostream& operator<<(std::ostream& os, Layout const& layout) +{ + return os << shape(layout) << ":" << stride(layout); +} +#endif + +// Generic 2D Layout to console table +template +CUTE_HOST_DEVICE +void +print_layout(Layout const& layout) // (m,n) -> idx +{ + CUTE_STATIC_ASSERT_V(rank(layout) == Int<2>{}); + + int idx_width = num_digits(cosize(layout)) + 2; + const char* delim = "+-----------------------"; + + print(layout); print("\n"); + + // Column indices + print(" "); + for (int n = 0; n < size<1>(layout); ++n) { printf(" %*d ", idx_width-2, n); } + printf("\n"); + + // Print out A m-by-n + for (int m = 0; m < size<0>(layout); ++m) { + // Header + print(" "); + for (int n = 0; n < size<1>(layout); ++n) { printf("%.*s", idx_width+1, delim); } + printf("+\n"); + // Values + printf("%2d ", m); // Row indices + for (int n = 0; n < size<1>(layout); ++n) { printf("| %*d ", idx_width-2, int(layout(m,n))); } + printf("|\n"); + } + // Footer + print(" "); + for (int n = 0; n < size<1>(layout); ++n) { printf("%.*s", idx_width+1, delim); } + printf("+\n"); +} + +// Generic ThrVal 2D Layout to console table +template +CUTE_HOST_DEVICE +void +print_layout(Layout const& layout, ThrID const& thrid) // (m,n) -> (tid,vid) and tid -> thr_idx +{ + CUTE_STATIC_ASSERT_V(rank(layout) == Int<2>{}); + + print(layout); print("\n"); + print(thrid); print("\n"); + + // Print out m-by-n + for (int m = 0; m < size<0>(layout); ++m) { + // Header + for (int n = 0; n < size<1>(layout); ++n) printf("+------"); + printf("+\n"); + // Values + for (int n = 0; n < size<1>(layout); ++n) printf("|%03d-%02d", int(thrid(layout(m,n) % size(thrid))), int(layout(m,n) / size(thrid))); + printf("|\n"); + } + // Footer + for (int n = 0; n < size<1>(layout); ++n) printf("+------"); + printf("+\n"); +} + +// Generic 2D Layout to Latex printer -- B&W 8-value color coding +template +CUTE_HOST_DEVICE +void +print_latex(Layout const& layout) // (m,n) -> idx +{ + CUTE_STATIC_ASSERT_V(rank(layout) == Int<2>{}); + + char const* latex_header = + "\\documentclass[convert]{standalone}\n" + "\\usepackage{tikz}\n\n" + "\\begin{document}\n" + "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},box/.style={rectangle,draw=black,thick,minimum size=1cm,anchor=center,font=\\Large}]\n\n"; + char const* latex_footer = + "\\end{tikzpicture}\n" + "\\end{document}\n"; + + char const* color_map[8] = {"black!00", + "black!40", + "black!20", + "black!60", + "black!10", + "black!50", + "black!30", + "black!70"}; + + // Header + printf("%% Layout: "); print(layout); printf("\n"); + + printf(latex_header); + + // Layout + for (int i = 0; i < size<0>(layout); ++i) { + for (int j = 0; j < size<1>(layout); ++j) { + int idx = layout(i,j); + + printf("\\node[box,fill=%s] at (%d,%d) {%d};\n", + color_map[idx % 8], + i, j, + idx); + } + } + + // Labels + for (int i = 0, j = -1; i < size<0>(layout); ++i) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, i); + } + for (int j = 0, i = -1; j < size<1>(layout); ++j) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, j); + } + + // Footer + printf(latex_footer); +} + +// Generic ThrVal 2D Layout to Latex TIKZ -- 8-value color coded by thread +template +CUTE_HOST_DEVICE +void +print_latex(Layout const& layout, ThrID const& thr) // (m,n) -> (tid,vid) and tid -> thr_idx +{ + CUTE_STATIC_ASSERT_V(rank(layout) == Int<2>{}); + + char const* latex_header = + "\\documentclass[convert]{standalone}\n" + "\\usepackage{tikz}\n\n" + "\\begin{document}\n" + "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},box/.style={rectangle,draw=black,thick,minimum size=1cm,anchor=center}]\n\n"; + char const* latex_footer = + "\\end{tikzpicture}\n" + "\\end{document}\n"; + + char const* color_map[8] = {"{rgb,255:red,175;green,175;blue,255}", + "{rgb,255:red,175;green,255;blue,175}", + "{rgb,255:red,255;green,255;blue,175}", + "{rgb,255:red,255;green,175;blue,175}", + "{rgb,255:red,210;green,210;blue,255}", + "{rgb,255:red,210;green,255;blue,210}", + "{rgb,255:red,255;green,255;blue,210}", + "{rgb,255:red,255;green,210;blue,210}"}; + + // Header + printf("%% layout: "); print(layout); printf("\n"); + printf("%% thrid: "); print(thr); printf("\n\n"); + + printf(latex_header); + + // Layout + for (int i = 0; i < size<0>(layout); ++i) { + for (int j = 0; j < size<1>(layout); ++j) { + int thrid = layout(i,j) % size(thr); + int val_idx = layout(i,j) / size(thr); + int thr_idx = thr(thrid); + + printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n", + color_map[thr_idx % 8], + i, j, + thr_idx, val_idx); + } + } + + // Labels + for (int i = 0, j = -1; i < size<0>(layout); ++i) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, i); + } + for (int j = 0, i = -1; j < size<1>(layout); ++j) { + printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, j); + } + + // Footer + printf(latex_footer); +} + +} // end namespace cute + +// +// Extended Layouts +// + +#include diff --git a/csrc/mctlass/include/cute/numeric/arithmetic_tuple.hpp b/csrc/mctlass/include/cute/numeric/arithmetic_tuple.hpp new file mode 100644 index 0000000..dd59e55 --- /dev/null +++ b/csrc/mctlass/include/cute/numeric/arithmetic_tuple.hpp @@ -0,0 +1,436 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include +#include +#include +#include + +namespace cute +{ + +template +struct ArithmeticTuple : tuple +{ + template + CUTE_HOST_DEVICE constexpr + ArithmeticTuple(ArithmeticTuple const& u) + : tuple(static_cast const&>(u)) {} + + template + CUTE_HOST_DEVICE constexpr + ArithmeticTuple(tuple const& u) + : tuple(u) {} + + template + CUTE_HOST_DEVICE constexpr + ArithmeticTuple(U const&... u) + : tuple(u...) {} +}; + +template +struct is_tuple> : true_type {}; + +template +CUTE_HOST_DEVICE constexpr +auto +make_arithmetic_tuple(T const&... t) { + return ArithmeticTuple(t...); +} + +template +CUTE_HOST_DEVICE constexpr +auto +as_arithmetic_tuple(tuple const& t) { + return ArithmeticTuple(t); +} + +// +// Numeric operators +// + +// Addition +template +CUTE_HOST_DEVICE constexpr +auto +operator+(ArithmeticTuple const& t, ArithmeticTuple const& u) { + constexpr int R = cute::max(int(sizeof...(T)), int(sizeof...(U))); + return transform_apply(append(t,Int<0>{}), append(u,Int<0>{}), plus{}, [](auto const&... a){ return make_arithmetic_tuple(a...); }); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator+(ArithmeticTuple const& t, tuple const& u) { + constexpr int R = cute::max(int(sizeof...(T)), int(sizeof...(U))); + return transform_apply(append(t,Int<0>{}), append(u,Int<0>{}), plus{}, [](auto const&... a){ return make_arithmetic_tuple(a...); }); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator+(tuple const& t, ArithmeticTuple const& u) { + constexpr int R = cute::max(int(sizeof...(T)), int(sizeof...(U))); + return transform_apply(append(t,Int<0>{}), append(u,Int<0>{}), plus{}, [](auto const&... a){ return make_arithmetic_tuple(a...); }); +} + +// +// Special cases +// + +template +CUTE_HOST_DEVICE constexpr +auto +operator+(constant, ArithmeticTuple const& u) { + return u; +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator+(ArithmeticTuple const& t, constant) { + return t; +} + +// +// ArithmeticTupleIterator +// + +template +struct ArithmeticTupleIterator +{ + ArithTuple coord_; + + CUTE_HOST_DEVICE constexpr + ArithmeticTupleIterator() : coord_() {} + CUTE_HOST_DEVICE constexpr + ArithmeticTupleIterator(ArithTuple const& coord) : coord_(coord) {} + + CUTE_HOST_DEVICE constexpr + ArithTuple const& operator*() const { return coord_; } + + template + CUTE_HOST_DEVICE constexpr + auto operator+(Coord const& c) const { + return ArithmeticTupleIterator(coord_ + c); + } + + template + CUTE_HOST_DEVICE constexpr + auto operator[](Coord const& c) const { return *(*this + c); } +}; + +template +CUTE_HOST_DEVICE void print(ArithmeticTupleIterator const& iter) { + printf("ArithTuple"); print(iter.coord_); +} + +// +// ArithmeticTuple "basis" elements +// + +// Abstract value: +// A ScaledBasis is a (at least) rank-N0 ArithmeticTuple: +// (_0,_0,...,T,_0,...) + +template +struct ScaledBasis : private tuple +{ + CUTE_HOST_DEVICE constexpr + ScaledBasis(T const& t = {}) : tuple(t) {} + + CUTE_HOST_DEVICE constexpr + decltype(auto) value() { return get<0>(static_cast &>(*this)); } + CUTE_HOST_DEVICE constexpr + decltype(auto) value() const { return get<0>(static_cast const&>(*this)); } + + CUTE_HOST_DEVICE static constexpr + auto mode() { return Int{}; } +}; + +template +struct is_scaled_basis : false_type {}; +template +struct is_scaled_basis> : true_type {}; + +template +struct is_integral> : true_type {}; + +template +CUTE_HOST_DEVICE constexpr auto +basis_value(T const& e) { + return e; +} + +template +CUTE_HOST_DEVICE constexpr auto +basis_value(ScaledBasis const& e) { + return basis_value(e.value()); +} + +namespace detail { + +template +struct Basis; + +template <> +struct Basis<> { + using type = Int<1>; +}; + +template +struct Basis { + using type = ScaledBasis::type, N>; +}; + +} // end namespace detail + +template +using E = typename detail::Basis::type; + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +as_arithmetic_tuple(T const& t, seq, seq) { + return make_arithmetic_tuple((void(I),Int<0>{})..., t, (void(J),Int<0>{})...); +} + +template +CUTE_HOST_DEVICE constexpr +auto +as_arithmetic_tuple(ArithmeticTuple const& t, seq, seq) { + return make_arithmetic_tuple(get(t)..., (void(J),Int<0>{})...); +} + +} // end namespace detail + +// Turn a ScaledBases into a rank-M ArithmeticTuple +// with N prefix 0s: (_0,_0,...N...,_0,T,_0,...,_0,_0) +template +CUTE_HOST_DEVICE constexpr +auto +as_arithmetic_tuple(ScaledBasis const& t) { + static_assert(M > N, "Mismatched ranks"); + return detail::as_arithmetic_tuple(t.value(), make_seq{}, make_seq{}); +} + +// Turn an ArithmeticTuple into a rank-M ArithmeticTuple +// with postfix 0s: (t0,t1,t2,...,_0,...,_0,_0) +template +CUTE_HOST_DEVICE constexpr +auto +as_arithmetic_tuple(ArithmeticTuple const& t) { + static_assert(M >= sizeof...(T), "Mismatched ranks"); + return detail::as_arithmetic_tuple(t, make_seq{}, make_seq{}); +} + +// Return... +template +CUTE_HOST_DEVICE constexpr +auto +make_basis_like(Shape const& shape) +{ + if constexpr (is_integral::value) { + return Int<1>{}; + } else { + // Generate bases for each rank of shape + return transform(tuple_seq{}, [&](auto I) { + // Generate bases for each rank of shape_i and add an i on front + constexpr int i = decltype(I)::value; // NOTE: nvcc workaround + return transform_leaf(make_basis_like(get(shape)), [&](auto e) { return ScaledBasis{}; }); + }); + } + + CUTE_GCC_UNREACHABLE; +} + +// Equality +template +CUTE_HOST_DEVICE constexpr +auto +operator==(ScaledBasis, Int) { + return false_type{}; +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator==(Int, ScaledBasis) { + return false_type{}; +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator==(ScaledBasis const& t, ScaledBasis const& u) { + return bool_constant{} && t.value() == u.value(); +} + +// Multiplication +template ::value)> +CUTE_HOST_DEVICE constexpr +auto +operator*(A const& a, ScaledBasis const& e) { + return ScaledBasis{a*e.value()}; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +auto +operator*(ScaledBasis const& e, B const& b) { + return ScaledBasis{e.value()*b}; +} + +// Addition +template +CUTE_HOST_DEVICE constexpr +auto +operator+(ScaledBasis const& t, ArithmeticTuple const& u) { + constexpr int R = cute::max(N+1, int(sizeof...(U))); + return as_arithmetic_tuple(t) + as_arithmetic_tuple(u); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator+(ArithmeticTuple const& t, ScaledBasis const& u) { + constexpr int R = cute::max(int(sizeof...(T)), M+1); + return as_arithmetic_tuple(t) + as_arithmetic_tuple(u); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator+(ScaledBasis const& t, ScaledBasis const& u) { + constexpr int R = cute::max(N+1,M+1); + return as_arithmetic_tuple(t) + as_arithmetic_tuple(u); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator+(constant, ScaledBasis const& u) { + return u; +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator+(ScaledBasis const& t, constant) { + return t; +} + +// +// Display utilities +// + +template +CUTE_HOST_DEVICE void print(ScaledBasis const& e) { + print(e.value()); printf("@%d", N); +} + +#if !defined(__MACACC_RTC__) +template +CUTE_HOST std::ostream& operator<<(std::ostream& os, ScaledBasis const& e) { + return os << e.value() << "@" << N; +} +#endif + +} // end namespace cute + + +namespace CUTE_STL_NAMESPACE +{ + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> + : CUTE_STL_NAMESPACE::tuple_element> +{}; + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> + : CUTE_STL_NAMESPACE::tuple_element> +{}; + +} // end namespace CUTE_STL_NAMESPACE + +#ifdef CUTE_STL_NAMESPACE_IS_CUDA_STD +namespace std +{ + +#if defined(__MACACC_RTC__) +template +struct tuple_size; + +template +struct tuple_element; +#endif + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> + : CUTE_STL_NAMESPACE::tuple_element> +{}; + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> + : CUTE_STL_NAMESPACE::tuple_element> +{}; + +} // end namespace std +#endif // CUTE_STL_NAMESPACE_IS_CUDA_STD diff --git a/csrc/mctlass/include/cute/numeric/bfloat.hpp b/csrc/mctlass/include/cute/numeric/bfloat.hpp new file mode 100644 index 0000000..5016871 --- /dev/null +++ b/csrc/mctlass/include/cute/numeric/bfloat.hpp @@ -0,0 +1,53 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +// #include +#include + +namespace cute { + +using mctlass::bfloat16_t; + +// +// Display utilities +// + +#if !defined(__MACACC_RTC__) +CUTE_HOST std::ostream& operator<<(std::ostream& os, bfloat16_t const& v) +{ + return os << float(v); +} +#endif + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/numeric/complex.hpp b/csrc/mctlass/include/cute/numeric/complex.hpp new file mode 100644 index 0000000..9fe2038 --- /dev/null +++ b/csrc/mctlass/include/cute/numeric/complex.hpp @@ -0,0 +1,193 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +//#if defined(__MACA_ARCH__) +//# include +//#else +//# include +//#endif + +// Suppress warnings for code in Thrust headers. + +#if defined(_MSC_VER) + // We check for MSVC first, because MSVC also defines __GNUC__. + // It's common for non-GCC compilers that emulate GCC's behavior + // to define __GNUC__. + // + // thrust/complex.h triggers MSVC's warning on conversion + // from double to float (or const float) ("possible loss of data"). + // MSVC treats this as an error by default (at least with + // MCTLASS's default CMake configuration). +#pragma warning( push ) +#pragma warning( disable : 4244 ) +#elif defined(__GNUC__) + // With GCC + CUDA 11.4, builds show spurious "-Wconversion" + // warnings on line 656 of thrust/detail/type_traits.h. +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wconversion" +#endif + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#if defined(_MSC_VER) +#pragma warning( pop ) +#elif defined(__GNUC__) +#pragma GCC diagnostic pop +#endif + +#include + +namespace cute +{ + +//#if defined(__MACA_ARCH__) +//template +//using complex = cuda::std::complex; +//#else +//template +//using complex = std::complex; +//#endif + +//template +//using complex = thrust::complex; + +#if defined(__MACACC_RTC__) +using cuda::std::complex; +#else +using thrust::complex; +#endif + +template +CUTE_HOST_DEVICE +T real(complex const& z) { + return z.real(); +} + +template +CUTE_HOST_DEVICE +T imag(complex const& z) { + return z.imag(); +} + +template +CUTE_HOST_DEVICE +complex conj(complex const& z) { + return complex(real(z), -imag(z)); +} + +// cute::conj forwards scalars +template +CUTE_HOST_DEVICE +T conj(T z) { + return z; +} + +//CUTE_HOST_DEVICE constexpr +//float conj(float z) { return z; } +//CUTE_HOST_DEVICE constexpr +//double conj(double z) { return z; } + +/// Fused multiply-add for complex numbers +template +CUTE_HOST_DEVICE constexpr +void +fma(complex & d, + complex const& a, + complex const& b, + complex const& c) +{ + d.real(c.real() + a.real() * b.real()); + d.imag(c.imag() + a.real() * b.imag()); + d.real(d.real() - a.imag() * b.imag()); + d.imag(d.imag() + a.imag() * b.real()); +} + +/// Fused multiply-add for triplets +template +CUTE_HOST_DEVICE constexpr +void +fma(complex const& a, + complex const& b, + complex & c) +{ + return fma(c, a, b, c); +} + +/// Used to determine the real-valued underlying type of a numeric type T +template +struct RealType { + using Type = T; +}; + +/// Partial specialization for complex-valued type +template +struct RealType> { + using Type = T; +}; + +////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct is_complex { + static bool const value = false; +}; + +template +struct is_complex> { + static bool const value = true; +}; + +////////////////////////////////////////////////////////////////////////////////////////////////// +// Display utilities + +#if !defined(__MACACC_RTC__) +template +CUTE_HOST std::ostream& operator<<(std::ostream& os, complex const& z) +{ + T _r = z.real(); + T _i = z.imag(); + + if (bool(_i)) { + return os << _r << "+i" << _i; + } else { + return os << _r; + } +} +#endif + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/numeric/float8.hpp b/csrc/mctlass/include/cute/numeric/float8.hpp new file mode 100644 index 0000000..96197d6 --- /dev/null +++ b/csrc/mctlass/include/cute/numeric/float8.hpp @@ -0,0 +1,43 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +// #include +#include + +namespace cute { + +using mctlass::float_e4m3_t; +using mctlass::float_e5m2_t; + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/numeric/half.hpp b/csrc/mctlass/include/cute/numeric/half.hpp new file mode 100644 index 0000000..e883f36 --- /dev/null +++ b/csrc/mctlass/include/cute/numeric/half.hpp @@ -0,0 +1,41 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +// #include +#include + +namespace cute { + +using mctlass::half_t; + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/numeric/int.hpp b/csrc/mctlass/include/cute/numeric/int.hpp new file mode 100644 index 0000000..f15677e --- /dev/null +++ b/csrc/mctlass/include/cute/numeric/int.hpp @@ -0,0 +1,135 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include +#include + +namespace cute +{ + +// +// Signed integers +// + +using int2_t = cute::int2b_t; +using int4_t = cute::int4b_t; +using int8_t = CUTE_STL_NAMESPACE::int8_t; +using int16_t = CUTE_STL_NAMESPACE::int16_t; +using int32_t = CUTE_STL_NAMESPACE::int32_t; +using int64_t = CUTE_STL_NAMESPACE::int64_t; + +template struct int_bit; +template <> struct int_bit< 2> { using type = cute::int2b_t; }; +template <> struct int_bit< 4> { using type = cute::int4b_t; }; +template <> struct int_bit< 8> { using type = int8_t; }; +template <> struct int_bit< 16> { using type = int16_t; }; +template <> struct int_bit< 32> { using type = int32_t; }; +template <> struct int_bit< 64> { using type = int64_t; }; + +template +using int_bit_t = typename int_bit::type; + +template +using int_byte = int_bit<8*N>; + +template +using int_byte_t = typename int_byte::type; + +// +// Unsigned integers +// + +using uint1_t = cute::uint1b_t; +using uint2_t = cute::uint2b_t; +using uint4_t = cute::uint4b_t; +using uint8_t = CUTE_STL_NAMESPACE::uint8_t; +using uint16_t = CUTE_STL_NAMESPACE::uint16_t; +using uint32_t = CUTE_STL_NAMESPACE::uint32_t; +using uint64_t = CUTE_STL_NAMESPACE::uint64_t; +using uint128_t = cute::uint128_t; + +template struct uint_bit; +template <> struct uint_bit< 1> { using type = cute::uint1b_t; }; +template <> struct uint_bit< 2> { using type = cute::uint2b_t; }; +template <> struct uint_bit< 4> { using type = cute::uint4b_t; }; +template <> struct uint_bit< 8> { using type = uint8_t; }; +template <> struct uint_bit< 16> { using type = uint16_t; }; +template <> struct uint_bit< 32> { using type = uint32_t; }; +template <> struct uint_bit< 64> { using type = uint64_t; }; +template <> struct uint_bit<128> { using type = cute::uint128_t; }; + +template +using uint_bit_t = typename uint_bit::type; + +template +using uint_byte = uint_bit<8*N>; + +template +using uint_byte_t = typename uint_byte::type; + +// +// sizeof_bytes +// + +template +struct sizeof_bytes { + static constexpr size_t value = sizeof(T); +}; +template +static constexpr int sizeof_bytes_v = sizeof_bytes::value; + +// +// sizeof_bits +// + +template +struct sizeof_bits { + static constexpr size_t value = sizeof(T) * 8; +}; +template <> +struct sizeof_bits { + static constexpr size_t value = 1; +}; +template +struct sizeof_bits> { + static constexpr size_t value = Bits; +}; +template +static constexpr int sizeof_bits_v = sizeof_bits::value; + +} // namespace cute diff --git a/csrc/mctlass/include/cute/numeric/integer_sequence.hpp b/csrc/mctlass/include/cute/numeric/integer_sequence.hpp new file mode 100644 index 0000000..1e2d659 --- /dev/null +++ b/csrc/mctlass/include/cute/numeric/integer_sequence.hpp @@ -0,0 +1,151 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +#include +#include + +namespace cute +{ + +using CUTE_STL_NAMESPACE::integer_sequence; +using CUTE_STL_NAMESPACE::make_integer_sequence; + +namespace detail { + +template +struct range_impl; + +template +struct range_impl, Begin> { + using type = integer_sequence; +}; + +template +struct reverse_impl; + +template +struct reverse_impl> { + using type = integer_sequence; +}; + +} // end namespace detail + +template +using make_integer_range = typename detail::range_impl< + T, + make_integer_sequence 0) ? (End-Begin) : 0>, + Begin>::type; + +template +using make_integer_sequence_reverse = typename detail::reverse_impl< + make_integer_sequence>::type; + +// +// Common aliases +// + +// int_sequence + +template +using int_sequence = integer_sequence; + +template +using make_int_sequence = make_integer_sequence; + +template +using make_int_rsequence = make_integer_sequence_reverse; + +template +using make_int_range = make_integer_range; + +// index_sequence + +template +using index_sequence = integer_sequence; + +template +using make_index_sequence = make_integer_sequence; + +template +using make_index_rsequence = make_integer_sequence_reverse; + +template +using make_index_range = make_integer_range; + +// +// Shortcuts +// + +template +using seq = int_sequence; + +template +using make_seq = make_int_sequence; + +template +using make_rseq = make_int_rsequence; + +template +using make_range = make_int_range; + +template +using tuple_seq = make_seq>::value>; + +template +using tuple_rseq = make_rseq>::value>; + +// +// Specialize cute::tuple-traits for std::integer_sequence +// + +template +struct tuple_size> + : cute::integral_constant +{}; + +template +struct tuple_element> +{ + constexpr static T idx[sizeof...(Is)] = {Is...}; + using type = cute::integral_constant; +}; + +template +CUTE_HOST_DEVICE constexpr +tuple_element_t> +get(integer_sequence) { + static_assert(I < sizeof...(Ints), "Index out of range"); + return {}; +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/numeric/integer_subbyte.hpp b/csrc/mctlass/include/cute/numeric/integer_subbyte.hpp new file mode 100644 index 0000000..0949bbc --- /dev/null +++ b/csrc/mctlass/include/cute/numeric/integer_subbyte.hpp @@ -0,0 +1,233 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include +#include + +namespace cute { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct integer_subbyte +{ + /// Storage type + using Storage = uint8_t; + + /// Number of bits + static_assert(Bits <= 8*sizeof(Storage), "Require a subbyte of bits in integer_subbyte"); + + /// External type + using xint_t = typename conditional::type; + + /// Bitmask for truncation from larger integers + static constexpr Storage bits_mask_ = Storage((1 << Bits) - 1); + /// Bitmask for the sign bit + static constexpr Storage sign_mask_ = Storage((Signed ? 1 : 0) << (Bits - 1)); + + // + // Data members + // + + Storage storage; + + // + // Methods + // + + /// No operation + CUTE_HOST_DEVICE constexpr + integer_subbyte() {} + + /// Conversion from integer type + CUTE_HOST_DEVICE constexpr + integer_subbyte(int value) // NOTE: Sign extension? + : storage(reinterpret_cast(value) & bits_mask_) {} + + CUTE_HOST_DEVICE constexpr + integer_subbyte(unsigned value) + : storage(reinterpret_cast(value) & bits_mask_) {} + + /// Convert to int or unsigned + CUTE_HOST_DEVICE constexpr + operator xint_t() const { + if (sign_mask_ & storage) { // Sign extend + return xint_t(storage) | ~xint_t(bits_mask_); + } else { + return xint_t(storage); + } + } + + /// Equality + CUTE_HOST_DEVICE constexpr + bool operator==(integer_subbyte const& rhs) const { + return storage == rhs.storage; + } + + /// Inequality + CUTE_HOST_DEVICE constexpr + bool operator!=(integer_subbyte const& rhs) const { + return storage != rhs.storage; + } + + /// Less than or equal + CUTE_HOST_DEVICE constexpr + bool operator<=(integer_subbyte const& rhs) const { + if (sign_mask_ & storage) { + return !(rhs.storage < storage); + } else { + return storage < rhs.storage; + } + } + + /// Less than + CUTE_HOST_DEVICE constexpr + bool operator<(integer_subbyte const& rhs) const { + if (sign_mask_ & storage) { + return !(rhs.storage <= storage); + } else { + return storage < rhs.storage; + } + } + + /// Greater than or equal + CUTE_HOST_DEVICE constexpr + bool operator>=(integer_subbyte const& rhs) const { + return !(*this < rhs); + } + + /// Greater than + CUTE_HOST_DEVICE constexpr + bool operator>(integer_subbyte const& rhs) const { + return !(*this <= rhs); + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// 1-bit unsigned integer type +using uint1b_t = integer_subbyte<1, false>; + +/// 2-bit integer type +using int2b_t = integer_subbyte<2, true>; + +/// 2-bit unsigned integer type +using uint2b_t = integer_subbyte<2, false>; + +/// 4-bit integer type +using int4b_t = integer_subbyte<4, true>; + +/// 4-bit unsigned integer type +using uint4b_t = integer_subbyte<4, false>; + +/// 1-bit binary type +using bin1_t = bool; + +} // namespace cute + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +#if !defined(__MACACC_RTC__) + +#include + +namespace CUTE_STL_NAMESPACE { + +template <> +struct numeric_limits { + CUTE_HOST_DEVICE static constexpr + cute::uint1b_t const lowest() noexcept { return 0; } + CUTE_HOST_DEVICE static constexpr + cute::uint1b_t const min() noexcept { return 0; } + CUTE_HOST_DEVICE static constexpr + cute::uint1b_t const max() noexcept { return 1; } + static constexpr bool is_integer = true; + static constexpr bool is_signed = false; +}; + +template <> +struct numeric_limits { + CUTE_HOST_DEVICE static constexpr + cute::int2b_t lowest() noexcept { return -2; } + CUTE_HOST_DEVICE static constexpr + cute::int2b_t min() noexcept { return -2; } + CUTE_HOST_DEVICE static constexpr + cute::int2b_t max() noexcept { return 1; } + static constexpr bool is_integer = true; + static constexpr bool is_signed = true; +}; + +template <> +struct numeric_limits { + CUTE_HOST_DEVICE static constexpr + cute::uint2b_t const lowest() noexcept { return 0; } + CUTE_HOST_DEVICE static constexpr + cute::uint2b_t const min() noexcept { return 0; } + CUTE_HOST_DEVICE static constexpr + cute::uint2b_t const max() noexcept { return 3; } + static constexpr bool is_integer = true; + static constexpr bool is_signed = false; +}; + +template <> +struct numeric_limits { + CUTE_HOST_DEVICE static constexpr + cute::int4b_t lowest() noexcept { return -8; } + CUTE_HOST_DEVICE static constexpr + cute::int4b_t min() noexcept { return -8; } + CUTE_HOST_DEVICE static constexpr + cute::int4b_t max() noexcept { return 7; } + static constexpr bool is_integer = true; + static constexpr bool is_signed = true; +}; + +template <> +struct numeric_limits { + CUTE_HOST_DEVICE static constexpr + cute::uint4b_t const lowest() noexcept { return 0; } + CUTE_HOST_DEVICE static constexpr + cute::uint4b_t const min() noexcept { return 0; } + CUTE_HOST_DEVICE static constexpr + cute::uint4b_t const max() noexcept { return 15; } + static constexpr bool is_integer = true; + static constexpr bool is_signed = false; +}; + +} // namespace std + +#endif // !defined(__MACACC_RTC__) diff --git a/csrc/mctlass/include/cute/numeric/integral_constant.hpp b/csrc/mctlass/include/cute/numeric/integral_constant.hpp new file mode 100644 index 0000000..9f085fb --- /dev/null +++ b/csrc/mctlass/include/cute/numeric/integral_constant.hpp @@ -0,0 +1,436 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include + +namespace cute +{ + +template +struct constant : CUTE_STL_NAMESPACE::integral_constant { + static constexpr T value = v; + using value_type = T; + using type = constant; + CUTE_HOST_DEVICE constexpr operator value_type() const noexcept { return value; } + CUTE_HOST_DEVICE constexpr value_type operator()() const noexcept { return value; } +}; + +template +using integral_constant = constant; + +template +using bool_constant = constant; + +using true_type = bool_constant; +using false_type = bool_constant; + +// +// Traits +// + +// Use std::is_integral to match built-in integral types (int, int64_t, unsigned, etc) +// Use cute::is_integral to match both built-in integral types AND constant + +template +struct is_integral : bool_constant::value> {}; +template +struct is_integral> : true_type {}; + +// is_static detects if an (abstract) value is defined completely by it's type (no members) + +template +struct is_static : bool_constant::value> {}; + +// is_constant detects if a type is a constant and if v is equal to a value + +template +struct is_constant : false_type {}; +template +struct is_constant > : bool_constant {}; +template +struct is_constant const > : bool_constant {}; +template +struct is_constant const&> : bool_constant {}; +template +struct is_constant &> : bool_constant {}; +template +struct is_constant &&> : bool_constant {}; + +// +// Specializations +// + +template +using Int = constant; + +using _m32 = Int<-32>; +using _m24 = Int<-24>; +using _m16 = Int<-16>; +using _m12 = Int<-12>; +using _m10 = Int<-10>; +using _m9 = Int<-9>; +using _m8 = Int<-8>; +using _m7 = Int<-7>; +using _m6 = Int<-6>; +using _m5 = Int<-5>; +using _m4 = Int<-4>; +using _m3 = Int<-3>; +using _m2 = Int<-2>; +using _m1 = Int<-1>; +using _0 = Int<0>; +using _1 = Int<1>; +using _2 = Int<2>; +using _3 = Int<3>; +using _4 = Int<4>; +using _5 = Int<5>; +using _6 = Int<6>; +using _7 = Int<7>; +using _8 = Int<8>; +using _9 = Int<9>; +using _10 = Int<10>; +using _12 = Int<12>; +using _16 = Int<16>; +using _24 = Int<24>; +using _32 = Int<32>; +using _64 = Int<64>; +using _96 = Int<96>; +using _128 = Int<128>; +using _160 = Int<160>; +using _192 = Int<192>; +using _224 = Int<224>; +using _256 = Int<256>; +using _512 = Int<512>; +using _1024 = Int<1024>; +using _2048 = Int<2048>; +using _4096 = Int<4096>; +using _8192 = Int<8192>; +using _16384 = Int<16384>; +using _32768 = Int<32768>; +using _65536 = Int<65536>; +using _131072 = Int<131072>; +using _262144 = Int<262144>; +using _524288 = Int<524288>; + +/***************/ +/** Operators **/ +/***************/ + +#define CUTE_LEFT_UNARY_OP(OP) \ + template \ + CUTE_HOST_DEVICE constexpr \ + constant \ + operator OP (constant) { \ + return {}; \ + } +#define CUTE_RIGHT_UNARY_OP(OP) \ + template \ + CUTE_HOST_DEVICE constexpr \ + constant \ + operator OP (constant) { \ + return {}; \ + } + +#define CUTE_BINARY_OP(OP) \ + template \ + CUTE_HOST_DEVICE constexpr \ + constant \ + operator OP (constant, constant) { \ + return {}; \ + } + +CUTE_LEFT_UNARY_OP(+); +CUTE_LEFT_UNARY_OP(-); +CUTE_LEFT_UNARY_OP(~); +CUTE_LEFT_UNARY_OP(!); +CUTE_LEFT_UNARY_OP(*); + +CUTE_BINARY_OP( +); +CUTE_BINARY_OP( -); +CUTE_BINARY_OP( *); +CUTE_BINARY_OP( /); +CUTE_BINARY_OP( %); +CUTE_BINARY_OP( &); +CUTE_BINARY_OP( |); +CUTE_BINARY_OP( ^); +CUTE_BINARY_OP(<<); +CUTE_BINARY_OP(>>); + +CUTE_BINARY_OP(&&); +CUTE_BINARY_OP(||); + +CUTE_BINARY_OP(==); +CUTE_BINARY_OP(!=); +CUTE_BINARY_OP( >); +CUTE_BINARY_OP( <); +CUTE_BINARY_OP(>=); +CUTE_BINARY_OP(<=); + +#undef CUTE_BINARY_OP +#undef CUTE_LEFT_UNARY_OP +#undef CUTE_RIGHT_UNARY_OP + +// +// Mixed static-dynamic special cases +// + +template ::value)> +CUTE_HOST_DEVICE constexpr +constant +operator*(constant, U) { + return {}; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +constant +operator*(U, constant) { + return {}; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +constant +operator/(constant, U) { + return {}; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +constant +operator%(U, constant) { + return {}; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +constant +operator%(U, constant) { + return {}; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +constant +operator%(constant, U) { + return {}; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +constant +operator&(constant, U) { + return {}; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +constant +operator&(U, constant) { + return {}; +} + +template ::value && !bool(t))> +CUTE_HOST_DEVICE constexpr +constant +operator&&(constant, U) { + return {}; +} + +template ::value && !bool(t))> +CUTE_HOST_DEVICE constexpr +constant +operator&&(U, constant) { + return {}; +} + +template ::value && bool(t))> +CUTE_HOST_DEVICE constexpr +constant +operator||(constant, U) { + return {}; +} + +template ::value && bool(t))> +CUTE_HOST_DEVICE constexpr +constant +operator||(U, constant) { + return {}; +} + +// +// Named functions from math.hpp +// + +#define CUTE_NAMED_UNARY_FN(OP) \ + template \ + CUTE_HOST_DEVICE constexpr \ + constant \ + OP (constant) { \ + return {}; \ + } + +#define CUTE_NAMED_BINARY_FN(OP) \ + template \ + CUTE_HOST_DEVICE constexpr \ + constant \ + OP (constant, constant) { \ + return {}; \ + } \ + \ + template ::value)> \ + CUTE_HOST_DEVICE constexpr \ + auto \ + OP (constant, U u) { \ + return OP(t,u); \ + } \ + \ + template ::value)> \ + CUTE_HOST_DEVICE constexpr \ + auto \ + OP (T t, constant) { \ + return OP(t,u); \ + } + +CUTE_NAMED_UNARY_FN(abs); +CUTE_NAMED_UNARY_FN(signum); +CUTE_NAMED_UNARY_FN(has_single_bit); + +CUTE_NAMED_BINARY_FN(max); +CUTE_NAMED_BINARY_FN(min); +CUTE_NAMED_BINARY_FN(shiftl); +CUTE_NAMED_BINARY_FN(shiftr); +CUTE_NAMED_BINARY_FN(gcd); +CUTE_NAMED_BINARY_FN(lcm); + +#undef CUTE_NAMED_UNARY_FN +#undef CUTE_NAMED_BINARY_FN + +// +// Other functions +// + +template +CUTE_HOST_DEVICE constexpr +constant +safe_div(constant, constant) { + static_assert(t % u == 0, "Static safe_div requires t % u == 0"); + return {}; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +auto +safe_div(constant, U u) { + return t / u; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +auto +safe_div(T t, constant) { + return t / u; +} + +// cute::true_type prefers standard conversion to std::true_type +// over user-defined conversion to bool +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +conditional_return(true_type, TrueType&& t, FalseType&&) { + return static_cast(t); +} + +// cute::false_type prefers standard conversion to std::false_type +// over user-defined conversion to bool +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +conditional_return(false_type, TrueType&&, FalseType&& f) { + return static_cast(f); +} + +// TrueType and FalseType must have a common type +template +CUTE_HOST_DEVICE constexpr +auto +conditional_return(bool b, TrueType const& t, FalseType const& f) { + return b ? t : f; +} + +// TrueType and FalseType don't require a common type +template +CUTE_HOST_DEVICE constexpr +auto +conditional_return(TrueType const& t, FalseType const& f) { + if constexpr (b) { + return t; + } else { + return f; + } +} + +// +// Display utilities +// + +template +CUTE_HOST_DEVICE void print(integral_constant const&) { + printf("_%d", N); +} + +#if !defined(__MACACC_RTC__) +template +CUTE_HOST std::ostream& operator<<(std::ostream& os, integral_constant const&) { + return os << "_" << N; +} +#endif + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/numeric/math.hpp b/csrc/mctlass/include/cute/numeric/math.hpp new file mode 100644 index 0000000..a90716a --- /dev/null +++ b/csrc/mctlass/include/cute/numeric/math.hpp @@ -0,0 +1,313 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +namespace cute +{ + +// +// Common Operations +// + +template ::value && + is_arithmetic::value)> +CUTE_HOST_DEVICE constexpr +auto +max(T const& t, U const& u) { + return t < u ? u : t; +} + +template ::value && + is_arithmetic::value)> +CUTE_HOST_DEVICE constexpr +auto +min(T const& t, U const& u) { + return t < u ? t : u; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +auto +abs(T const& t) { + if constexpr (is_signed::value) { + return t < T(0) ? -t : t; + } else { + return t; + } + + CUTE_GCC_UNREACHABLE; +} + +// +// C++17 operations +// + +// Greatest common divisor of two integers +template ::value && + CUTE_STL_NAMESPACE::is_integral::value)> +CUTE_HOST_DEVICE constexpr +auto +gcd(T t, U u) { + while (true) { + if (t == 0) { return u; } + u %= t; + if (u == 0) { return t; } + t %= u; + } +} + +// Least common multiple of two integers +template ::value && + CUTE_STL_NAMESPACE::is_integral::value)> +CUTE_HOST_DEVICE constexpr +auto +lcm(T const& t, U const& u) { + return (t / gcd(t,u)) * u; +} + +// +// C++20 operations +// + +// Checks if a number is an integral power of two +template +CUTE_HOST_DEVICE constexpr +bool +has_single_bit(T x) { + return x != 0 && (x & (x - 1)) == 0; +} + +// Smallest number of bits needed to represent the given value +// bit_width( 0b0000 ) = 0 +// bit_width( 0b0001 ) = 1 +// bit_width( 0b0010 ) = 2 +// bit_width( 0b0011 ) = 2 +// bit_width( 0b0100 ) = 3 +// bit_width( 0b0101 ) = 3 +// bit_width( 0b0110 ) = 3 +// bit_width( 0b0111 ) = 3 +template +CUTE_HOST_DEVICE constexpr +T +bit_width(T x) { + static_assert(is_unsigned::value, "Only to be used for unsigned types."); + constexpr int N = (numeric_limits::digits == 64 ? 6 : + (numeric_limits::digits == 32 ? 5 : + (numeric_limits::digits == 16 ? 4 : + (numeric_limits::digits == 8 ? 3 : (assert(false),0))))); + T r = 0; + for (int i = N - 1; i >= 0; --i) { + T shift = (x > ((T(1) << (T(1) << i))-1)) << i; + x >>= shift; + r |= shift; + } + return r + (x != 0); +} + +// Smallest integral power of two not less than the given value +// bit_ceil( 0b00000000 ) = 0b00000001 +// bit_ceil( 0b00000001 ) = 0b00000001 +// bit_ceil( 0b00000010 ) = 0b00000010 +// bit_ceil( 0b00000011 ) = 0b00000100 +// bit_ceil( 0b00000100 ) = 0b00000100 +// bit_ceil( 0b00000101 ) = 0b00001000 +// bit_ceil( 0b00000110 ) = 0b00001000 +// bit_ceil( 0b00000111 ) = 0b00001000 +// bit_ceil( 0b00001000 ) = 0b00001000 +// bit_ceil( 0b00001001 ) = 0b00010000 +template +CUTE_HOST_DEVICE constexpr +T +bit_ceil(T x) { + return x == 0 ? T(1) : (T(1) << bit_width(x - 1)); +} + +// Largest integral power of two not greater than the given value +// bit_floor( 0b00000000 ) = 0b00000000 +// bit_floor( 0b00000001 ) = 0b00000001 +// bit_floor( 0b00000010 ) = 0b00000010 +// bit_floor( 0b00000011 ) = 0b00000010 +// bit_floor( 0b00000100 ) = 0b00000100 +// bit_floor( 0b00000101 ) = 0b00000100 +// bit_floor( 0b00000110 ) = 0b00000100 +// bit_floor( 0b00000111 ) = 0b00000100 +// bit_floor( 0b00001000 ) = 0b00001000 +// bit_floor( 0b00001001 ) = 0b00001000 +template +CUTE_HOST_DEVICE constexpr +T +bit_floor(T x) { + return x == 0 ? 0 : (T(1) << (bit_width(x) - 1)); +} + +template +CUTE_HOST_DEVICE constexpr T rotl(T x, int s); +template +CUTE_HOST_DEVICE constexpr T rotr(T x, int s); + +// Computes the result of circular bitwise left-rotation +template +CUTE_HOST_DEVICE constexpr +T +rotl(T x, int s) { + constexpr int N = numeric_limits::digits; + return s == 0 ? x : s > 0 ? (x << s) | (x >> (N - s)) : rotr(x, -s); +} + +// Computes the result of circular bitwise right-rotation +template +CUTE_HOST_DEVICE constexpr +T +rotr(T x, int s) { + constexpr int N = numeric_limits::digits; + return s == 0 ? x : s > 0 ? (x >> s) | (x << (N - s)) : rotl(x, -s); +} + +// Counts the number of consecutive 0 bits, starting from the most significant bit +// countl_zero( 0b00000000 ) = 8 +// countl_zero( 0b11111111 ) = 0 +// countl_zero( 0b00011100 ) = 3 +template +CUTE_HOST_DEVICE constexpr +T +countl_zero(T x) { + return numeric_limits::digits - bit_width(x); +} + +// Counts the number of consecutive 1 bits, starting from the most significant bit +// countl_one( 0b00000000 ) = 0 +// countl_one( 0b11111111 ) = 8 +// countl_one( 0b11100011 ) = 3 +template +CUTE_HOST_DEVICE constexpr +T +countl_one(T x) { + return countl_zero(~x); +} + +// Counts the number of consecutive 0 bits, starting from the least significant bit +// countr_zero( 0b00000000 ) = 8 +// countr_zero( 0b11111111 ) = 0 +// countr_zero( 0b00011100 ) = 2 +template +CUTE_HOST_DEVICE constexpr +T +countr_zero(T x) { + return x == 0 ? numeric_limits::digits : bit_width(T(x & T(-x))) - 1; // bit_width of the LSB +} + +// Counts the number of consecutive 1 bits, starting from the least significant bit +// countr_one( 0b00000000 ) = 0 +// countr_one( 0b11111111 ) = 8 +// countr_one( 0b11100011 ) = 2 +template +CUTE_HOST_DEVICE constexpr +T +countr_one(T x) { + return countr_zero(~x); +} + +// Counts the number of 1 bits in an unsigned integer +// popcount( 0b00000000 ) = 0 +// popcount( 0b11111111 ) = 8 +// popcount( 0b00011101 ) = 4 +template +CUTE_HOST_DEVICE constexpr +int +popcount(T x) { + int c = 0; + while (x) { + ++c; + x &= x - 1; // clear the least significant bit set + } + return c; +} + +// +// Custom operations +// + +// Computes the result of bitwise left-shift +template +CUTE_HOST_DEVICE constexpr +T +shiftl(T x, int s) { + return s >= 0 ? (x << s) : (x >> -s); +} + +// Computes the result of bitwise right-shift +template +CUTE_HOST_DEVICE constexpr +T +shiftr(T x, int s) { + return s >= 0 ? (x >> s) : (x << -s); +} + +// Returns 1 if x > 0, -1 if x < 0, and 0 if x is zero. +template ::value)> +CUTE_HOST_DEVICE constexpr +int +signum(T const& x) { + return T(0) < x; +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +int +signum(T const& x) { + return (T(0) < x) - (x < T(0)); +} + +// Safe divide +// @pre t % u == 0 +// @result t / u +template ::value && + CUTE_STL_NAMESPACE::is_integral::value)> +CUTE_HOST_DEVICE constexpr +auto +safe_div(T const& t, U const& u) { + //assert(t % u == 0); + return t / u; +} + +} // namespace cute diff --git a/csrc/mctlass/include/cute/numeric/real.hpp b/csrc/mctlass/include/cute/numeric/real.hpp new file mode 100644 index 0000000..d85e304 --- /dev/null +++ b/csrc/mctlass/include/cute/numeric/real.hpp @@ -0,0 +1,56 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +namespace cute +{ + +/// Generic fused multiply-add +template +CUTE_HOST_DEVICE constexpr +void +fma(D& d, A const& a, B const& b, C const& c) +{ + d = a * b + c; +} + +/// Fused multiply-add for triplets +template +CUTE_HOST_DEVICE constexpr +void +fma(A const& a, B const& b, C& c) +{ + return fma(c, a, b, c); +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/numeric/tfloat.hpp b/csrc/mctlass/include/cute/numeric/tfloat.hpp new file mode 100644 index 0000000..39956b2 --- /dev/null +++ b/csrc/mctlass/include/cute/numeric/tfloat.hpp @@ -0,0 +1,53 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +// #include +#include + +namespace cute { + +using mctlass::tfloat32_t; + +// +// Display utilities +// + +#if !defined(__MACACC_RTC__) +CUTE_HOST std::ostream& operator<<(std::ostream& os, tfloat32_t const& v) +{ + return os << float(v); +} +#endif + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/numeric/uint128.hpp b/csrc/mctlass/include/cute/numeric/uint128.hpp new file mode 100644 index 0000000..af8a36e --- /dev/null +++ b/csrc/mctlass/include/cute/numeric/uint128.hpp @@ -0,0 +1,259 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#include +#include +#include +#include +#endif + +#include + +/// Optionally enable GCC's built-in type +#if defined(__x86_64) && !defined(__MACA_ARCH__) +# if defined(__GNUC__) && 0 +# define CUTE_UINT128_NATIVE +# elif defined(_MSC_VER) +# define CUTE_INT128_ARITHMETIC +# include +# endif +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace cute { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///! Unsigned 128b integer type +struct alignas(16) uint128_t +{ + /// Size of one part of the uint's storage in bits + static constexpr int storage_bits_ = 64; + + struct hilo + { + uint64_t lo; + uint64_t hi; + }; + + // Use a union to store either low and high parts or, if present, a built-in 128b integer type. + union + { + struct hilo hilo_; + +#if defined(CUTE_UINT128_NATIVE) + unsigned __int128 native; +#endif // defined(CUTE_UINT128_NATIVE) + }; + + // + // Methods + // + + /// Default ctor + CUTE_HOST_DEVICE constexpr + uint128_t() : hilo_{0, 0} {} + + /// Constructor from uint64 + CUTE_HOST_DEVICE constexpr + uint128_t(uint64_t lo_) : hilo_{lo_, 0} {} + + /// Constructor from two 64b unsigned integers + CUTE_HOST_DEVICE constexpr + uint128_t(uint64_t lo_, uint64_t hi_) : hilo_{lo_, hi_} {} + + /// Optional constructor from native value +#if defined(CUTE_UINT128_NATIVE) + uint128_t(unsigned __int128 value) : native(value) { } +#endif + + /// Lossily cast to uint64 + CUTE_HOST_DEVICE constexpr + explicit operator uint64_t() const + { + return hilo_.lo; + } + + template + CUTE_HOST_DEVICE constexpr + static void exception() + { + //static_assert(sizeof(Dummy) == 0, "Not implemented exception!"); + //abort(); + //printf("uint128 not implemented!\n"); + } + + /// Add + CUTE_HOST_DEVICE constexpr + uint128_t operator+(uint128_t const& rhs) const + { + uint128_t y; +#if defined(CUTE_UINT128_NATIVE) + y.native = native + rhs.native; +#else + y.hilo_.lo = hilo_.lo + rhs.hilo_.lo; + y.hilo_.hi = hilo_.hi + rhs.hilo_.hi + (!y.hilo_.lo && (rhs.hilo_.lo)); +#endif + return y; + } + + /// Subtract + CUTE_HOST_DEVICE constexpr + uint128_t operator-(uint128_t const& rhs) const + { + uint128_t y; +#if defined(CUTE_UINT128_NATIVE) + y.native = native - rhs.native; +#else + y.hilo_.lo = hilo_.lo - rhs.hilo_.lo; + y.hilo_.hi = hilo_.hi - rhs.hilo_.hi - (rhs.hilo_.lo && y.hilo_.lo > hilo_.lo); +#endif + return y; + } + + /// Multiply by unsigned 64b integer yielding 128b integer + CUTE_HOST_DEVICE constexpr + uint128_t operator*(uint64_t const& rhs) const + { + uint128_t y; +#if defined(CUTE_UINT128_NATIVE) + y.native = native * rhs; +#elif defined(CUTE_INT128_ARITHMETIC) + // Multiply by the low part + y.hilo_.lo = _umul128(hilo_.lo, rhs, &y.hilo_.hi); + + // Add the high part and ignore the overflow + uint64_t overflow; + y.hilo_.hi += _umul128(hilo_.hi, rhs, &overflow); +#else + exception(); +#endif + return y; + } + + /// Divide 128b operation by 64b operation yielding a 64b quotient + CUTE_HOST_DEVICE constexpr + uint64_t operator/(uint64_t const& divisor) const + { + uint64_t quotient = 0; +#if defined(CUTE_UINT128_NATIVE) + quotient = uint64_t(native / divisor); +#elif defined(CUTE_INT128_ARITHMETIC) + // implemented using MSVC's arithmetic intrinsics + uint64_t remainder = 0; + quotient = _udiv128(hilo_.hi, hilo_.lo, divisor, &remainder); +#else + exception(); +#endif + return quotient; + } + + /// Divide 128b operation by 64b operation yielding a 64b quotient + CUTE_HOST_DEVICE constexpr + uint64_t operator%(uint64_t const& divisor) const + { + uint64_t remainder = 0; +#if defined(CUTE_UINT128_NATIVE) + remainder = uint64_t(native % divisor); +#elif defined(CUTE_INT128_ARITHMETIC) + // implemented using MSVC's arithmetic intrinsics + (void)_udiv128(hilo_.hi, hilo_.lo, divisor, &remainder); +#else + exception(); +#endif + return remainder; + } + + /// Computes the quotient and remainder in a single method. + CUTE_HOST_DEVICE constexpr + uint64_t divmod(uint64_t &remainder, uint64_t divisor) const + { + uint64_t quotient = 0; +#if defined(CUTE_UINT128_NATIVE) + quotient = uint64_t(native / divisor); + remainder = uint64_t(native % divisor); +#elif defined(CUTE_INT128_ARITHMETIC) + // implemented using MSVC's arithmetic intrinsics + quotient = _udiv128(hilo_.hi, hilo_.lo, divisor, &remainder); +#else + exception(); +#endif + return quotient; + } + + /// Left-shifts a 128b unsigned integer + CUTE_HOST_DEVICE constexpr + uint128_t operator<<(int sh) const + { + if (sh == 0) { + return *this; + } + else if (sh >= storage_bits_) { + return uint128_t(0, hilo_.lo << (sh - storage_bits_)); + } + else { + return uint128_t( + (hilo_.lo << sh), + (hilo_.hi << sh) | uint64_t(hilo_.lo >> (storage_bits_ - sh)) + ); + } + } + + /// Right-shifts a 128b unsigned integer + CUTE_HOST_DEVICE constexpr + uint128_t operator>>(int sh) const + { + if (sh == 0) { + return *this; + } + else if (sh >= storage_bits_) { + return uint128_t((hilo_.hi >> (sh - storage_bits_)), 0); + } + else { + return uint128_t( + (hilo_.lo >> sh) | (hilo_.hi << (storage_bits_ - sh)), + (hilo_.hi >> sh) + ); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace cute + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/cute/pointer.hpp b/csrc/mctlass/include/cute/pointer.hpp new file mode 100644 index 0000000..89c8d7a --- /dev/null +++ b/csrc/mctlass/include/cute/pointer.hpp @@ -0,0 +1,325 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include +#include + +namespace cute +{ + +// +// has_dereference to determine if a type is a pointer concept +// + +template +struct has_dereference : false_type { +}; + +template +struct has_dereference())>> : true_type { +}; + +// +// Pointer categories +// + +template +struct is_gmem : false_type {}; + +template +struct is_smem : false_type {}; + +// Anything that is not gmem or smem is rmem +template +struct is_rmem : bool_constant< not (is_gmem::value || is_smem::value)> {}; + +// +// A very simplified wrapper for pointers -- use for constructing tagged pointers +// +template +struct device_ptr +{ + using value_type = T; + + CUTE_HOST_DEVICE constexpr + device_ptr(T* ptr) : ptr_(ptr) {} + + CUTE_HOST_DEVICE constexpr + T* get() const { return ptr_; } + + CUTE_HOST_DEVICE constexpr + T& operator*() const { return *ptr_; } + + template + CUTE_HOST_DEVICE constexpr + T& operator[](Index const& i) const { return ptr_[i]; } + + template + CUTE_HOST_DEVICE constexpr + DerivedType operator+(Index const& i) const { return {ptr_ + i}; } + + CUTE_HOST_DEVICE constexpr friend + ptrdiff_t operator-(device_ptr const& a, + device_ptr const& b) { + return a.ptr_ - b.ptr_; + } + + T* ptr_; +}; + +// +// gmem_ptr +// + +template +struct gmem_ptr : device_ptr> { + using device_ptr>::device_ptr; +}; + +template +CUTE_HOST_DEVICE constexpr +gmem_ptr +make_gmem_ptr(T* ptr) { + return {ptr}; +} + +template +CUTE_HOST_DEVICE constexpr +gmem_ptr +make_gmem_ptr(void* ptr) { + return {reinterpret_cast(ptr)}; +} + +template +struct is_gmem> : true_type {}; + +// +// smem_ptr +// + +template +struct smem_ptr : device_ptr> { + using device_ptr>::device_ptr; +}; + +template +CUTE_HOST_DEVICE constexpr +smem_ptr +make_smem_ptr(T* ptr) { + return {ptr}; +} + +template +CUTE_HOST_DEVICE constexpr +smem_ptr +make_smem_ptr(void* ptr) { + return {reinterpret_cast(ptr)}; +} + +template +struct is_smem> : true_type {}; + +// +// rmem_ptr +// + +template +struct rmem_ptr : device_ptr> { + using device_ptr>::device_ptr; +}; + +template +CUTE_HOST_DEVICE constexpr +rmem_ptr +make_rmem_ptr(T* ptr) { + return {ptr}; +} + +template +CUTE_HOST_DEVICE constexpr +rmem_ptr +make_rmem_ptr(void* ptr) { + return {reinterpret_cast(ptr)}; +} + +template +struct is_rmem> : true_type {}; + +// +// counting iterator -- quick and dirty +// + +struct counting +{ + using index_type = int; + using value_type = index_type; + + CUTE_HOST_DEVICE constexpr + counting() : n_(0) {} + CUTE_HOST_DEVICE constexpr + counting(index_type const& n) : n_(n) {} + + CUTE_HOST_DEVICE constexpr + index_type operator[](index_type const& i) const { return n_ + i; } + + CUTE_HOST_DEVICE constexpr + index_type const& operator*() const { return n_; } + + CUTE_HOST_DEVICE constexpr + counting operator+(index_type const& i) const { return {n_ + i}; } + CUTE_HOST_DEVICE constexpr + counting& operator++() { ++n_; return *this; } + + CUTE_HOST_DEVICE constexpr + bool operator==(counting const& other) const { return n_ == other.n_; } + CUTE_HOST_DEVICE constexpr + bool operator!=(counting const& other) const { return n_ != other.n_; } + + CUTE_HOST_DEVICE constexpr + bool operator< (counting const& other) const { return n_ < other.n_; } + + index_type n_; +}; + +// +// recast +// + +template +CUTE_HOST_DEVICE constexpr +auto +recast(T* ptr) { + return reinterpret_cast(ptr); +} + +template +CUTE_HOST_DEVICE constexpr +auto +recast(T const* ptr) { + return reinterpret_cast(ptr); +} + +template +CUTE_HOST_DEVICE constexpr +auto +recast(gmem_ptr const& ptr) { + return make_gmem_ptr(recast(ptr.ptr_)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +recast(gmem_ptr const& ptr) { + return make_gmem_ptr(recast(ptr.ptr_)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +recast(smem_ptr const& ptr) { + return make_smem_ptr(recast(ptr.ptr_)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +recast(smem_ptr const& ptr) { + return make_smem_ptr(recast(ptr.ptr_)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +recast(rmem_ptr const& ptr) { + return make_rmem_ptr(recast(ptr.ptr_)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +recast(rmem_ptr const& ptr) { + return make_rmem_ptr(recast(ptr.ptr_)); +} + +// +// Display utilities +// + +template +CUTE_HOST_DEVICE void print(T const* const ptr) +{ + printf("raw_ptr_%db(%p)", int(8*sizeof(T)), ptr); +} + +template +CUTE_HOST_DEVICE void print(gmem_ptr const& ptr) +{ + printf("gmem_ptr_%db(%p)", int(8*sizeof(T)), ptr.get()); +} + +template +CUTE_HOST_DEVICE void print(smem_ptr const& ptr) +{ + printf("smem_ptr_%db(%p)", int(8*sizeof(T)), ptr.get()); +} + +template +CUTE_HOST_DEVICE void print(rmem_ptr const& ptr) +{ + printf("rmem_ptr_%db(%p)", int(8*sizeof(T)), ptr.get()); +} + +#if !defined(__MACACC_RTC__) +template +CUTE_HOST std::ostream& operator<<(std::ostream& os, gmem_ptr const& ptr) +{ + return os << "gmem_ptr_" << int(8*sizeof(T)) << "b"; +} + +template +CUTE_HOST std::ostream& operator<<(std::ostream& os, smem_ptr const& ptr) +{ + return os << "smem_ptr_" << int(8*sizeof(T)) << "b"; +} + +template +CUTE_HOST std::ostream& operator<<(std::ostream& os, rmem_ptr const& ptr) +{ + return os << "rmem_ptr_" << int(8*sizeof(T)) << "b"; +} + +#endif // !defined(__MACACC_RTC__) + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/stride.hpp b/csrc/mctlass/include/cute/stride.hpp new file mode 100644 index 0000000..515bb7b --- /dev/null +++ b/csrc/mctlass/include/cute/stride.hpp @@ -0,0 +1,442 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +namespace cute +{ + +/** crd2idx maps a coordinate within to an index + * This is computed as follows: + * [coord, shape, and stride are all integers => step forward by stride] + * op(c, s, d) => c * d + * [coord is integer, shape and stride are tuple => divmod coord for each mode] + * op(c, (s,S), (d,D)) => op(c % prod(s), s, d) + op(c / prod(s), (S), (D)) + * [coord, shape, and stride are all tuples => consider each mode independently] + * op((c,C), (s,S), (d,D)) => op(c, s, d) + op((C), (S), (D)) + */ + +template +CUTE_HOST_DEVICE constexpr +auto +crd2idx(Coord const& coord, + Shape const& shape, + Stride const& stride); + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +crd2idx_ttt(Coord const& coord, + Shape const& shape, + Stride const& stride, seq) +{ + return (... + crd2idx(get(coord), get(shape), get(stride))); +} + +template +CUTE_HOST_DEVICE constexpr +auto +crd2idx_itt(CInt const& coord, + STuple const& shape, + DTuple const& stride, seq) +{ + if constexpr (sizeof...(Is) == 0) { // Avoid recursion and mod on single/last iter + return crd2idx(coord, get(shape), get(stride)); + } else { // General case + return crd2idx(coord % product(get(shape)), get(shape), get(stride)) + + crd2idx_itt(coord / product(get(shape)), shape, stride, seq{}); + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +crd2idx(Coord const& coord, + Shape const& shape, + Stride const& stride) +{ + if constexpr (is_tuple::value) { + if constexpr (is_tuple::value) { // tuple tuple tuple + static_assert(tuple_size::value == tuple_size< Shape>::value, "Mismatched Ranks"); + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + return detail::crd2idx_ttt(coord, shape, stride, tuple_seq{}); + } else { // tuple "int" "int" + static_assert(sizeof(Coord) == 0, "Invalid parameters"); + } + } else { + if constexpr (is_tuple::value) { // "int" tuple tuple + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + return detail::crd2idx_itt(coord, shape, stride, tuple_seq{}); + } else { // "int" "int" "int" + return coord * stride; + } + } + + CUTE_GCC_UNREACHABLE; +} + +// +// If we know Stride is default [CompactColMajor], then we can take shortcuts +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +crd2idx_horner(CTuple const& coord, + STuple const& shape, seq) +{ + if constexpr (sizeof...(Is) == 0) { // No recursion on single/last iter + return get(coord); + } else { // General case + return get(coord) + get(shape) * crd2idx_horner(coord, shape, seq{}); + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +crd2idx(Coord const& coord, + Shape const& shape) +{ + static_assert(decltype(congruent(coord,shape))::value, "Mismatched Ranks"); + if constexpr (is_tuple::value) { + // Flatten and apply Horner's method + auto flat_coord = flatten(coord); + auto flat_shape = flatten(shape); + return detail::crd2idx_horner(flat_coord, flat_shape, tuple_seq{}); + } else { + return coord; + } + + CUTE_GCC_UNREACHABLE; +} + +/** idx2crd splits an index to a coordinate within . + * + * This is computed as follows: + * [index, shape, and stride are all integers => determine 1D coord] + * op(i, s, d) => (i / d) % s + * [index is integer, shape and stride are tuple => determine component for each mode] + * op(i, (s,S), (d,D)) => (op(i, s, d), op(i, S, D)...) + * [index, shape, and stride are all tuples => consider each mode independently] + * op((i,I), (s,S), (d,D)) => (op(i, s, d), op((I), (S), (D))) + * + * NOTE: This only works for compact shape+stride layouts. A more general version would + * apply to all surjective layouts + */ + +template +CUTE_HOST_DEVICE constexpr +auto +idx2crd(Index const& idx, + Shape const& shape, + Stride const& stride) +{ + if constexpr (is_tuple::value) { + if constexpr (is_tuple::value) { // tuple tuple tuple + static_assert(tuple_size::value == tuple_size< Shape>::value, "Mismatched Ranks"); + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + return transform(idx, shape, stride, [](auto const& i, auto const& s, auto const& d){ return idx2crd(i,s,d); }); + } else { // tuple "int" "int" + static_assert(sizeof(Index) == 0, "Invalid parameters"); + } + } else { + if constexpr (is_tuple::value) { + if constexpr (is_tuple::value) { // "int" tuple tuple + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + return transform(shape, stride, [&](auto const& s, auto const& d){ return idx2crd(idx,s,d); }); + } else { // "int" tuple "int" + return transform(shape, compact_col_major(shape, stride), [&](auto const& s, auto const& d){ return idx2crd(idx,s,d); }); + } + } else { // "int" "int" "int" + if constexpr (is_constant<1, Shape>::value) { + // Skip potential stride-0 division + return Int<0>{}; + } else { + return (idx / stride) % shape; + } + } + } + + CUTE_GCC_UNREACHABLE; +} + +// +// If we know Stride is default [CompactColMajor], then we can take shortcuts +// + +//(idx / 1) % s0 +//(idx / s0) % s1 +//(idx / (s0 * s1)) % s2 +//... + +template +CUTE_HOST_DEVICE constexpr +auto +idx2crd(Index const& idx, + Shape const& shape) +{ + if constexpr (is_tuple::value) { + if constexpr (is_tuple::value) { // tuple tuple + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + return transform(idx, shape, [](auto const& i, auto const& s) { return idx2crd(i,s); }); + } else { // tuple "int" + static_assert(sizeof(Index) == 0, "Invalid parameters"); + } + } else { + if constexpr (is_tuple::value) { // "int" tuple + return idx2crd(idx, shape, compact_col_major(shape)); + } else { // "int" "int" + return idx; + } + } + + CUTE_GCC_UNREACHABLE; +} + +// +// crd2crd +// + +template +CUTE_HOST_DEVICE constexpr +auto +crd2crd(Coord const& coord, + SShape const& src_shape, + DShape const& dst_shape) +{ + if constexpr (is_tuple::value && is_tuple::value && is_tuple::value) { + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + return transform(coord, src_shape, dst_shape, [](auto const& c, auto const& s, auto const& d) { return crd2crd(c,s,d); }); + } else { + // assert(size(src_shape) == size(dst_shape)) + return idx2crd(crd2idx(coord, src_shape), dst_shape); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Compact Major +// + +// Tags for common layouts and dispatching +struct LayoutLeft; // Col-major layout mapping; leftmost extent has stride 1 +using GenColMajor = LayoutLeft; // Alias + +struct LayoutRight; // Row-major layout mapping; rightmost extent has stride 1 +using GenRowMajor = LayoutRight; // Alias + +namespace detail { + +// GGC8.5 WAR -- Use of lambdas in unevaluated contexts. Instead use function objects. +template +struct CompactLambda; + +// @pre is_integral +// Return (result, current * product(shape)) to enable recurrence +template +CUTE_HOST_DEVICE constexpr +auto +compact(Shape const& shape, + Current const& current) +{ + if constexpr (is_tuple::value) { // Shape::tuple Current::int + using Lambda = CompactLambda; // Append or Prepend + using Seq = typename Lambda::template seq; // Seq or RSeq + return cute::detail::fold(shape, cute::make_tuple(cute::make_tuple(), current), Lambda{}, Seq{}); + } else { // Shape::int Current::int + if constexpr (is_constant<1, Shape>::value) { + return cute::make_tuple(Int<0>{}, current); // If current is dynamic, this could save a reg + } else { + return cute::make_tuple(current, current * shape); + } + } + + CUTE_GCC_UNREACHABLE; +} + +// GCC8.5 WAR -- Specialization LayoutLeft +template <> +struct CompactLambda +{ + template + CUTE_HOST_DEVICE constexpr auto + operator()(Init const& init, Shape const& si) { + auto result = detail::compact(si, get<1>(init)); + return cute::make_tuple(append(get<0>(init), get<0>(result)), get<1>(result)); // Append + } + + template + using seq = tuple_seq; // Seq +}; + +// GCC8.5 WAR -- Specialization LayoutRight +template <> +struct CompactLambda +{ + template + CUTE_HOST_DEVICE constexpr auto + operator()(Init const& init, Shape const& si) { + auto result = detail::compact(si, get<1>(init)); + return cute::make_tuple(prepend(get<0>(init), get<0>(result)), get<1>(result)); // Prepend + } + + template + using seq = tuple_rseq; // RSeq +}; + +} // end namespace detail + +template , + __CUTE_REQUIRES(is_tuple::value || is_integral::value)> +CUTE_HOST_DEVICE constexpr +auto +compact_major(Shape const& shape, + Current const& current = {}) +{ + if constexpr (is_tuple::value) { // Shape::tuple Current::tuple + static_assert(is_tuple::value, "Invalid parameters"); + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + // Recurse to apply to the terminals of current + return transform(shape, current, [&](auto const& s, auto const& c){ return compact_major(s,c); }); + } else { + return get<0>(detail::compact(shape, current)); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Compact Col Major +// + +struct LayoutLeft { + template + using Apply = decltype(compact_major(declval())); +}; + +template > +CUTE_HOST_DEVICE constexpr +auto +compact_col_major(Shape const& shape, + Current const& current = {}) +{ + return compact_major(shape, current); +} + +// +// Compact Row Major +// + +struct LayoutRight { + template + using Apply = decltype(compact_major(declval())); +}; + +template > +CUTE_HOST_DEVICE constexpr +auto +compact_row_major(Shape const& shape, + Current const& current = {}) +{ + return compact_major(shape, current); +} + +// +// Compact Order -- compute a compact stride based on an ordering of the modes +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +compact_order(Shape const& shape, Order const& order, + OrigShape const& orig_shape, OrigOrder const& orig_order) +{ + if constexpr (is_tuple::value) { + return transform(shape, order, [&](auto const& x, auto const& y) { return compact_order(x, y, orig_shape, orig_order); }); + } else { + auto d = product(transform(orig_shape, orig_order, + [&](auto const& s, auto const& o) { + return conditional_return(o < order, product(s), Int<1>{}); + })); + return compact_col_major(shape, d); + } + + CUTE_GCC_UNREACHABLE; +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +compact_order(Shape const& shape, Order const& order) +{ + static_assert(is_congruent::value, "Need congruence of shape and order."); + return detail::compact_order(shape, order, flatten_to_tuple(shape), flatten_to_tuple(order)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +compact_order(Shape const& shape, GenColMajor const& major) +{ + return compact_major(shape); +} + +template +CUTE_HOST_DEVICE constexpr +auto +compact_order(Shape const& shape, GenRowMajor const& major) +{ + return compact_major(shape); +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/swizzle.hpp b/csrc/mctlass/include/cute/swizzle.hpp new file mode 100644 index 0000000..d35abda --- /dev/null +++ b/csrc/mctlass/include/cute/swizzle.hpp @@ -0,0 +1,527 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include +#include +#include +#include + +namespace cute +{ + +// A generic Swizzle functor +/* 0bxxxxxxxxxxxxxxxYYYxxxxxxxZZZxxxx + * ^--^ MBase is the number of least-sig bits to keep constant + * ^-^ ^-^ BBits is the number of bits in the mask + * ^---------^ SShift is the distance to shift the YYY mask + * (pos shifts YYY to the right, neg shifts YYY to the left) + * + * e.g. Given + * 0bxxxxxxxxxxxxxxxxYYxxxxxxxxxZZxxx + * the result is + * 0bxxxxxxxxxxxxxxxxYYxxxxxxxxxAAxxx where AA = ZZ xor YY + */ +template +struct Swizzle +{ + static constexpr int num_bits = BBits; + static constexpr int num_base = MBase; + static constexpr int num_shft = SShift; + + static_assert(num_base >= 0, "MBase must be positive."); + static_assert(num_bits >= 0, "BBits must be positive."); + static_assert(abs(num_shft) >= num_bits, "abs(SShift) must be more than BBits."); + + // using 'int' type here to avoid unintentially casting to unsigned... unsure. + using bit_msk = cute::constant; + using yyy_msk = cute::constant; + using zzz_msk = cute::constant; + using msk_sft = cute::constant; + + static constexpr uint32_t swizzle_code = uint32_t(yyy_msk{} | zzz_msk{}); + + template ::value)> + CUTE_HOST_DEVICE constexpr static + auto + apply(Offset const& offset) + { + return offset ^ shiftr(offset & yyy_msk{}, msk_sft{}); // ZZZ ^= YYY + } + + template ::value)> + CUTE_HOST_DEVICE constexpr + auto + operator()(Offset const& offset) const + { + return apply(offset); + } +}; + +template +using SwizzleXor = Swizzle; + +// +// make_swizzle<0b1000, 0b0100>() -> Swizzle<1,2,1> +// make_swizzle<0b11000000, 0b00000110>() -> Swizzle<2,1,5> +// + +template +CUTE_HOST_DEVICE constexpr +auto +make_swizzle() +{ + constexpr uint32_t BZ = popcount(Y); // Number of swizzle bits + constexpr uint32_t BY = popcount(Z); // Number of swizzle bits + static_assert(BZ == BY, "Number of bits in Y and Z don't match"); + constexpr uint32_t TZ_Y = countr_zero(Y); // Number of trailing zeros in Y + constexpr uint32_t TZ_Z = countr_zero(Z); // Number of trailing zeros in Z + constexpr uint32_t M = cute::min(TZ_Y, TZ_Z) % 32; + constexpr int32_t S = int32_t(TZ_Y) - int32_t(TZ_Z); // Difference in trailing zeros + static_assert((Y | Z) == Swizzle::swizzle_code, "Something went wrong."); + return Swizzle{}; +} + +template +CUTE_HOST_DEVICE constexpr +auto +composition(Swizzle, Swizzle) +{ + static_assert(S0 == S1, "Can only merge swizzles of the same shift."); + constexpr uint32_t Y = Swizzle::yyy_msk::value ^ Swizzle::yyy_msk::value; + constexpr uint32_t Z = Swizzle::zzz_msk::value ^ Swizzle::zzz_msk::value; + return make_swizzle(); + + //return ComposedFn, Swizzle>{}; +} + +// +// Upcast and Downcast +// + +template +CUTE_HOST_DEVICE constexpr +auto +upcast(Swizzle const& swizzle) +{ + static_assert(has_single_bit(N), "N must be a power of two"); + constexpr int log2_n = bit_width(uint32_t(N)) - 1; + constexpr int NewM = M - log2_n; + if constexpr (NewM >= 0) { + return Swizzle{}; + } else { + return Swizzle{}; + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +downcast(Swizzle const& swizzle) +{ + static_assert(has_single_bit(N), "N must be a power of two"); + constexpr int log2_n = bit_width(uint32_t(N)) - 1; + return Swizzle{}; +} + +template +CUTE_HOST_DEVICE constexpr +auto +recast(Swizzle const& swizzle) +{ + if constexpr (sizeof_bits::value == sizeof_bits::value) { + return swizzle; + } else if constexpr (sizeof_bits::value > sizeof_bits::value) { + static_assert(sizeof_bits::value % sizeof_bits::value == 0, "NewType must be a multiple of OldType"); + return upcast::value/sizeof_bits::value>(swizzle); + } else if constexpr (sizeof_bits::value < sizeof_bits::value) { + static_assert(sizeof_bits::value % sizeof_bits::value == 0, "NewType must be a divisor of OldType"); + return downcast::value/sizeof_bits::value>(swizzle); + } +} + +// +// Utility for slicing and swizzle "offsets" +// + +// For swizzle functions, it is often needed to keep track of which bits are +// consumed and which bits are free. Furthermore, it is useful to know whether +// each of these bits is known statically or dynamically. + +// MixedBits is an integer class where some bits are known statically and some +// bits are known dynamically. These sets of bits are disjoint and it is known +// statically which bits are known dynamically. + +// MixedBits can only be manipulated through bitwise operations + +// Abstract value: StaticInt | (dynamic_int_ & StaticFlags) +template // 0: static, 1: dynamic +struct MixedBits +{ + // Representation invariants + static_assert(StaticFlags != 0, "Should be at least one dynamic bit in MixedBits."); + static_assert((StaticInt & StaticFlags) == 0, "No static/dynamic overlap allowed in MixedBits."); + // assert((dynamic_int_ & ~F) == 0); + + DynamicType dynamic_int_; +}; + +template +CUTE_HOST_DEVICE constexpr +auto +make_mixed_bits(constant const&, DynamicType const& d, constant const&) +{ + static_assert(is_integral::value); + if constexpr (is_static::value) { + static_assert((s & DynamicType::value & f) == 0, "No static/dynamic overlap allowed."); + return constant{} | (d & constant{}); // Just return a static int + } else if constexpr (f == 0) { + return constant{}; // Just return a static int + } else { + return MixedBits{d & f}; // MixedBits + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Explicit conversion for now -- consider casting on plus or minus +// + +template +CUTE_HOST_DEVICE constexpr +auto +to_integral(MixedBits const& m) +{ + //return S | (m.dynamic_int_ & F); + return S | m.dynamic_int_; +} + +// Any cute::is_integral +template ::value)> +CUTE_HOST_DEVICE constexpr +auto +to_integral(I const& i) +{ + return i; +} + +// +// Operators +// + +// Equality +template +CUTE_HOST_DEVICE constexpr +auto +operator==(MixedBits const& m, constant const&) +{ + return (S0 == (S1 & ~F0)) && (m.dynamic_int_ == (S1 & F0)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator==(constant const& s, MixedBits const& m) +{ + return m == s; +} + +// Bitwise AND +template +CUTE_HOST_DEVICE constexpr +auto +operator&(MixedBits const& m0, MixedBits const& m1) +{ + // Truth table for (S0,D0,F0) & (S1,D1,F1) -> (S,D,F) + // S0D0F0 | 0X0 | 001 | 011 | 1X0 | + // S1D1F1 + // 0X0 | 0X0 | 0X0 | 0X0 | 0X0 | + // 001 | 0X0 | 001 | 001 | 001 | + // 011 | 0X0 | 001 | 011 | 011 | + // 1X0 | 0X0 | 001 | 011 | 1X0 | + + return make_mixed_bits(constant{}, + //(S0 | m0.dynamic_int_) & (S1 | m1.dynamic_int_), + ((S1 & F0) & m0.dynamic_int_) | ((S0 & F1) & m1.dynamic_int_) | (m0.dynamic_int_ & m1.dynamic_int_), + constant{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator&(MixedBits const& m, constant const&) +{ + return make_mixed_bits(constant{}, + m.dynamic_int_, + constant{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator&(constant const& s, MixedBits const& m) +{ + return m & s; +} + +// Bitwise OR +template +CUTE_HOST_DEVICE constexpr +auto +operator|(MixedBits const& m0, MixedBits const& m1) +{ + // Truth table for (S0,D0,F0) | (S1,D1,F1) -> (S,D,F) + // S0D0F0 | 0X0 | 001 | 011 | 1X0 | + // S1D1F1 + // 0X0 | 0X0 | 001 | 011 | 1X0 | + // 001 | 001 | 001 | 011 | 1X0 | + // 011 | 011 | 011 | 011 | 1X0 | + // 1X0 | 1X0 | 1X0 | 1X0 | 1X0 | + + return make_mixed_bits(constant{}, + ((~S1 & F0) & m0.dynamic_int_) | ((~S0 & F1) & m1.dynamic_int_), + constant{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator|(MixedBits const& m, constant const&) +{ + return make_mixed_bits(constant{}, + m.dynamic_int_, + constant{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator|(constant const& s, MixedBits const& m) +{ + return m | s; +} + +// Bitwise XOR +template +CUTE_HOST_DEVICE constexpr +auto +operator^(MixedBits const& m0, MixedBits const& m1) +{ + // Truth table for (S0,D0,F0) ^ (S1,D1,F1) -> (S,D,F) + // S0D0F0 | 0X0 | 001 | 011 | 1X0 | + // S1D1F1 + // 0X0 | 0X0 | 001 | 011 | 1X0 | + // 001 | 001 | 001 | 011 | 011 | + // 011 | 011 | 011 | 001 | 001 | + // 1X0 | 1X0 | 011 | 001 | 0X0 | + + return make_mixed_bits(constant{}, + (S0 | m0.dynamic_int_) ^ (S1 | m1.dynamic_int_), + constant{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator^(MixedBits const& m, constant const&) +{ + return make_mixed_bits(constant{}, + (S0 | m.dynamic_int_) ^ S1, + constant{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +operator^(constant const& s, MixedBits const& m) +{ + return m ^ s; +} + +// +// upcast and downcast +// + +template +CUTE_HOST_DEVICE constexpr +auto +safe_div(MixedBits const& m, constant const& s) +{ + static_assert(has_single_bit(S1), "Only divide MixedBits by powers of two."); + return make_mixed_bits(safe_div(constant{}, s), + safe_div(m.dynamic_int_, s), + safe_div(constant{}, s)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +upcast(MixedBits const& m) +{ + static_assert(has_single_bit(N), "Only divide MixedBits by powers of two."); + return safe_div(m, constant{}); +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +auto +upcast(T const& m) +{ + return safe_div(m, constant{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +downcast(MixedBits const& m) +{ + static_assert(has_single_bit(N), "Only scale MixedBits by powers of two."); + return make_mixed_bits(constant{}, + m.dynamic_int_ * N, + constant{}); +} + +template ::value)> +CUTE_HOST_DEVICE constexpr +auto +downcast(T const& m) +{ + return m * constant{}; +} + +// +// Convert a Pow2Layout+Coord to a MixedBits +// + +template +CUTE_HOST_DEVICE constexpr +auto +to_mixed_bits(Shape const& shape, Stride const& stride, Coord const& coord) +{ + if constexpr (is_tuple::value && is_tuple::value && is_tuple::value) { + static_assert(tuple_size::value == tuple_size::value, "Mismatched ranks"); + static_assert(tuple_size::value == tuple_size::value, "Mismatched ranks"); + return transform_apply(shape, stride, coord, [](auto const& s, auto const& d, auto const& c) { return to_mixed_bits(s,d,c); }, + [](auto const&... a) { return (a ^ ...); }); + } else if constexpr (is_integral::value && is_integral::value && is_integral::value) { + static_assert(decltype(shape*stride)::value == 0 || has_single_bit(decltype(shape*stride)::value), "Requires pow2 shape*stride."); + return make_mixed_bits(Int<0>{}, coord * stride, (shape - Int<1>{}) * stride); + } else { + static_assert(is_integral::value && is_integral::value && is_integral::value, "Either Shape, Stride, and Coord must be all tuples, or they must be all integral (in the sense of cute::is_integral)."); + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +to_mixed_bits(Layout const& layout, Coord const& coord) +{ + return to_mixed_bits(layout.shape(), layout.stride(), idx2crd(coord, layout.shape())); +} + +// +// Display utilities +// + +template +CUTE_HOST_DEVICE void print(MixedBits const& m) +{ + printf("M_%u|(%u&%u)=%u", S, uint32_t(m.dynamic_int_), F, to_integral(m)); +} + +#if !defined(__MACACC_RTC__) +template +CUTE_HOST std::ostream& operator<<(std::ostream& os, MixedBits const& m) +{ + return os << "M_" << S << "|(" << uint32_t(m.dynamic_int_) << "&" << F << ")=" << to_integral(m); +} + +template +CUTE_HOST_DEVICE void print(Swizzle const&) +{ + print("S<%d,%d,%d>", B, M, S); +} + +template +CUTE_HOST std::ostream& operator<<(std::ostream& os, Swizzle const&) +{ + return os << "S<" << B << "," << M << "," << S << ">"; +} +#endif // !defined(__MACACC_RTC__) + +/* + * input: + * : How many elements does a thread move at a time + * For example: copy128B, fp16 -> 8, copy128B, fp8 ->16 + * : Same as swizzle + * output: + * offset corresponding to original address + * + * usage: + * auto tAgA = make_tensor(tAgA_noSwizzle.data() + get_swizzle_offset<8,3,3,3>(tidx), layout(tAgA_noSwizzle)); + */ +template +CUTE_HOST_DEVICE int get_swizzle_offset(int tidx) { + + static_assert(MBase >= 0, "MBase must be positive."); + static_assert(BBits >= 0, "BBits must be positive."); + static_assert(SShift >= BBits, "SShift must be more than BBits."); + + int Mbase_thread = (1 << MBase) / elementPerThread; + // eg. (0,0,0) + if (Mbase_thread <= 0) { + return 0; + } + int thread_per_row = Mbase_thread * (1 << SShift); + int row_id = tidx / thread_per_row; + int col_id = tidx % thread_per_row / Mbase_thread; + int col_id_new = col_id ^ (row_id % (1 << BBits)); + return (col_id_new - col_id) * (1 << MBase); +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/swizzle_layout.hpp b/csrc/mctlass/include/cute/swizzle_layout.hpp new file mode 100644 index 0000000..d5cfbbd --- /dev/null +++ b/csrc/mctlass/include/cute/swizzle_layout.hpp @@ -0,0 +1,1012 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +#include + +/* This implements a ComposedLayout of the form + * InvolutionFn o OffsetPlus o Layout + * where the InvolutionFn need not be linear (hence the need for the Offset). + * + * This ComposedLayout provides similar coordinate-to-index mapping and layout manipulations, + * but is not considered a "normal" layout. + * For example, this layout provides size() functions, but does not provide stride() functions. + * + * Furthermore, for known InvolutionFns, this layout attempts to decay itself + * to a normal-layout with dynamic or static strides. + * This is possible by determining the subdomain of the Involution function + * that is identity and testing if the right Layout's codomain is contained + * within it. + */ + +namespace cute +{ + +// A Layout of non-trivially composable functions: F o I o L +template +struct ComposedLayout + : private cute::tuple // EBO for static layouts +{ + CUTE_HOST_DEVICE constexpr + ComposedLayout(InvolutionFn const& fn = {}, + IntermediateOffset const& offset = {}, + Layout const& layout = {}) + : cute::tuple(fn, offset, layout) + {} + + // + // Accessors + // + + static constexpr int rank = Layout::rank; + + CUTE_HOST_DEVICE constexpr + decltype(auto) + swizzle_fn() const { + return get<0>(static_cast const&>(*this)); + } + + CUTE_HOST_DEVICE constexpr + decltype(auto) + offset_fn() const { + return get<1>(static_cast const&>(*this)); + } + + CUTE_HOST_DEVICE constexpr + decltype(auto) + layout_fn() const { + return get<2>(static_cast const&>(*this)); + } + + CUTE_HOST_DEVICE constexpr + decltype(auto) + layout() const { + return *this; + } + + CUTE_HOST_DEVICE constexpr + decltype(auto) + shape() const { + return layout_fn().shape(); + } + + // Doesn't really make sense to ask for the strides of this "layout" + CUTE_HOST_DEVICE constexpr + decltype(auto) + stride() const = delete; + + // + // Mappings + // + + // Map a logical coordinate to a linear index (Coord has no Underscore slice operators) + // OR + // Slice the layout and return the sublayout (Coord has an Underscore slice op) + template + CUTE_HOST_DEVICE constexpr + auto + operator()(Coord const& coord) const { + if constexpr (has_underscore::value) { + return slice(coord, *this); + } else { + return swizzle_fn()(to_integral(offset_fn()) + layout_fn()(coord)); // (F o L)(c) + } + + CUTE_GCC_UNREACHABLE; + } + + // Map a 1D linear coordinate to a flat ND logical coordinate + template ::value)> + CUTE_HOST_DEVICE constexpr + auto + operator[](Int const& linear_idx) const { + return get_flat_coord(linear_idx); + } + + // Convenience function for multi-dimensional coordinates + template + CUTE_HOST_DEVICE constexpr + auto + operator()(Coord0 const& c0, Coord1 const& c1, Coords const&... cs) const { + return operator()(make_coord(c0,c1,cs...)); + } + + // + // Compose + // + + template + CUTE_HOST_DEVICE constexpr + auto + compose(OtherLayout const& other) const { + return composition(*this, other); + } + + template + CUTE_HOST_DEVICE constexpr + auto + compose(Layouts const&... layouts) const { + return composition(*this, make_tile(layouts...)); + } + + template + CUTE_HOST_DEVICE constexpr + auto + with_shape(OtherShape const& shape) const { + return composition(*this, make_layout(shape)); + } + + template + CUTE_HOST_DEVICE constexpr + auto + with_shape(Shapes const&... shapes) const { + return composition(*this, make_layout(make_shape(shapes...))); + } + + // + // Tile + // + + template + CUTE_HOST_DEVICE constexpr + auto + tile(OtherLayout const& other) const { + return tiled_divide(*this, other); + } + + template + CUTE_HOST_DEVICE constexpr + auto + tile(Layouts const&... layouts) const { + return tiled_divide(*this, make_tile(layouts...)); + } + + // + // Utility + // + + // + // Index to Coordinate + // + + // NOTE Only valid for compact layouts + + // Return the (hierarchical) ND logical coordinate corresponding to the linear index + // @post this->crd2idx(@a result) == idx + // @post congruent(@a result, shape()) + template ::value)> + CUTE_HOST_DEVICE constexpr + auto + get_hier_coord(IInt const& idx) const { + return layout_fn().get_hier_coord(swizzle_fn()(idx) - to_integral(offset_fn())); // (L^-1 o F)(k) + } + + // Return the (flat) ND logical coordinate corresponding to the linear index + // @post this->crd2idx(@a result) == idx + // @post rank(@a result) == rank(shape()) && depth(@a result) == 1 + template ::value)> + CUTE_HOST_DEVICE constexpr + auto + get_flat_coord(IInt const& idx) const { + return layout_fn().get_flat_coord(swizzle_fn()(idx) - to_integral(offset_fn())); // (L^-1 o F)(k) + } + + // Return the generalized column-major 1D logical coordinate corresponding to the linear index + // @post this->crd2idx(@a result) == idx + // @post is_integral::value + template ::value)> + CUTE_HOST_DEVICE constexpr + auto + get_1d_coord(IInt const& idx) const { + return layout_fn().get_1d_coord(swizzle_fn()(idx) - to_integral(offset_fn())); // (L^-1 o F)(k) + } +}; + +template +struct is_layout> : true_type {}; + +template +struct is_composed_layout : false_type {}; +template +struct is_composed_layout> : true_type {}; + +// +// Constructors +// + +template +CUTE_HOST_DEVICE constexpr +auto +make_layout(Swizzle const& sxor) +{ + return composition(sxor, Layout,Int<1>>{}); +} + +template +CUTE_HOST_DEVICE constexpr +auto +make_layout(ComposedLayout const& a, Layout const& b) +{ + return composition(a.swizzle_fn(), a.offset_fn(), make_layout(a.layout_fn(), b)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +make_layout(Layout const& a, ComposedLayout const& b) +{ + return composition(b.swizzle_fn(), b.offset_fn(), make_layout(a, b.layout_fn())); +} + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +transfer_swizzle(Layout const& old_layout, + Layout const& new_layout) +{ + // Our goal is to determine a new swizzle for the strides in new_layout for consistent vectorizations + + // This is accomplished by identifying + // S o L :=: S? o L* + // We identify the "active" portion of S by computing (P o L)(c*) where P is a projection generated by S + // Then that active identifier is transformed through the layouts: + // L*(L[(P o L)(c*)]) + // which is a new swizzle identifier for S?, the new swizzle + + // Projections of the swizzle layout for composition, P + auto swizzle_only_zy = make_layout(make_shape (Int<(1 << M)>{}, Int<(1 << B)>{}, Int<(1 << (abs(S)-B))>{}, Int<(1 << B )>{}, Int<1>{}), + make_stride( Int<0>{}, Int<(1 << M)>{}, Int<0>{}, Int<(1 << (M+abs(S)))>{}, Int<0>{})); + + // Compose with the tile to get the swizzle projection, P o L [The Z and Y contributing portions of L] + auto layout_only_zy = composition(swizzle_only_zy, old_layout); + // Transform the end coordinate to get the active bits of the swizzle, (P o L)(c*) + auto swizzle_active_bits = layout_only_zy(size(layout_only_zy)-Int<1>{}); + + // Get the Z bit and the Y bits -- keep only those that are active in Z *and* Y + auto zzz_msk = typename Swizzle::zzz_msk{}; + auto yyy_msk = typename Swizzle::yyy_msk{}; + auto msk_sft = typename Swizzle::msk_sft{}; + auto active_Z = swizzle_active_bits & shiftr(swizzle_active_bits, msk_sft) & zzz_msk; + auto active_Y = swizzle_active_bits & shiftr(swizzle_active_bits, -msk_sft) & yyy_msk; + + // Pass the identifiers through the old layout and new layout to make a new swizzle identifier, L*(L[(P o L)(c*)]) + auto new_active_Z = new_layout(old_layout.get_1d_coord(active_Z)); + auto new_active_Y = new_layout(old_layout.get_1d_coord(active_Y)); + + // Use this new swizzle identifier to construct the new swizzle for new_layout + // (this also makes sure it's a "valid" swizzle that Swizzle can represent) + return composition(make_swizzle(), new_layout); +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +make_fragment_like(ComposedLayout,Offset,Layout> const& layout) +{ + return detail::transfer_swizzle(layout.layout_fn(), make_fragment_like(layout.layout_fn())); +} + +// +// Utilities +// + +// Return the layout of a mode +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +layout(ComposedLayout const& clayout) +{ + return composition(clayout.swizzle_fn(), clayout.offset_fn(), layout(clayout.layout_fn())); +} + +// Return the shape of a mode +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +shape(ComposedLayout const& layout) +{ + return shape(layout.layout_fn()); +} + +// Doesn't make sense to directly ask for the strides of this "layout" +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +stride(ComposedLayout const& layout) = delete; + +// Return the number of elements in a mode +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +size(ComposedLayout const& layout) +{ + return size(layout.layout_fn()); +} + +// Return the number of modes +template +CUTE_HOST_DEVICE constexpr +auto +rank(ComposedLayout const& layout) +{ + return rank(layout.layout_fn()); +} + +// Return the depth of the layout +template +CUTE_HOST_DEVICE constexpr +auto +depth(ComposedLayout const& layout) +{ + return depth(layout.layout_fn()); +} + +// Return the codomain size of a mode +template +CUTE_HOST_DEVICE constexpr +auto +cosize(ComposedLayout const& layout) +{ + return cosize(layout.layout_fn()); +} + +// +// Operations to manipulate Layouts like a tuple of pairs +// + +template +CUTE_HOST_DEVICE constexpr +auto +get(ComposedLayout const& a) +{ + return composition(a.swizzle_fn(), a.offset_fn(), get(a.layout_fn())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +take(ComposedLayout const& a) +{ + return composition(a.swizzle_fn(), a.offset_fn(), take(a.layout_fn())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +flatten(ComposedLayout const& a) +{ + return composition(a.swizzle_fn(), a.offset_fn(), flatten(a.layout_fn())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +append(ComposedLayout const& a, X const& x) +{ + return composition(a.swizzle_fn(), a.offset_fn(), append(a.layout_fn(), x)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +group(ComposedLayout const& a) +{ + return composition(a.swizzle_fn(), a.offset_fn(), group(a.layout_fn())); +} + +// +// Slice a ComposedLayout +// + +namespace detail { + +template +CUTE_HOST_DEVICE constexpr +auto +make_swizzle_strides(true_type, + IntZ const& Z, + IntY const& Y, + Offset const& offset, + int_sequence) +{ + // Below is an optimized/compressed version of: + //return cute::make_tuple((swizzle(offset + Z*Int<(1 << I)>{}) - swizzle(offset))...); + // with knowledge of Swizzle, I... ranges for each B bits, + // and the layout won't slice along z-bits that are already set + + // y\z 0 1 + // 0 Z DC + // 1 -Z DC + + return cute::make_tuple(conditional_return((offset & (Y << Int{})) == Int<0>{}, Z << Int{}, -(Z << Int{}))...); +} + +template +CUTE_HOST_DEVICE constexpr +auto +make_swizzle_strides(false_type, + IntZ const& Z, + IntY const& Y, + Offset const& offset, + int_sequence) +{ + // Below is an optimized/compressed version of: + //return cute::make_tuple((swizzle(offset + Y*Int<(1 << I)>{}) - swizzle(offset))...); + // with knowledge of Swizzle, I... ranges for each B bits, + // and the layout won't slice along y-bits that are already set + + // y\z 0 1 + // 0 Y+Z Y-Z + // 1 DC DC + + return cute::make_tuple(conditional_return((offset & (Z << Int{})) == Int<0>{}, (Y+Z) << Int{}, (Y-Z) << Int{})...); +} + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr +auto +slice_and_offset(Coord const& coord, ComposedLayout,Offset,Layout> const& layout) +{ + if constexpr (all_underscore::value) { + // Skip the expensive/complicated attempt to decay to a normal layout and just reshape + return cute::make_tuple(composition(layout.swizzle_fn(), layout.offset_fn(), slice(coord, layout.layout_fn())), Int<0>{}); + } else { + + // Projections of the swizzle layout for composition + auto sw = make_layout(make_shape(Int<(1 << M)>{}, Int<(1 << B)>{}, Int<(1 << (abs(S)-B))>{}, Int<(1 << B)>{}, Int<1>{})); + + auto swizzle_anti_zy = make_layout(shape(sw), + make_stride(stride<0>(sw), Int<0>{}, stride<2>(sw), Int<0>{}, size(sw))); + auto swizzle_only_zy = make_layout(shape(sw), + make_stride( Int<0>{}, stride<1>(sw), Int<0>{}, stride<3>(sw), Int<0>{})); + + // The portion of the layout that is not yet consumed + auto sliced_layout = slice(coord, layout.layout_fn()); + + // If the sliced_layout hits two bits that are swizzled together, then don't attempt to decay + + // Compose with the layout to get the swizzle projection, P o L [The Z and Y contributing portions of L] + // (this also tests that shape/stride of layout compose with swizzle) + auto sliced_layout_only_zy = composition(swizzle_only_zy, sliced_layout); + // Transform the end coordinate to get the active bits of the swizzle, (P o L)(c*) + auto swizzle_active_bits = sliced_layout_only_zy(size(sliced_layout_only_zy)-Int<1>{}); + // Determine if any active bits collide under the swizzle + auto hit_ZandY = !(swizzle_active_bits & ~layout.swizzle_fn()(swizzle_active_bits)); + + // The portion of the layout that we are consuming now + auto diced_layout = dice(coord, layout.layout_fn()); + auto diced_coord = dice(coord, coord); + + auto diced_layout_anti_zy = composition(swizzle_anti_zy, diced_layout); + auto diced_layout_only_zy = composition(swizzle_only_zy, diced_layout); + + // New swizzle and offset + auto swizzle = layout.swizzle_fn(); + // offset_only_zy interacts with swizzle and gets accumulated with layout.offset_fn() + // being careful about the static/dynamic contributions from diced_layout and diced_coord + auto offset_only_zy = layout.offset_fn() ^ to_mixed_bits(diced_layout_only_zy, diced_coord); + // offset_anti_zy always gets passed through, no interaction with swizzle + auto offset_anti_zy = diced_layout_anti_zy(diced_coord); + + // If Layout's codomain hits on Y AND Z, then it's not reducible + // If Layout's codomain hits on Y XOR Z, then it's dynamic-normal + // If Layout's codomain hits on neither Y NOR Z, then it's static-normal + + // Test the sliced layout for hit_X & hit_Y for potential decay + if constexpr (is_constant::value) + { // Hits on Y AND Z, so it's not reducible + return cute::make_tuple(composition(swizzle, offset_only_zy, sliced_layout), offset_anti_zy); + } else + { // Misses on Y or Z, so it's static-normal or dynamic-normal + + // Lowest bit of the Z and Y masks + auto Z = typename Swizzle::zzz_msk{} & -typename Swizzle::zzz_msk{}; + auto Y = typename Swizzle::yyy_msk{} & -typename Swizzle::yyy_msk{}; + auto stride_lo = detail::make_swizzle_strides(Z < Y, Z, Y, offset_only_zy, make_int_sequence{}); + auto stride_hi = detail::make_swizzle_strides(Z > Y, Z, Y, offset_only_zy, make_int_sequence{}); + + // Construct a (dynamic) layout that we can perform the composition with + auto swizzle_layout = make_layout(make_shape (Int<(1 << M)>{}, repeat(Int<2>{}), Int<(1 << (abs(S)-B))>{}, repeat(Int<2>{}), Int< 1>{}), + make_stride(Int< 1>{}, stride_lo, Int<(1 << (M+B))>{}, stride_hi , Int<(1 << (M+B+abs(S)))>{})); + + // Decay to a normal layout with offset + return cute::make_tuple(composition(swizzle_layout, sliced_layout), + swizzle(to_integral(offset_only_zy)) + offset_anti_zy); + } + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +slice(Coord const& coord, ComposedLayout const& layout) +{ + return get<0>(slice_and_offset(coord, layout)); +} + +// +// composition +// + +template +CUTE_HOST_DEVICE constexpr +auto +composition(Swizzle const& sxor, + Offset const& offset, + Layout const& layout) +{ + return ComposedLayout>{sxor, offset, layout}; +} + +template +CUTE_HOST_DEVICE constexpr +auto +composition(Swizzle const& sxor, + Offset const& offset, + ComposedLayout const& layout) +{ + // Assume disjoint swizzles and offsets for commutivity + return composition(composition(sxor,layout.swizzle_fn()), offset ^ layout.offset_fn(), layout.layout_fn()); +} + +// Ignore identity case +template +CUTE_HOST_DEVICE constexpr +auto +composition(Swizzle<0,M,S> const&, + Int<0> const&, + Layout const& layout) +{ + return layout; +} + +template +CUTE_HOST_DEVICE constexpr +auto +composition(Swizzle const& sxor, + Layout const& layout) +{ + return composition(sxor, Int<0>{}, layout); +} + +template +CUTE_HOST_DEVICE constexpr +auto +composition(ComposedLayout const& a, + LayoutOrTile const& b) +{ + return composition(a.swizzle_fn(), a.offset_fn(), composition(a.layout_fn(), b)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +composition(Layout const& a, + Swizzle const& b) +{ + // Get the Z bits and the Y bits + auto active_Y = a(typename Swizzle::yyy_msk{}); + auto active_Z = a(typename Swizzle::zzz_msk{}); + + // Works in simple cases... but could be greatly generalized + + return composition(make_swizzle(), a); +} + +template +CUTE_HOST_DEVICE constexpr +auto +composition(Layout const& a, + ComposedLayout const& b) +{ + CUTE_STATIC_ASSERT_V(b.offset_fn() == Int<0>{}, "Require Swizzle offset == 0."); + + return composition(composition(a, b.swizzle_fn()), b.layout_fn()); +} + +template +CUTE_HOST_DEVICE constexpr +auto +composition(ComposedLayout const& a, + ComposedLayout const& b) +{ + auto asb = composition(a.layout_fn(), b); + + return composition(composition(a.swizzle_fn(),asb.swizzle_fn()), asb.offset_fn(), asb.layout_fn()); +} + +// +// complement +// + +template +CUTE_HOST_DEVICE constexpr +auto +complement(ComposedLayout const& layout, CoSizeHi const& cosize_hi) +{ + // Assume there is no swizzle component in the complement + return complement(layout.layout_fn(), cosize_hi); +} + +template +CUTE_HOST_DEVICE constexpr +auto +complement(ComposedLayout const& layout) +{ + return complement(layout, cosize(layout)); +} + +// +// inverse +// + +template +CUTE_HOST_DEVICE constexpr +auto +right_inverse(ComposedLayout const& layout) +{ + CUTE_STATIC_ASSERT_V(layout.offset_fn() == Int<0>{}, "Requires 0-offset."); + return composition(right_inverse(layout.layout_fn()), layout.swizzle_fn()); +} + +template +CUTE_HOST_DEVICE constexpr +auto +left_inverse(ComposedLayout const& layout) +{ + CUTE_STATIC_ASSERT_V(layout.offset_fn() == Int<0>{}, "Requires 0-offset."); + return composition(left_inverse(layout.layout_fn()), layout.swizzle_fn()); +} + +// +// Other operations +// + +template +CUTE_HOST_DEVICE constexpr +auto +max_common_vector(ComposedLayout,Offset,SLayout> const& a, + Layout const& b) +{ + // This assumes that Offset is in the YZ domain of the Swizzle... + return cute::min(Int<(1 << M)>{}, max_common_vector(a.layout_fn(), b)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +max_common_vector(Layout const& a, + ComposedLayout,Offset,SLayout> const& b) +{ + return max_common_vector(b, a); +} + +template +CUTE_HOST_DEVICE constexpr +auto +max_common_vector(ComposedLayout,Offset0,SLayout0> const& a, + ComposedLayout,Offset1,SLayout1> const& b) +{ + auto result = coalesce(composition(a, right_inverse(b))); + + if constexpr (is_constant<1, decltype(stride<0>(result.layout_fn()))>::value) { + return shape<0>(result); + } else { + return Int<1>{}; + } + + CUTE_GCC_UNREACHABLE; +} + +template +CUTE_HOST_DEVICE constexpr +auto +zip(ComposedLayout const& a) +{ + return composition(a.swizzle_fn(), a.offset_fn(), zip(a.layout_fn())); +} + +// Partitions + +template +CUTE_HOST_DEVICE constexpr +auto +logical_divide(ComposedLayout const& a, + Tile const& b) +{ + return composition(a.swizzle_fn(), a.offset_fn(), logical_divide(a.layout_fn(), b)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tile_unzip(ComposedLayout const& a, + Tile const& b) +{ + return composition(a.swizzle_fn(), a.offset_fn(), tile_unzip(a.layout_fn(), b)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tiled_divide(ComposedLayout const& a, + Tile const& b) +{ + return composition(a.swizzle_fn(), a.offset_fn(), tiled_divide(a.layout_fn(), b)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +zipped_divide(ComposedLayout const& a, + Tile const& b) +{ + return composition(a.swizzle_fn(), a.offset_fn(), zipped_divide(a.layout_fn(), b)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +logical_product(ComposedLayout const& a, + Tile const& b) +{ + return composition(a.swizzle_fn(), a.offset_fn(), logical_product(a.layout_fn(), b)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tiled_product(ComposedLayout const& a, + Tile const& b) +{ + return composition(a.swizzle_fn(), a.offset_fn(), tiled_product(a.layout_fn(), b)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +blocked_product(ComposedLayout const& a, + Tile const& b) +{ + return composition(a.swizzle_fn(), a.offset_fn(), blocked_product(a.layout_fn(), b)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +raked_product(ComposedLayout const& a, + Tile const& b) +{ + return composition(a.swizzle_fn(), a.offset_fn(), raked_product(a.layout_fn(), b)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tile_to_shape(ComposedLayout const& layout, + Shape const& trg_shape, + ModeOrder const& ord_shape = {}) +{ + return composition(layout.swizzle_fn(), layout.offset_fn(), tile_to_shape(layout.layout_fn(), trg_shape, ord_shape)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +filter(ComposedLayout const& layout, Shape const& trg_profile) +{ + return composition(layout.swizzle_fn(), layout.offset_fn(), filter(layout.layout_fn(), trg_profile)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +coalesce(ComposedLayout const& layout) +{ + return composition(layout.swizzle_fn(), layout.offset_fn(), coalesce(layout.layout_fn())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +coalesce(ComposedLayout const& layout, Shape const& trg_profile) +{ + return composition(layout.swizzle_fn(), layout.offset_fn(), coalesce(layout.layout_fn(), trg_profile)); +} + +/////////////////////////////////////////////////////////////////////////////// +// ComposedLayout as second argument is often more difficult... + +template +CUTE_HOST_DEVICE constexpr +auto +logical_product(Layout const& block, + ComposedLayout,Offset,LayoutT> const& tile) +{ + CUTE_STATIC_ASSERT_V(tile.offset_fn() == Int<0>{}, "Require Swizzle offset == 0."); + // The new layout -- if swizzle wasn't an issue, this is the result + // our goal is to determine a new swizzle for these strides + auto new_layout = logical_product(block, tile.layout_fn()); + + // This is accomplished by identifying + // S o L :=: S? o L* + // We identify the "active" portion of S by computing (P o L)(c*) where P is a projection generated by S + // Then that active identifier is transformed through the layouts: + // L*(L[(P o L)(c*)]) + // which is a new swizzle identifier for S?, the new swizzle + + // Projections of the swizzle layout for composition, P + auto swizzle_only_zy = make_layout(make_shape (Int<(1 << M)>{}, Int<(1 << B)>{}, Int<(1 << (abs(S)-B))>{}, Int<(1 << B )>{}, Int<1>{}), + make_stride( Int<0>{}, Int<(1 << M)>{}, Int<0>{}, Int<(1 << (M+abs(S)))>{}, Int<0>{})); + + // Compose with the tile to get the swizzle projection, P o L [The Z and Y contributing portions of L] + auto layout_only_zy = composition(swizzle_only_zy, tile.layout_fn()); + // Transform the end coordinate to get the active bits of the swizzle, (P o L)(c*) + auto swizzle_active_bits = layout_only_zy(size(layout_only_zy)-Int<1>{}); + // Get the Z bit and the Y bits + auto active_Z = swizzle_active_bits & typename Swizzle::zzz_msk{}; + auto active_Y = swizzle_active_bits & typename Swizzle::yyy_msk{}; + + // Pass the identifiers through the old layout and new layout to make a new swizzle identifier, L*(L[(P o L)(c*)]) + auto new_active_Z = new_layout(Int<0>{}, tile.layout_fn()[active_Z]); + auto new_active_Y = new_layout(Int<0>{}, tile.layout_fn()[active_Y]); + + // Use this new swizzle identifier to construxt the new swizzle for new_layout + // (this also makes sure it's a "valid" swizzle that Swizzle can represent) + return composition(make_swizzle(), new_layout); +} + +template +CUTE_HOST_DEVICE constexpr +auto +tiled_product(Layout const& block, + ComposedLayout const& tile) +{ + /// Avoid swizzle slice + auto result = logical_product(block, tile); + return composition(result.swizzle_fn(), result.offset_fn(), result.layout_fn()(_, repeat>(_))); +} + +template +CUTE_HOST_DEVICE constexpr +auto +blocked_product(Layout const& block, + ComposedLayout const& layout) +{ + constexpr int R = cute::max(rank_v, rank_v); + auto padded_block = append(block, Layout<_1,_0>{}); + auto padded_layout = append(layout, Layout<_1,_0>{}); + + auto result = logical_product(padded_block, padded_layout); + + return composition(result.swizzle_fn(), + result.offset_fn(), + coalesce(zip(get<0>(result.layout_fn()), get<1>(result.layout_fn())), repeat(Int<1>{}))); +} + +// +// Upcast and Downcast +// + +template +CUTE_HOST_DEVICE constexpr +auto +upcast(ComposedLayout const& layout) +{ + return composition(upcast(layout.swizzle_fn()), upcast(layout.offset_fn()), upcast(layout.layout_fn())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +downcast(ComposedLayout const& layout) +{ + return composition(downcast(layout.swizzle_fn()), downcast(layout.offset_fn()), downcast(layout.layout_fn())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +recast(ComposedLayout const& layout) +{ + if constexpr (sizeof(NewType) == sizeof(OldType)) { + return layout; + } else if constexpr (sizeof(NewType) > sizeof(OldType)) { + static_assert(sizeof(NewType) % sizeof(OldType) == 0, "NewType must be a multiple of OldType"); + return upcast(layout); + } else if constexpr (sizeof(NewType) < sizeof(OldType)) { + static_assert(sizeof(OldType) % sizeof(NewType) == 0, "NewType must be a divisor of OldType"); + return downcast(layout); + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Display utilities +// + +template +CUTE_HOST_DEVICE void print(ComposedLayout const& layout) +{ + print(layout.swizzle_fn()); print(" o "); print(layout.offset_fn()); print(" o "); print(layout.layout_fn()); +} + +#if !defined(__MACACC_RTC__) +template +CUTE_HOST std::ostream& operator<<(std::ostream& os, ComposedLayout const& layout) +{ + return os << layout.swizzle_fn() << " o " << layout.offset_fn() << " o " << layout.layout_fn(); +} +#endif + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/swizzle_ptr.hpp b/csrc/mctlass/include/cute/swizzle_ptr.hpp new file mode 100644 index 0000000..69cce6c --- /dev/null +++ b/csrc/mctlass/include/cute/swizzle_ptr.hpp @@ -0,0 +1,284 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +#include +#include +#include + +#include +#include +#include + +/* This implements a swizzle pointer of the form + * InvolutionFn o PtrAdd + * where the InvolutionFn need not be linear. + * + * This differs subtly from swizzle_layout because the smem pointer is used + * as the offset. That means that swizzle_layout will implement position-independent + * swizzle layouts, while swizzle_ptr implements position-dependent swizzle tensors. + * Arch chose to design hardware with position-dependent swizzles. + * + * For clarity: + * NormalLayout : DeRef <- PtrAdd <- [Layout] + * ComposedLayout: DeRef <- PtrAdd <- [Swizzle <- OffsetAdd <- Layout] + * SwizzlePtr : [DeRef <- Swizzle <- PtrAdd] <- Layout + * + * Furthermore, for known swizzles, this pointer attempts to decay itself + * to a normal-pointer with a new layout containing dynamic or static strides. + * This is possible by determining the subdomain of the InvolutionFn + * that is identity and testing if the Layout's codomain is contained + * within it. + */ + +namespace cute +{ + +template +struct smem_ptr_swizzle +{ + static_assert(is_empty::value, "Swizzle can't have state."); + + CUTE_HOST_DEVICE constexpr + T* get() const + { + return ptr_; + } + + CUTE_HOST_DEVICE constexpr static + Swizzle get_swizzle() + { + return {}; + } + + CUTE_HOST_DEVICE constexpr static + T* apply_swizzle(T* ptr) + { + return reinterpret_cast(Swizzle::apply(reinterpret_cast(ptr))); + } + + CUTE_HOST_DEVICE constexpr + T& operator*() const + { + return *apply_swizzle(get()); + } + + template + CUTE_HOST_DEVICE constexpr + T& operator[](Int const& i) const + { + return *apply_swizzle(get() + i); + } + + template + CUTE_HOST_DEVICE constexpr + smem_ptr_swizzle operator+(Int const& i) const + { + return {ptr_ + i}; + } + + T* ptr_; +}; + +template +struct is_smem> : true_type {}; + +// Make a swizzle pointer +template +CUTE_HOST_DEVICE constexpr +auto +make_smem_ptr(T* ptr, Swizzle const& swizzle) +{ + return smem_ptr_swizzle{ptr}; +} + +// A model of a nullptr smem_ptr with B == sizeof_bits::value +// That represents an unset pointer. This is a placeholder type that is waiting for an smem_ptr +template +struct smem_ptr_flag_bits : Int<0> {}; + +using smem_ptr_flag = smem_ptr_flag_bits<1>; + +// A flagged construction method to transform ComposedLayout +// Make a swizzle pointer tensor and check that the intended type size matches +template +CUTE_HOST_DEVICE constexpr +auto +make_tensor(smem_ptr const& ptr, + ComposedLayout,Layout> const& layout) +{ + static_assert(B == sizeof_bits::value, "Expected a B-bit pointer type."); + return make_tensor(make_smem_ptr(ptr.get(), layout.swizzle_fn()), + layout.layout_fn()); +} + +// Specialization for immediate decay +template +CUTE_HOST_DEVICE constexpr +auto +make_tensor(smem_ptr_swizzle>& p, Layout const& layout) +{ + return make_tensor(make_smem_ptr(p.ptr_), layout); +} + +template +CUTE_HOST_DEVICE constexpr +auto +make_tensor(smem_ptr_swizzle> const& p, Layout const& layout) +{ + return make_tensor(make_smem_ptr(p.ptr_), layout); +} + +// NOTE: To preserve smem_ptr_flag_bits under recast ops +template +CUTE_HOST_DEVICE constexpr +auto +upcast(ComposedLayout,Layout> const& layout) +{ + return composition(layout.swizzle_fn(), smem_ptr_flag_bits{}, upcast(layout.layout_fn())); +} + +template +CUTE_HOST_DEVICE constexpr +auto +downcast(ComposedLayout,Layout> const& layout) +{ + return composition(layout.swizzle_fn(), smem_ptr_flag_bits{}, downcast(layout.layout_fn())); +} + +// +// Recast +// Swizzle operates on the pointer address, so it doesn't care about the type +// + +template +CUTE_HOST_DEVICE constexpr +auto +recast(smem_ptr_swizzle const& ptr) +{ + return smem_ptr_swizzle{recast(ptr.ptr_)}; +} + +template +CUTE_HOST_DEVICE constexpr +auto +recast(smem_ptr_swizzle const& ptr) +{ + return smem_ptr_swizzle{recast(ptr.ptr_)}; +} + +// +// Conversion with swizzle_layout +// + +template +CUTE_HOST_DEVICE +auto +as_position_independent_swizzle_layout(ComposedLayout,Layout> const& layout) +{ + return composition(recast,uint_bit_t>(layout.swizzle_fn()), Int<0>{}, layout.layout_fn()); +} + +template +CUTE_HOST_DEVICE +auto +as_position_independent_swizzle_tensor(Tensor>, Layout> const& tensor) +{ + { + uint32_t address = cast_smem_ptr_to_uint(tensor.data().get()); + uint32_t mask = ((uint32_t(1) << Swizzle::num_base) - 1) & (Swizzle::swizzle_code); + assert((address & mask) == 0); // Alignment to the Base, Z, and Y of Swizzle + } + auto new_swizzle = recast,uint_bit_t>>(tensor.data().get_swizzle()); + return make_tensor(make_smem_ptr(tensor.data().get()), composition(new_swizzle, Int<0>{}, tensor.layout())); +} + +template +CUTE_HOST_DEVICE +auto +as_position_independent_swizzle_tensor(Tensor>, Layout>& tensor) +{ + { + uint32_t address = cast_smem_ptr_to_uint(tensor.data().get()); + uint32_t mask = ((uint32_t(1) << Swizzle::num_base) - 1) & (Swizzle::swizzle_code); + assert((address & mask) == 0); // Alignment to the Base, Z, and Y of Swizzle + } + auto new_swizzle = recast,uint_bit_t>>(tensor.data().get_swizzle()); + return make_tensor(make_smem_ptr(tensor.data().get()), composition(new_swizzle, Int<0>{}, tensor.layout())); +} + +template +CUTE_HOST_DEVICE +auto +as_position_independent_swizzle_tensor(Tensor>, Layout>&& tensor) +{ + return as_position_independent_swizzle_tensor(tensor); +} + +// +// Print +// + +// Capture and cast smem_ptr_flag Layouts to offset-0 layouts +template +CUTE_HOST_DEVICE +void +print_latex(ComposedLayout,Layout> const& layout) +{ + auto new_swizzle = recast,uint_bit_t>(layout.swizzle_fn()); + print_latex(composition(new_swizzle, Int<0>{}, layout.layout_fn())); +} + +template +CUTE_HOST_DEVICE void print(smem_ptr_flag_bits const& ptr) +{ + printf("smem_ptr_%db(unset)", B); +} + +template +CUTE_HOST_DEVICE void print(smem_ptr_swizzle> const& ptr) +{ + printf("smem_ptr_S<%d,%d,%d>_%db(%p)", B, M, S, int(8*sizeof(T)), ptr.get()); +} + +#if !defined(__MACACC_RTC__) +template +CUTE_HOST std::ostream& operator<<(std::ostream& os, smem_ptr_swizzle> const&) +{ + return os << "smem_ptr_S<" << B << "," << M << "," << S << ">_" << int(8*sizeof(T)) << "b"; +} +#endif + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/tensor.hpp b/csrc/mctlass/include/cute/tensor.hpp new file mode 100644 index 0000000..f7882cf --- /dev/null +++ b/csrc/mctlass/include/cute/tensor.hpp @@ -0,0 +1,956 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace cute +{ + +// +// Engine -- owning or non-owning data store +// + +// concept Engine { +// using value_type = ; +// iterator begin(); +// }; + +template +using ArrayEngine = typename conditional<(sizeof_bits::value % 8 == 0), + array_aligned, + array_subbyte>::type; + +template +struct ViewEngine +{ + using value_type = typename cute::remove_cvref())>::type; + + using iterator = Iterator; + iterator storage_; + + CUTE_HOST_DEVICE constexpr + iterator const& + begin() const { + return storage_; + } + + CUTE_HOST_DEVICE constexpr + iterator& + begin() { + return storage_; + } +}; + +template +struct is_rmem> : is_rmem {}; +template +struct is_smem> : is_smem {}; +template +struct is_gmem> : is_gmem {}; +template +struct ConstViewEngine +{ + using value_type = typename cute::remove_cvref())>::type; + + using iterator = Iterator; + iterator storage_; + + CUTE_HOST_DEVICE constexpr + iterator const& + begin() const { + return storage_; + } +}; + +template +struct is_rmem> : is_rmem {}; +template +struct is_smem> : is_smem {}; +template +struct is_gmem> : is_gmem {}; +// +// Tensor +// + +template +struct Tensor +{ + using value_type = typename Engine::value_type; + //using pointer = typename engine_traits::pointer; + //using const_pointer = typename engine_traits::const_pointer; + //using reference = typename engine_traits::reference; + //using const_reference = typename engine_traits::const_reference; + + using engine_type = Engine; + using layout_type = Layout; + + CUTE_HOST_DEVICE constexpr + Tensor() {} + + template + CUTE_HOST_DEVICE constexpr + Tensor(Ptr const& ptr, Layout const& layout) + : rep_(layout, ptr) { + } + + // + // Accessors + // + + static constexpr int rank = Layout::rank; + + CUTE_HOST_DEVICE constexpr + decltype(auto) + tensor() const { + return *this; + } + + CUTE_HOST_DEVICE constexpr + decltype(auto) + layout() const { + return get<0>(rep_); + } + + CUTE_HOST_DEVICE constexpr + decltype(auto) + engine() const { + return get<1>(rep_); + } + + CUTE_HOST_DEVICE constexpr + decltype(auto) + engine() { + return get<1>(rep_); + } + + CUTE_HOST_DEVICE constexpr + decltype(auto) + data() const { + return engine().begin(); + } + + CUTE_HOST_DEVICE constexpr + decltype(auto) + data() { + return engine().begin(); + } + + CUTE_HOST_DEVICE constexpr + decltype(auto) + shape() const { + return layout().shape(); + } + + CUTE_HOST_DEVICE constexpr + auto + size() const { + return cute::size(shape()); + } + + CUTE_HOST_DEVICE constexpr + decltype(auto) + stride() const { + return layout().stride(); + } + + // + // Indexing op() and op[] + // + + // Index into this tensor like an array by computing the offset via layout() + template + CUTE_HOST_DEVICE constexpr + decltype(auto) + operator[](Coord const& coord) { + return data()[layout()(coord)]; + } + + template + CUTE_HOST_DEVICE constexpr + decltype(auto) + operator[](Coord const& coord) const { + return data()[layout()(coord)]; + } + + template + CUTE_HOST_DEVICE constexpr + decltype(auto) + operator()(Coord const& coord) { + if constexpr (has_underscore::value) { + auto const& [sliced_layout,offset] = slice_and_offset(coord, layout()); + return make_tensor(data() + offset, sliced_layout); + } else { + return data()[layout()(coord)]; + } + + CUTE_GCC_UNREACHABLE; + } + + template + CUTE_HOST_DEVICE constexpr + decltype(auto) + operator()(Coord const& coord) const { + if constexpr (has_underscore::value) { + auto const& [sliced_layout,offset] = slice_and_offset(coord, layout()); + return make_tensor(data() + offset, sliced_layout); + } else { + return data()[layout()(coord)]; + } + + CUTE_GCC_UNREACHABLE; + } + + // op() convenience function for multi-dimensional coordinates + template + CUTE_HOST_DEVICE constexpr + decltype(auto) + operator()(Coord0 const& c0, Coord1 const& c1, Coords const&... cs) { + return operator()(make_coord(c0,c1,cs...)); + } + + template + CUTE_HOST_DEVICE constexpr + decltype(auto) + operator()(Coord0 const& c0, Coord1 const& c1, Coords const&... cs) const { + return operator()(make_coord(c0,c1,cs...)); + } + + // + // Compose + // + + template + CUTE_HOST_DEVICE constexpr + auto + compose(Layouts const&... layouts) { + return make_tensor(data(), layout().compose(layouts...)); + } + + template + CUTE_HOST_DEVICE constexpr + auto + compose(Layouts const&... layouts) const { + return make_tensor(data(), layout().compose(layouts...)); + } + + // + // Tile + // + + template + CUTE_HOST_DEVICE constexpr + auto + tile(Layouts const&... layouts) { + return make_tensor(data(), layout().tile(layouts...)); + } + + template + CUTE_HOST_DEVICE constexpr + auto + tile(Layouts const&... layouts) const { + return make_tensor(data(), layout().tile(layouts...)); + } + + // + // Utility + // + + template ::value)> + CUTE_HOST_DEVICE constexpr + auto + get_1d_coord(Int const& linear_idx) const { + return layout().get_1d_coord(linear_idx); + } + + template ::value)> + CUTE_HOST_DEVICE constexpr + auto + get_hier_coord(Int const& linear_idx) const { + return layout().get_hier_coord(linear_idx); + } + + template ::value)> + CUTE_HOST_DEVICE constexpr + auto + get_flat_coord(Int const& linear_idx) const { + return layout().get_flat_coord(linear_idx); + } + + cute::tuple rep_; +}; + + +template +struct is_tensor : false_type {}; +template +struct is_tensor> : true_type {}; + +template +struct is_rmem> : is_rmem {}; +template +struct is_smem> : is_smem {}; +template +struct is_gmem> : is_gmem {}; +// Customization point for creation of owning and non-owning Tensors +template +struct MakeTensor +{ + template ::value && + is_layout::value)> + CUTE_HOST_DEVICE constexpr auto + operator()(Layout const& layout) const + { + static_assert(is_static::value, "Dynamic owning tensors not supported"); + using Engine = ArrayEngine>; + return Tensor(); + } + + template ::value && + is_layout::value)> + CUTE_HOST_DEVICE constexpr auto + operator()(T const& iter, Layout const& layout) + { + using Engine = ViewEngine; + return Tensor(iter, layout); + } + + template ::value)> + CUTE_HOST_DEVICE constexpr auto + operator()(LayoutArg const& arg, LayoutArgs const&... args) const + { + return operator()(make_layout(arg, args...)); + } + + template ::value)> + CUTE_HOST_DEVICE constexpr auto + operator()(T const& iter, LayoutArg const& arg, LayoutArgs const&... args) + { + return operator()(iter, make_layout(arg, args...)); + } +}; + +// +// make_tensor +// + +// Make an owning Tensor that will allocate a static array +// e.g. make_tensor(Int<12>{}) +template +CUTE_HOST_DEVICE constexpr +auto +make_tensor(Args const&... args) +{ + return MakeTensor{}(args...); +} + +// Make a non-owning Tensor that will use a pointer (view) +// e.g. make_tensor(vec.data(), 12) +template +CUTE_HOST_DEVICE constexpr +auto +make_tensor(Iterator const& iter, Args const&... args) +{ + return MakeTensor{}(iter, args...); +} + +// +// make_tensor_like +// Make a register tensor the same type and shape and (if possible) order as another tensor +// + +template +CUTE_HOST_DEVICE constexpr +auto +make_tensor_like(Layout const& layout) +{ + if constexpr (is_static::value) { + return make_tensor(make_ordered_layout(layout)); + } else { + return make_tensor(make_layout(layout.shape())); + } +} + +template +CUTE_HOST_DEVICE constexpr +auto +make_tensor_like(Tensor const& tensor) +{ + return make_tensor_like(tensor.layout()); +} + +template +CUTE_HOST_DEVICE constexpr +auto +make_tensor_like(Tensor const& tensor) +{ + return make_tensor_like(tensor.layout()); +} + +// +// make_fragment_like -- +// Make a tensor the same shape and (if possible) order as another tensor, with special +// consideration of the 0th mode. The 0th mode is commonly used for MMA_Atoms or Copy_Atoms +// so this allocates the 0th mode with LayoutLeft regardless of the reference layout. +// + +template +CUTE_HOST_DEVICE constexpr +auto +make_fragment_like(Layout const& layout) +{ + return make_tensor(make_fragment_like(layout)); +} + +template +CUTE_HOST_DEVICE constexpr +auto +make_fragment_like(Tensor const& tensor) +{ + return make_fragment_like(tensor.layout()); +} + +template +CUTE_HOST_DEVICE constexpr +auto +make_fragment_like(Tensor const& tensor) +{ + return make_fragment_like(tensor.layout()); +} + +// +// make_identity_tensor +// + +template +CUTE_HOST_DEVICE constexpr +auto +make_identity_tensor(Shape const& shape) +{ + return make_tensor(ArithmeticTupleIterator(as_arithmetic_tuple(repeat_like(shape, Int<0>{}))), + make_identity_layout(shape)); +} + +// +// Utilities +// + +// Return the subtensor of a mode +template >::value)> +CUTE_HOST_DEVICE constexpr +decltype(auto) +tensor(Tensor&& tensor) +{ + return std::forward(tensor); +} + +template >::value)> +CUTE_HOST_DEVICE constexpr +decltype(auto) +tensor(Tensor&& tensor) +{ + return make_tensor(std::forward(tensor).data(), get(tensor.layout())); +} + +// Return the subtensor of a range of modes +template >::value)> +CUTE_HOST_DEVICE constexpr +decltype(auto) +take(Tensor&& tensor) +{ + return make_tensor(std::forward(tensor).data(), take(tensor.layout())); +} + +// Return the layout of a mode +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +layout(Tensor const& tensor) +{ + return layout(tensor.layout()); +} + +// Return the shape of a mode +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +shape(Tensor const& tensor) +{ + return shape(tensor.layout()); +} + +// Return the stride of a mode +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +stride(Tensor const& tensor) +{ + return stride(tensor.layout()); +} + +// Return the number of elements in a mode +template +CUTE_HOST_DEVICE constexpr +decltype(auto) +size(Tensor const& tensor) +{ + return size(tensor.layout()); +} + +// Return the rank of a mode +template +CUTE_HOST_DEVICE constexpr +auto +rank(Tensor const& tensor) +{ + return rank(tensor.layout()); +} + +// Return the depth of a mode +template +CUTE_HOST_DEVICE constexpr +auto +depth(Tensor const& tensor) +{ + return depth(tensor.layout()); +} + +// +// Operations to manipulate Tensors like a Layout +// + +template >::value)> +CUTE_HOST_DEVICE constexpr +auto +flatten(Tensor&& tensor) +{ + return make_tensor(std::forward(tensor).data(), flatten(tensor.layout())); +} + +template >::value)> +CUTE_HOST_DEVICE constexpr +auto +coalesce(Tensor&& tensor) +{ + return make_tensor(std::forward(tensor).data(), coalesce(tensor.layout())); +} + +template >::value)> +CUTE_HOST_DEVICE constexpr +auto +coalesce(Tensor&& tensor, Profile const& profile) +{ + return make_tensor(std::forward(tensor).data(), coalesce(tensor.layout(), profile)); +} + +// Group the modes [B,E) into a single mode +// e.g. group<2,4>(make_tensor(Layout>{})) +// => make_tensor(Layout,_5,_6>>{}) +template >::value)> +CUTE_HOST_DEVICE constexpr +auto +group_modes(Tensor&& tensor) +{ + return make_tensor(std::forward(tensor).data(), + group(tensor.layout())); +} + +// +// Recast +// + +// NOTE: This is very dangerous to do +// -- doesn't check dynamic integer divisibility +// -- doesn't check alignment + +// A tagged version for dispatching +template >::value)> +CUTE_HOST_DEVICE constexpr +auto +recast(Tensor&& tensor, type_list) +{ + using OldType = typename remove_cvref_t::value_type; + auto old_layout = tensor.layout(); + auto new_layout = recast(old_layout); + + // If this is an upcast of a normal Layout with static negative strides, then offset as well + if constexpr (sizeof(OldType) < sizeof(NewType) && not is_composed_layout::value) { + auto shape_diff = transform(flatten(old_layout.shape()), flatten(new_layout.shape()), minus{}); + auto extent_diff = transform(shape_diff, flatten(old_layout.stride()), multiplies{}); + auto offset = fold(extent_diff, Int<0>{}, [](auto const& i, auto const& a) { return i + cute::min(a,Int<0>{}); }); + + return make_tensor(recast(std::forward(tensor).data() + offset), new_layout); + } else { + return make_tensor(recast(std::forward(tensor).data() ), new_layout); + } + + CUTE_GCC_UNREACHABLE; +} + +template >::value)> +CUTE_HOST_DEVICE constexpr +auto +recast(Tensor&& tensor) +{ + return recast(std::forward(tensor), type_list{}); +} + +// +// max_common_vector +// + +/* Return Int such that N is the maximum number of continguous elements + * that logically correspond in the tensors of @a a and @a b. This is, + * the number of elements that could reasonably be vectorized into a single load/store. + * + * @returns Int with N >= 0 + * + * A return value of Int<0> indicates that no such conclusion can be made and no + * vectorization should be attempted. + */ +template +CUTE_HOST_DEVICE constexpr +auto +max_common_vector(Tensor const& a, + Tensor const& b) +{ + using SrcType = typename Tensor::value_type; + using DstType = typename Tensor::value_type; + + using SrcRef = decltype(*(a.data())); + using DstRef = decltype(*(b.data())); + + // Determine if vectorization candidates at all + if constexpr (// Should be the same value_types, else the copy is also performing a cast + sizeof(SrcType) == sizeof(DstType) && + // The types should be trivially copyable so that vectorization is valid + is_trivially_copyable::value && + is_trivially_copyable::value && + // Should be load/storing real data, rather than implicit iterators or such + is_reference::value && + is_reference::value) + { + return max_common_vector(a.layout(), b.layout()); + } else { + return Int<0>{}; + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Key algebraic operations +// + +template >::value)> +CUTE_HOST_DEVICE constexpr +auto +logical_divide(Tensor && tensor, + Tile const& tile) +{ + return make_tensor(std::forward(tensor).data(), + logical_divide(tensor.layout(), tile)); +} + +// zipped_divide is logical_divide with modes gathered into standard form ((BLK_A,BLK_B),(a,b)) +template >::value)> +CUTE_HOST_DEVICE constexpr +auto +zipped_divide(Tensor && tensor, + Tile const& tile) // Layout or Tile +{ + return make_tensor(std::forward(tensor).data(), + zipped_divide(tensor.layout(), tile)); +} + +// tiled_divide is logical_divide with the second output mode flattened ((BLK_A,BLK_B),a,b) +template >::value)> +CUTE_HOST_DEVICE constexpr +auto +tiled_divide(Tensor && tensor, + Tile const& tile) // Layout or Tile +{ + return make_tensor(std::forward(tensor).data(), + tiled_divide(tensor.layout(), tile)); +} + +// logical_product on a Tensor doesn't make sense since it often increases cosize + +// +// Logicial Divide utilities: local_partition and local_tile +// + +template >::value)> +CUTE_HOST_DEVICE constexpr +auto +local_partition(Tensor && tensor, + Tile const& tile, + Coord const& coord) +{ + constexpr int R1 = decltype(rank(tensor))::value; + + // Split the modes of tensor according to the modes of tile + // zipped_divide returns something like ((VEC_A,VEC_B,...),(a,b,...)) + + // The_coord is the coord into the first mode, flatten the rest + return zipped_divide(std::forward(tensor), tile)(coord, repeat(_)); +} + +template >::value)> +CUTE_HOST_DEVICE constexpr +auto +local_partition(Tensor && tensor, + Tile const& tile, + Coord const& coord, + Projection const& proj) +{ + return local_partition(std::forward(tensor), + dice(proj, tile), + dice(proj, coord)); +} + +// Special case with Layout and Integral that extracts the coord first +// e.g. local_partition(tensor, ThrLayout, threadIdx.x) +template >::value && + is_integral::value)> +CUTE_HOST_DEVICE +auto +local_partition(Tensor && tensor, + Layout const& tile, + Index const& index) +{ + return local_partition(std::forward(tensor), + product_each(shape(tile)), + tile.get_flat_coord(index)); +} + +// Special case with Layout and Integral that extracts the coord first +// e.g. local_partition(tensor, ThrLayout, threadIdx.x, Step<_1,X,_1>{}) +template >::value && + is_integral::value)> +CUTE_HOST_DEVICE +auto +local_partition(Tensor && tensor, + Layout const& tile, + Index const& index, + Projection const& proj) +{ + return local_partition(std::forward(tensor), + dice(proj, product_each(shape(tile))), + dice(proj, tile).get_flat_coord(index)); +} + +template >::value)> +CUTE_HOST_DEVICE constexpr +auto +local_tile(Tensor && tensor, + Tile const& tile, + Coord const& coord) +{ + constexpr int R0 = decltype(rank(tile))::value; + constexpr int R1 = decltype(rank(tensor))::value; + + // Split the modes of tensor according to the modes of tile + // zipped_divide returns something like ((VEC_A,VEC_B,...),(a,b,...)) + + // The padded_coord is the coord into the second mode, flatten the rest + return zipped_divide(std::forward(tensor), tile)(repeat(_), append(coord,_)); +} + +template >::value)> +CUTE_HOST_DEVICE +auto +local_tile(Tensor && tensor, + Tile const& tile, + Coord const& coord, + Proj const& proj) +{ + return local_tile(std::forward(tensor), + dice(proj, tile), + dice(proj, coord)); +} + +// +// Display utilities +// + +template +CUTE_HOST_DEVICE void print_tensor(Tensor const& tensor) +{ + auto format = get_format(tensor(0)); + using type = typename decltype(format)::type; + + if constexpr (Layout::rank == 1) + { + for (int m = 0; m < size(tensor); ++m) { + printf(format.format, format.digits, type(tensor(m))); + printf("\n"); + } + } else + if constexpr (Layout::rank == 2) + { + for (int m = 0; m < size<0>(tensor); ++m) { + for (int n = 0; n < size<1>(tensor); ++n) { + printf(format.format, format.digits, type(tensor(m,n))); + } + printf("\n"); + } + } else + if constexpr (Layout::rank == 3) + { + print_tensor(tensor(_,_,0)); + for (int k = 1; k < size<2>(tensor); ++k) { + for (int i = 0; i < format.digits*size<1>(tensor); ++i) { print("-"); } print("\n"); + print_tensor(tensor(_,_,k)); + } + } else + if constexpr (Layout::rank == 4) + { + print_tensor(tensor(_,_,_,0)); + for (int p = 1; p < size<3>(tensor); ++p) { + for (int i = 0; i < format.digits*size<1>(tensor); ++i) { print("="); } print("\n"); + print_tensor(tensor(_,_,_,p)); + } + } +} + +template +CUTE_HOST_DEVICE void print(Tensor const& tensor) +{ + print(tensor.layout()); print("\n"); + print_tensor(tensor); +} + +#if !defined(__MACACC_RTC__) +template +CUTE_HOST std::ostream& print_tensor_os(std::ostream& os, Tensor const& tensor) +{ + int digits = 9; + + if constexpr (Layout::rank == 1) + { + for (int m = 0; m < size(tensor); ++m) { + os << std::setw(digits) << tensor(m) << std::endl; + } + } else + if constexpr (Layout::rank == 2) + { + for (int m = 0; m < size<0>(tensor); ++m) { + for (int n = 0; n < size<1>(tensor); ++n) { + os << std::setw(digits) << tensor(m,n); + } + os << std::endl; + } + } else + if constexpr (Layout::rank == 3) + { + print_tensor_os(os, tensor(_,_,0)); + for (int k = 1; k < size<2>(tensor); ++k) { + for (int i = 0; i < digits*size<1>(tensor); ++i) { os << "-"; } os << std::endl; + print_tensor_os(os, tensor(_,_,k)); + } + } else + if constexpr (Layout::rank == 4) + { + print_tensor_os(os, tensor(_,_,_,0)); + for (int p = 1; p < size<3>(tensor); ++p) { + for (int i = 0; i < digits*size<1>(tensor); ++i) { os << "="; } os << std::endl; + print_tensor_os(os, tensor(_,_,_,p)); + } + } + + return os; +} + +template +CUTE_HOST std::ostream& operator<<(std::ostream& os, Tensor const& tensor) +{ + os << tensor.layout() << std::endl; + return print_tensor_os(os, tensor); +} +#endif // !defined(__MACACC_RTC__) + +} // end namespace cute + +// +// Extended Engines +// + +#include + +// +// Tensor Algorithms +// + +#include +#include +#include +#include +#include +#include diff --git a/csrc/mctlass/include/cute/tensor_predicate.hpp b/csrc/mctlass/include/cute/tensor_predicate.hpp new file mode 100644 index 0000000..730f219 --- /dev/null +++ b/csrc/mctlass/include/cute/tensor_predicate.hpp @@ -0,0 +1,63 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +namespace cute +{ + +template +struct ConstantTensor +{ + template + CUTE_HOST_DEVICE constexpr + T const& + operator()(Coords const&...) const { + return val_; + } + + T val_; +}; + +struct TrivialPredTensor +{ + template + CUTE_HOST_DEVICE constexpr + true_type + operator()(Coords const&...) const { + return {}; + } +}; + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/tile.hpp b/csrc/mctlass/include/cute/tile.hpp new file mode 100644 index 0000000..b2fa2e8 --- /dev/null +++ b/csrc/mctlass/include/cute/tile.hpp @@ -0,0 +1,58 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include + +namespace cute +{ + +// +// A Tile is not a Layout, it's a tuple of Layouts or Tiles or Underscores +// + +template +using Tile = tuple; + +template +using is_tile = is_tuple; + +template +CUTE_HOST_DEVICE constexpr +auto +make_tile(Layouts const&... layouts) +{ + return Tile(layouts...); +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/underscore.hpp b/csrc/mctlass/include/cute/underscore.hpp new file mode 100644 index 0000000..19a8967 --- /dev/null +++ b/csrc/mctlass/include/cute/underscore.hpp @@ -0,0 +1,150 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include +#include +#include + +namespace cute +{ + +// For slicing +struct Underscore : Int<0> {}; + +CUTE_INLINE_CONSTANT Underscore _; + +// Treat Underscore as an integral like integral_constant +template <> +struct is_integral : true_type {}; + +template +struct is_underscore : false_type {}; +template <> +struct is_underscore : true_type {}; + +// Tuple trait for detecting static member element +template +struct has_elem : false_type {}; +template +struct has_elem : true_type {}; +template +struct has_elem::value> > + : has_elem > {}; +template +struct has_elem> + : disjunction, Elem>...> {}; + +// Tuple trait for detecting static member element +template +struct all_elem : false_type {}; +template +struct all_elem : true_type {}; +template +struct all_elem::value> > + : all_elem > {}; +template +struct all_elem> + : conjunction, Elem>...> {}; + +// Tuple trait for detecting Underscore member +template +using has_underscore = has_elem; + +template +using all_underscore = all_elem; + +template +using has_int1 = has_elem>; + +template +using has_int0 = has_elem>; + +// +// Slice keeps only the elements of Tuple B that are paired with an Underscore +// + +template +CUTE_HOST_DEVICE constexpr +auto +slice(A const& a, B const& b) +{ + if constexpr (is_tuple::value) { + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + return filter_tuple(a, b, [](auto const& x, auto const& y) { return slice(x,y); }); + } else if constexpr (is_underscore::value) { + return cute::tuple{b}; + } else { + return cute::tuple<>{}; + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Dice keeps only the elements of Tuple B that are paired with an Int +// + +template +CUTE_HOST_DEVICE constexpr +auto +dice(A const& a, B const& b) +{ + if constexpr (is_tuple::value) { + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + return filter_tuple(a, b, [](auto const& x, auto const& y) { return dice(x,y); }); + } else if constexpr (is_underscore::value) { + return cute::tuple<>{}; + } else { + return cute::tuple{b}; + } + + CUTE_GCC_UNREACHABLE; +} + +// +// Display utilities +// + +CUTE_HOST_DEVICE void print(Underscore const&) { + printf("_"); +} + +#if !defined(__MACACC_RTC__) +CUTE_HOST std::ostream& operator<<(std::ostream& os, Underscore const&) { + return os << "_"; +} +#endif + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/util/debug.hpp b/csrc/mctlass/include/cute/util/debug.hpp new file mode 100644 index 0000000..0e6de14 --- /dev/null +++ b/csrc/mctlass/include/cute/util/debug.hpp @@ -0,0 +1,163 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +/** + * \file + * \brief Debugging and logging functionality + */ + +#include + +#include + +namespace cute +{ + +/****************************************************************************** + * Debug and logging macros + ******************************************************************************/ + +/** + * Formats and prints the given message to stdout + */ +#if !defined(CUTE_LOG) +# if !defined(__MACA_ARCH__) +# define CUTE_LOG(format, ...) printf(format, __VA_ARGS__) +# else +# define CUTE_LOG(format, ...) \ + printf("[block (%d,%d,%d), thread (%d,%d,%d)]: " format, \ + blockIdx.x, blockIdx.y, blockIdx.z, \ + threadIdx.x, threadIdx.y, threadIdx.z, \ + __VA_ARGS__); +# endif +#endif + +/** + * Formats and prints the given message to stdout only if DEBUG is defined + */ +#if !defined(CUTE_LOG_DEBUG) +# ifdef DEBUG +# define CUTE_LOG_DEBUG(format, ...) CUTE_LOG(format, __VA_ARGS__) +# else +# define CUTE_LOG_DEBUG(format, ...) +# endif +#endif + +/** + * \brief Perror macro with exit + */ +#if !defined(CUTE_ERROR_EXIT) +# define CUTE_ERROR_EXIT(e) \ + do { \ + mcError_t code = (e); \ + if (code != mcSuccess) { \ + fprintf(stderr, "<%s:%d> %s:\n %s: %s\n", \ + __FILE__, __LINE__, #e, \ + mcGetErrorName(code), mcGetErrorString(code)); \ + fflush(stderr); \ + exit(0); \ + } \ + } while (0) +#endif + +#if !defined(CUTE_CHECK_LAST) +# define CUTE_CHECK_LAST() CUTE_ERROR_EXIT(mcPeekAtLastError()); CUTE_ERROR_EXIT(mcDeviceSynchronize()) +#endif + +#if !defined(CUTE_CHECK_ERROR) +# define CUTE_CHECK_ERROR(e) CUTE_ERROR_EXIT(e) +#endif + +// A dummy function that uses compilation failure to print a type +template +CUTE_HOST_DEVICE void +print_type() { + static_assert(sizeof(T) < 0, "Printing type T."); +} +template +CUTE_HOST_DEVICE void +print_type(T&&) { + static_assert(sizeof(T) < 0, "Printing type T."); +} + +// +// Device-specific helpers +// +// e.g. +// if (thread0()) print(...); +// if (block0()) print(...); +// if (thread(42)) print(...); + +CUTE_HOST_DEVICE +bool +block(int bid) +{ +#if defined(__MACA_ARCH__) + return blockIdx.x + blockIdx.y*gridDim.x + blockIdx.z*gridDim.x*gridDim.y == bid; +#else + return true; +#endif +} + +CUTE_HOST_DEVICE +bool +thread(int tid, int bid) +{ +#if defined(__MACA_ARCH__) + return (threadIdx.x + threadIdx.y*blockDim.x + threadIdx.z*blockDim.x*blockDim.y == tid) && block(bid); +#else + return true; +#endif +} + +CUTE_HOST_DEVICE +bool +thread(int tid) +{ + return thread(tid,0); +} + +CUTE_HOST_DEVICE +bool +thread0() +{ + return thread(0,0); +} + +CUTE_HOST_DEVICE +bool +block0() +{ + return block(0); +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/util/print.hpp b/csrc/mctlass/include/cute/util/print.hpp new file mode 100644 index 0000000..320b4f5 --- /dev/null +++ b/csrc/mctlass/include/cute/util/print.hpp @@ -0,0 +1,141 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +#include +#include + +// +// CUDA compatible print and printf +// + +namespace cute +{ + +CUTE_HOST_DEVICE +int +num_digits(int x) +{ + return (x < 10 ? 1 : + (x < 100 ? 2 : + (x < 1000 ? 3 : + (x < 10000 ? 4 : + (x < 100000 ? 5 : + (x < 1000000 ? 6 : + (x < 10000000 ? 7 : + (x < 100000000 ? 8 : + (x < 1000000000 ? 9 : + 10))))))))); +} + +template +struct format_and_size { + using type = T; + char const* format; + int digits; +}; + +CUTE_HOST_DEVICE +format_and_size +get_format(bool) { + return {"%*d", 3}; +} + +CUTE_HOST_DEVICE +format_and_size +get_format(int32_t) { + return {"%*d", 5}; +} + +CUTE_HOST_DEVICE +format_and_size +get_format(uint32_t) { + return {"%*d", 5}; +} + +CUTE_HOST_DEVICE +format_and_size +get_format(int64_t) { + return {"%*d", 5}; +} + +CUTE_HOST_DEVICE +format_and_size +get_format(uint64_t) { + return {"%*d", 5}; +} + +CUTE_HOST_DEVICE +format_and_size +get_format(half_t) { + return {"%*.2f", 8}; +} + +CUTE_HOST_DEVICE +format_and_size +get_format(float) { + return {"%*.2e", 10}; +} + +CUTE_HOST_DEVICE +format_and_size +get_format(double) { + return {"%*.3e", 11}; +} + +// +// print dispatcher +// + +CUTE_HOST_DEVICE +void +print(char const& c) { + printf("%c", c); +} + +template ::value)> +CUTE_HOST_DEVICE +void +print(T const& a) { + printf("%d", int(a)); +} + +template +CUTE_HOST_DEVICE +void +print(char const* format, T const&... t) { + printf(format, t...); +} + +} // end namespace cute diff --git a/csrc/mctlass/include/cute/util/type_traits.hpp b/csrc/mctlass/include/cute/util/type_traits.hpp new file mode 100644 index 0000000..3342278 --- /dev/null +++ b/csrc/mctlass/include/cute/util/type_traits.hpp @@ -0,0 +1,207 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#if defined(__MACACC_RTC__) +#include +#include +#include +#include +#include +#else +#include +#include // tuple_size, tuple_element +#include // ptrdiff_t +#include // uintptr_t +#include // numeric_limits +#endif + +#include + +namespace cute +{ + using CUTE_STL_NAMESPACE::enable_if; + using CUTE_STL_NAMESPACE::enable_if_t; +} + +#define __CUTE_REQUIRES(...) typename cute::enable_if<(__VA_ARGS__)>::type* = nullptr +#define __CUTE_REQUIRES_V(...) typename cute::enable_if::type* = nullptr + +namespace cute +{ + +// +using CUTE_STL_NAMESPACE::conjunction; +using CUTE_STL_NAMESPACE::conjunction_v; + +using CUTE_STL_NAMESPACE::disjunction; +using CUTE_STL_NAMESPACE::disjunction_v; + +using CUTE_STL_NAMESPACE::negation; +using CUTE_STL_NAMESPACE::negation_v; + +using CUTE_STL_NAMESPACE::void_t; +using CUTE_STL_NAMESPACE::is_void_v; + +using CUTE_STL_NAMESPACE::is_base_of; +using CUTE_STL_NAMESPACE::is_base_of_v; + +// using CUTE_STL_NAMESPACE::true_type; +// using CUTE_STL_NAMESPACE::false_type; + +using CUTE_STL_NAMESPACE::conditional; +using CUTE_STL_NAMESPACE::conditional_t; + +using CUTE_STL_NAMESPACE::remove_cv_t; +using CUTE_STL_NAMESPACE::remove_reference_t; + +using CUTE_STL_NAMESPACE::extent; +using CUTE_STL_NAMESPACE::remove_extent; + +using CUTE_STL_NAMESPACE::decay; +using CUTE_STL_NAMESPACE::decay_t; + +using CUTE_STL_NAMESPACE::is_reference; +using CUTE_STL_NAMESPACE::is_trivially_copyable; + +using CUTE_STL_NAMESPACE::is_same; +using CUTE_STL_NAMESPACE::is_same_v; + +using CUTE_STL_NAMESPACE::is_arithmetic; +using CUTE_STL_NAMESPACE::is_unsigned; +using CUTE_STL_NAMESPACE::is_signed; +// using CUTE_STL_NAMESPACE::is_integral; + +using CUTE_STL_NAMESPACE::is_empty; + +using CUTE_STL_NAMESPACE::invoke_result_t; + +// +using CUTE_STL_NAMESPACE::declval; + +// +using CUTE_STL_NAMESPACE::numeric_limits; + +// +using CUTE_STL_NAMESPACE::ptrdiff_t; + +// +using CUTE_STL_NAMESPACE::uintptr_t; + +// C++20 +// using std::remove_cvref; +template +struct remove_cvref { + using type = remove_cv_t>; +}; + +// C++20 +// using std::remove_cvref_t; +template +using remove_cvref_t = typename remove_cvref::type; + +// +// dependent_false +// +// @brief An always-false value that depends on one or more template parameters. +// See +// https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2019/p1830r1.pdf +// https://github.com/cplusplus/papers/issues/572 +// https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2022/p2593r0.html +template +inline constexpr bool dependent_false = false; + +// +// tuple_size, tuple_element +// +// @brief CuTe-local tuple-traits to prevent conflicts with other libraries. +// For cute:: types, we specialize std::tuple-traits, which is explicitly allowed. +// cute::tuple, cute::array, cute::array_subbyte, etc +// But CuTe wants to treat some external types as tuples as well. For those, +// we specialize cute::tuple-traits to avoid polluting external traits. +// dim3, uint3, etc + +template +struct tuple_size; + +template +struct tuple_size::type>> : CUTE_STL_NAMESPACE::integral_constant::value> {}; + +// S = : std::integral_constant::value> {}; + +template +constexpr size_t tuple_size_v = tuple_size::value; + +template +struct tuple_element; + +template +struct tuple_element::type>> : CUTE_STL_NAMESPACE::tuple_element {}; + +template +using tuple_element_t = typename tuple_element::type; + +// +// is_valid +// + +namespace detail { + +template ()(declval()...))> +CUTE_HOST_DEVICE constexpr auto +is_valid_impl(int) { return CUTE_STL_NAMESPACE::true_type{}; } + +template +CUTE_HOST_DEVICE constexpr auto +is_valid_impl(...) { return CUTE_STL_NAMESPACE::false_type{}; } + +template +struct is_valid_fn { + template + CUTE_HOST_DEVICE constexpr auto + operator()(Args&&...) const { return is_valid_impl(int{}); } +}; + +} // end namespace detail + +template +CUTE_HOST_DEVICE constexpr auto +is_valid(F&&) { + return detail::is_valid_fn{}; +} + +template +CUTE_HOST_DEVICE constexpr auto +is_valid(F&&, Args&&...) { + return detail::is_valid_impl(int{}); +} + +} // end namespace cute diff --git a/csrc/mctlass/include/mctlass/aligned_buffer.h b/csrc/mctlass/include/mctlass/aligned_buffer.h new file mode 100644 index 0000000..bf76214 --- /dev/null +++ b/csrc/mctlass/include/mctlass/aligned_buffer.h @@ -0,0 +1,129 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief AlignedBuffer is a container for trivially copyable elements suitable for use in + unions and shared memory. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Modifies semantics of mctlass::Array<> to provide guaranteed alignment. +template < + typename T, + int N, + int Align = 16 +> +struct AlignedBuffer { + + /// Internal storage type + using Storage = uint8_t; + + /// Number of logical elements held in buffer + static int const kCount = N; + + /// Alignment requirement in bytes + static int const kAlign = Align; + + /// Number of storage elements + static int const kBytes = + (sizeof_bits::value * N + 7) / 8; + +private: + + /// Internal storage + alignas(Align) Storage storage[kBytes]; + +public: + + // + // C++ standard members + // + + typedef T value_type; + typedef size_t size_type; + typedef ptrdiff_t difference_type; + typedef value_type *pointer; + typedef value_type const * const_pointer; + + using Array = Array; + using reference = typename Array::reference; + using const_reference = typename Array::const_reference; + +public: + + MCTLASS_HOST_DEVICE + pointer data() { + return reinterpret_cast(storage); + } + + MCTLASS_HOST_DEVICE + const_pointer data() const { + return reinterpret_cast(storage); + } + + MCTLASS_HOST_DEVICE + Storage * raw_data() { + return storage; + } + + MCTLASS_HOST_DEVICE + Storage const * raw_data() const { + return storage; + } + + + MCTLASS_HOST_DEVICE + constexpr bool empty() const { + return !kCount; + } + + MCTLASS_HOST_DEVICE + constexpr size_type size() const { + return kCount; + } + + MCTLASS_HOST_DEVICE + constexpr size_type max_size() const { + return kCount; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/arch/arch.h b/csrc/mctlass/include/mctlass/arch/arch.h new file mode 100644 index 0000000..209b5a0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/arch.h @@ -0,0 +1,111 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines tags for architecture-specific configurations. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace arch { + +#if defined(__MXCC__) || (defined(__clang__) && defined(__MACA__)) + +/// Computes laneId within a warp +MCTLASS_DEVICE +int LaneId() { + //int ret; + //asm ("mov.u32 %0, %%laneid;" : "=r"(ret) : ); + int ret = __lane_id(); + return ret; +} + +/// Computes SM number the thread is running on +MCTLASS_DEVICE +int SmId() { + int ret; + #if 0 + asm ("mov.u32 %0, %%smid;" : "=r"(ret) : ); + #endif + printf("arch.h L60 this function cannot working correctly now.\n"); + return ret; +} + +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// +struct Sm50 { + static int const kMinComputeCapability = 50; +}; +struct Sm60 { + static int const kMinComputeCapability = 60; +}; +struct Sm61 { + static int const kMinComputeCapability = 61; +}; +struct Sm70 { + static int const kMinComputeCapability = 70; +}; +struct Sm72 { + static int const kMinComputeCapability = 72; +}; +struct Sm75 { + static int const kMinComputeCapability = 75; +}; +struct Sm80 { + static int const kMinComputeCapability = 80; +}; +struct Sm86 { + static int const kMinComputeCapability = 86; +}; +struct Sm90 { + static int const kMinComputeCapability = 90; +}; + +/// Triggers a breakpoint on the device +MCTLASS_DEVICE +void device_breakpoint() { +#if defined(__MACA_ARCH__) + //asm volatile (" brkpt;\n"); + printf("arch.h L97 this function cannot working correctly now.\n"); +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/arch/barrier.h b/csrc/mctlass/include/mctlass/arch/barrier.h new file mode 100644 index 0000000..28146c8 --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/barrier.h @@ -0,0 +1,444 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Barrier Operations on SM90+ +*/ + +#pragma once + +#include +#include +namespace mctlass { +/// @brief +namespace arch { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 && (__CUDACC_VER_MAJOR__ >= 12) +#define CUDA_BARRIER_ENABLED 1 +#else +#define CUDA_BARRIER_ENABLED 0 +#endif + +class NamedBarrier { + + // Data Members: + + // Range = [1 , NUM_THREADS_PER_CTA] + // Range % warp-size (i.e 32) == 0 + uint32_t const num_threads_; + + // Range : [0, 15] + uint32_t const id_; + + public: + + MCTLASS_DEVICE + NamedBarrier(uint32_t num_threads, uint32_t id = 0) + : num_threads_(num_threads), id_(id) {} + + MCTLASS_DEVICE + void arrive_and_wait() const { + NamedBarrier::arrive_and_wait(num_threads_, id_); + } + + MCTLASS_DEVICE + void arrive() const { + NamedBarrier::arrive(num_threads_, id_); + } + + MCTLASS_DEVICE + void sync() const { + NamedBarrier::arrive_and_wait(); + } + + // Static variants + MCTLASS_DEVICE + static void arrive_and_wait(uint32_t num_threads, uint32_t barrier_id) { +#if CUDA_BARRIER_ENABLED + asm volatile("bar.sync %0, %1;" : : "r"(barrier_id), "r"(num_threads)); +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif + } + + MCTLASS_DEVICE + static void arrive(uint32_t num_threads, uint32_t barrier_id) { +#if CUDA_BARRIER_ENABLED + asm volatile("bar.arrive %0, %1;" : : "r"(barrier_id), "r"(num_threads)); +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif + } + + MCTLASS_DEVICE + static void sync(uint32_t num_threads, uint32_t barrier_id) { + NamedBarrier::arrive_and_wait(num_threads, barrier_id); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Hopper introduces a new cluster-wide barrier which handle with Cluster-wide arrive-wait behaviour. +// This is an extension to the Ampere arrive-wait barriers +// Note : Ampere arrive-wait Barriers have a larger max-arrive count (2^30) than Hopper arrive-wait Barriers (2^20). +struct ClusterBarrier { + + using ValueType = uint64_t; + +protected: + // Can never be initialized - can only be aliased to smem + ValueType barrier_; + +public: + + MCTLASS_DEVICE + ClusterBarrier() = delete; + + MCTLASS_DEVICE + void init(uint32_t arrive_count) const { + ClusterBarrier::init(&this->barrier_, arrive_count); + } + + MCTLASS_DEVICE + uint32_t test_wait(uint32_t phase, uint32_t pred=true) const { + return ClusterBarrier::test_wait(&this->barrier_, phase, pred); + } + + MCTLASS_DEVICE + uint32_t try_wait(uint32_t phase) const { + return ClusterBarrier::try_wait(&this->barrier_, phase); + } + + MCTLASS_DEVICE + void wait(uint32_t phase) const { + ClusterBarrier::wait(&this->barrier_, phase); + } + + // Barrier arrive on local smem + MCTLASS_DEVICE + void arrive() const { + ClusterBarrier::arrive(&this->barrier_); + } + + // Remote SMEM arrive with a perdicate (usually done to pick the thread doing the arrive) + MCTLASS_DEVICE + void arrive(uint32_t cta_id, uint32_t pred = true ) const { + ClusterBarrier::arrive(&this->barrier_, cta_id, pred); + } + + // + // Static Versions + // + MCTLASS_DEVICE + static void init(ValueType const* smem_ptr, uint32_t arrive_count) { +#if CUDA_BARRIER_ENABLED + uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); + asm volatile( + "{\n\t" + "mbarrier.init.shared.b64 [%1], %0; \n" + "}" + : + : "r"(arrive_count), "r"(smem_addr)); +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif + } + + // Static version of wait - in case we don't want to burn a register + MCTLASS_DEVICE + static void wait(ValueType const* smem_ptr, uint32_t phase) { +#if CUDA_BARRIER_ENABLED + uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); + // Arbitrarily large timer value after which try-wait expires and re-tries. + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1; \n\t" + "LAB_WAIT: \n\t" + "mbarrier.try_wait.parity.shared.b64 P1, [%0], %1, %2; \n\t" + "@P1 bra.uni DONE; \n\t" + "bra.uni LAB_WAIT; \n\t" + "DONE: \n\t" + "}" + : + : "r"(smem_addr), "r"(phase), "r"(ticks)); + +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif + } + + MCTLASS_DEVICE + static uint32_t test_wait(ValueType const* smem_ptr, uint32_t phase, uint32_t pred) { +#if CUDA_BARRIER_ENABLED + uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); + uint32_t waitComplete; + + asm volatile( + "{\n\t" + ".reg .pred P1; \n\t" + ".reg .pred P2; \n\t" + "setp.eq.u32 P2, %3, 1;\n\t" + "@P2 mbarrier.test_wait.parity.shared.b64 P1, [%1], %2; \n\t" + "selp.b32 %0, 1, 0, P1; \n\t" + "}" + : "=r"(waitComplete) + : "r"(smem_addr), "r"(phase), "r"(pred)); + + return waitComplete; +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif + return 0; + } + + MCTLASS_DEVICE + static uint32_t try_wait(ValueType const* smem_ptr, uint32_t phase) { +#if CUDA_BARRIER_ENABLED + uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); + uint32_t waitComplete; + + asm volatile( + "{\n\t" + ".reg .pred P1; \n\t" + "mbarrier.try_wait.parity.shared.b64 P1, [%1], %2; \n\t" + "selp.b32 %0, 1, 0, P1; \n\t" + "}" + : "=r"(waitComplete) + : "r"(smem_addr), "r"(phase)); + + return waitComplete; +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif + return 0; + } + + // Static Predicated version of the above - in case we know the address. + MCTLASS_DEVICE + static void arrive(ValueType const* smem_ptr, uint32_t cta_id, uint32_t pred) { +#if CUDA_BARRIER_ENABLED + uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 remAddr32;\n\t" + "setp.eq.u32 p, %2, 1;\n\t" + "@p mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "@p mbarrier.arrive.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + : + : "r"(smem_addr), "r"(cta_id), "r"(pred)); +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif + } + + // Barrier arrive on local smem + MCTLASS_DEVICE + static void arrive(ValueType const* smem_ptr) { +#if CUDA_BARRIER_ENABLED + uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); + uint64_t state = 0; + asm volatile( + "{\n\t" + "mbarrier.arrive.shared.b64 %1, [%0];\n\t" + "}" + : + : "r"(smem_addr), "l"(state)); +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif + } + + MCTLASS_DEVICE + static void invalidate(ValueType const* smem_ptr) { +#if CUDA_BARRIER_ENABLED + uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); + asm volatile( + "{\n\t" + "mbarrier.ival.shared.b64 [%0]; \n\t" + "}" + : + : "r"(smem_addr)); +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// SM90 also introduces a new type of cluster-barrier which supports sync. +// not just based on Arrive Count, but also transaction count (in bytes) +struct ClusterTransactionBarrier : public ClusterBarrier { + + MCTLASS_DEVICE + ClusterTransactionBarrier() = delete; + + // Performs an arrive operation + bytes reset + MCTLASS_DEVICE + void arrive_and_reset_bytes(uint32_t transaction_bytes) const { + ClusterTransactionBarrier::arrive_and_reset_bytes(&this->barrier_, transaction_bytes); + } + + // Performs an arrive operation + bytes reset + MCTLASS_DEVICE + void arrive_and_reset_bytes(uint32_t transaction_bytes, uint32_t cta_id) const { + ClusterTransactionBarrier::arrive_and_reset_bytes(&this->barrier_, transaction_bytes , cta_id, true); + } + + MCTLASS_DEVICE + void commit(uint32_t transaction_bytes, uint32_t pred = 1) const { + uint32_t cta_rank = cute::block_rank_in_cluster(); + ClusterTransactionBarrier::commit(&this->barrier_, cta_rank, transaction_bytes, pred); + } + + MCTLASS_DEVICE + void commit(uint32_t dst_cta_id, uint32_t transaction_bytes, uint32_t pred) const { + ClusterTransactionBarrier::commit(&this->barrier_, dst_cta_id, transaction_bytes, pred); + } + + // + // Static Versions + // + + // Performs an arrive operation + bytes reset + MCTLASS_DEVICE + static void arrive_and_reset_bytes(ValueType const* smem_ptr, uint32_t transaction_bytes) { +#if CUDA_BARRIER_ENABLED + uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); + asm volatile( + "{\n\t" + "mbarrier.arrive.expect_tx.shared.b64 _, [%1], %0; \n\t" + "}" + : + : "r"(transaction_bytes), "r"(smem_addr)); +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif + } + + // Performs an arrive operation + bytes reset for a remote cta_id in a Cluster + MCTLASS_DEVICE + static void arrive_and_reset_bytes( + ValueType const* smem_ptr, uint32_t transaction_bytes, uint32_t cta_id, uint32_t pred) { +#if CUDA_BARRIER_ENABLED + uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b32 remAddr32;\n\t" + "setp.eq.u32 p, %2, 1;\n\t" + "@p mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "@p mbarrier.arrive.expect_tx.shared::cluster.b64 _, [remAddr32], %3;\n\t" + "}" + : + : "r"(smem_addr), "r"(cta_id), "r"(pred), "r"(transaction_bytes)); +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif + } + + // Performs an bytes reset without doing an arrive operation + MCTLASS_DEVICE + static void reset_bytes(ValueType const* smem_ptr, uint32_t transaction_bytes) { +#if CUDA_BARRIER_ENABLED + uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); + asm volatile( + "{\n\t" + "mbarrier.expect_tx.shared.b64 [%1], %0; \n\t" + "}" + : + : "r"(transaction_bytes), "r"(smem_addr)); +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif + } + + // Increments transaction bytes in the barrier + MCTLASS_DEVICE + static void commit( + ValueType const* smem_ptr, uint32_t dst_cta_id, uint32_t transaction_bytes, uint32_t pred = 1) { +#if CUDA_BARRIER_ENABLED + uint32_t smem_addr = cute::cast_smem_ptr_to_uint(smem_ptr); + smem_addr = cute::set_block_rank(smem_addr, dst_cta_id); + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.eq.u32 p, %2, 1;\n\t" + "@p mbarrier.complete_tx.shared::cluster.relaxed.cluster.b64 [%1], %0;" + "}" + : + : "r"(transaction_bytes), "r"(smem_addr), "r"(pred)); +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif + } +}; + +// Helps with visibility of barrier init operations across warps / cta / cluster +// Available as a separate function so as to batch inits across barriers and fence once +// Note : It must be composed with an appropriate sync instruction with the right scope +// to ensure visibility eg. __syncthreads() or a cluster_arrive() + cluster_wait() +MCTLASS_DEVICE +void fence_barrier_init() { +#if CUDA_BARRIER_ENABLED + asm volatile( + "{\n\t" + "fence.mbarrier_init.release.cluster; \n" + "}" + ::); +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif +} + +// Issue a shared memory fence for async operations +MCTLASS_DEVICE +void fence_view_async_shared() { +#if CUDA_BARRIER_ENABLED + asm volatile ( + "{\n\t" + "fence.proxy.async.shared::cta; \n" + "}" + ::); +#elif defined(__CUDA_ARCH__) + asm volatile ("brkpt;\n" ::); +#endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////////////////////////// +} // end namespace arch +} // end namespace mctlass diff --git a/csrc/mctlass/include/mctlass/arch/cache_operation.h b/csrc/mctlass/include/mctlass/arch/cache_operation.h new file mode 100644 index 0000000..a885df9 --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/cache_operation.h @@ -0,0 +1,66 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Directives related to cache operations +*/ +#pragma once + +#include "mctlass/mctlass.h" + +namespace mctlass { +namespace arch { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Controls PTX cache operations +struct CacheOperation { + enum Kind { + /// Cache at all levels - accessed again + Always, + /// Cache at global level + Global, + /// Streaming - likely to be accessed once + Streaming, + /// Indicates the line will not be used again + LastUse, + /// Don't cache, and fetch again + Volatile, + /// Write back at all coherent levels + WriteBack, + /// Write through to system memory + WriteThrough + }; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/arch/memory.h b/csrc/mctlass/include/mctlass/arch/memory.h new file mode 100644 index 0000000..6dd977b --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/memory.h @@ -0,0 +1,684 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Architecture-specific operators on memory +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/arch/cache_operation.h" + +namespace mctlass { +namespace arch { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Fragment type to store loaded data + typename AccessType, + /// The bytes of loading + int LoadBytes, + /// Cache operation + CacheOperation::Kind cache_op = CacheOperation::Always + > +struct global_load; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Specializations +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#if (((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 4)) || \ + (__CUDACC_VER_MAJOR__ > 11)) && \ + defined(__MACA_ARCH__) && \ + ! (defined(__clang__) && defined(__MACA__)) + #define MCTLASS_ENABLE_L2_PREFETCH 1 +#else + #define MCTLASS_ENABLE_L2_PREFETCH 0 +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// The redundant mov PTX instruction is used to enforce the compiler to +// keep the initializing code before ld.global +template +struct global_load { + MCTLASS_DEVICE + global_load(AccessType &D, void const *ptr, bool pred_guard) { +#if 0 + uint4 *data = reinterpret_cast(&D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %9, 0;\n" + " mov.b32 %0, %10;\n" + " mov.b32 %1, %11;\n" + " mov.b32 %2, %12;\n" + " mov.b32 %3, %13;\n" + " mov.b32 %4, %14;\n" + " mov.b32 %5, %15;\n" + " mov.b32 %6, %16;\n" + " mov.b32 %7, %17;\n" +#if MCTLASS_ENABLE_L2_PREFETCH + " @p ld.global.L2::128B.v4.u32 {%0, %1, %2, %3}, [%8];\n" + " @p ld.global.L2::128B.v4.u32 {%4, %5, %6, %7}, [%18];\n" +#else + " @p ld.global.v4.u32 {%0, %1, %2, %3}, [%8];\n" + " @p ld.global.v4.u32 {%4, %5, %6, %7}, [%18];\n" +#endif + "}\n" + : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w), + "=r"(data[1].x), "=r"(data[1].y), "=r"(data[1].z), "=r"(data[1].w) + : "l"(ptr), "r"((int)pred_guard), "r"(data[0].x), "r"(data[0].y), + "r"(data[0].z), "r"(data[0].w), "r"(data[1].x), "r"(data[1].y), + "r"(data[1].z), "r"(data[1].w), "l"(((uint8_t *)ptr) + 16)); +#else + if (pred_guard) D = *(reinterpret_cast(ptr)); +#endif + } +}; + +template +struct global_load { + MCTLASS_DEVICE + global_load(AccessType &D, void const *ptr, bool pred_guard) { +#if 0 + uint4 *data = reinterpret_cast(&D); + + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %9, 0;\n" + " mov.b32 %0, %10;\n" + " mov.b32 %1, %11;\n" + " mov.b32 %2, %12;\n" + " mov.b32 %3, %13;\n" + " mov.b32 %4, %14;\n" + " mov.b32 %5, %15;\n" + " mov.b32 %6, %16;\n" + " mov.b32 %7, %17;\n" + " @p ld.global.lu.v4.u32 {%0, %1, %2, %3}, [%8];\n" + " @p ld.global.lu.v4.u32 {%4, %5, %6, %7}, [%18];\n" + "}\n" + : "=r"(data[0].x), "=r"(data[0].y), "=r"(data[0].z), "=r"(data[0].w), + "=r"(data[1].x), "=r"(data[1].y), "=r"(data[1].z), "=r"(data[1].w) + : "l"(ptr), "r"((int)pred_guard), "r"(data[0].x), "r"(data[0].y), + "r"(data[0].z), "r"(data[0].w), "r"(data[1].x), "r"(data[1].y), + "r"(data[1].z), "r"(data[1].w), "l"(((uint8_t *)ptr) + 16)); +#else + if (pred_guard) D = *(reinterpret_cast(ptr)); +#endif + } +}; + +template +struct global_load { + MCTLASS_DEVICE + global_load(AccessType &D, void const *ptr, bool pred_guard) { +#if 0 + uint4 &data = reinterpret_cast(D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %5, 0;\n" + " mov.b32 %0, %6;\n" + " mov.b32 %1, %7;\n" + " mov.b32 %2, %8;\n" + " mov.b32 %3, %9;\n" +#if MCTLASS_ENABLE_L2_PREFETCH + " @p ld.global.L2::128B.v4.u32 {%0, %1, %2, %3}, [%4];\n" +#else + " @p ld.global.v4.u32 {%0, %1, %2, %3}, [%4];\n" +#endif + "}\n" + : "=r"(data.x), "=r"(data.y), "=r"(data.z), "=r"(data.w) + : "l"(ptr), "r"((int)pred_guard), "r"(data.x), "r"(data.y), "r"(data.z), "r"(data.w)); +#else + if (pred_guard) D = *(reinterpret_cast(ptr)); +#endif + } +}; + +template +struct global_load { + MCTLASS_DEVICE + global_load(AccessType &D, void const *ptr, bool pred_guard) { +#if 0 + uint4 &data = reinterpret_cast(D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %5, 0;\n" + " mov.b32 %0, %6;\n" + " mov.b32 %1, %7;\n" + " mov.b32 %2, %8;\n" + " mov.b32 %3, %9;\n" + " @p ld.global.lu.v4.u32 {%0, %1, %2, %3}, [%4];\n" + "}\n" + : "=r"(data.x), "=r"(data.y), "=r"(data.z), "=r"(data.w) + : "l"(ptr), "r"((int)pred_guard), "r"(data.x), "r"(data.y), "r"(data.z), "r"(data.w)); +#else + if (pred_guard) D = *(reinterpret_cast(ptr)); +#endif + } +}; + +template +struct global_load { + MCTLASS_DEVICE + global_load(AccessType &D, void const *ptr, bool pred_guard) { +#if 0 + uint2 &data = reinterpret_cast(D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %3, 0;\n" + " mov.b32 %0, %4;\n" + " mov.b32 %1, %5;\n" +#if MCTLASS_ENABLE_L2_PREFETCH + " @p ld.global.L2::128B.v2.u32 {%0, %1}, [%2];\n" +#else + " @p ld.global.v2.u32 {%0, %1}, [%2];\n" +#endif + "}\n" + : "=r"(data.x), "=r"(data.y) + : "l"(ptr), "r"((int)pred_guard), "r"(data.x), "r"(data.y)); +#else + if (pred_guard) D = *(reinterpret_cast(ptr)); +#endif + } +}; + +template +struct global_load { + MCTLASS_DEVICE + global_load(AccessType &D, void const *ptr, bool pred_guard) { +#if 0 + uint2 &data = reinterpret_cast(D); + + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %3, 0;\n" + " mov.b32 %0, %4;\n" + " mov.b32 %1, %5;\n" + " @p ld.global.lu.v2.u32 {%0, %1}, [%2];\n" + "}\n" + : "=r"(data.x), "=r"(data.y) + : "l"(ptr), "r"((int)pred_guard), "r"(data.x), "r"(data.y)); +#else + if (pred_guard) D = *(reinterpret_cast(ptr)); +#endif + } +}; + +template +struct global_load { + MCTLASS_DEVICE + global_load(AccessType &D, void const *ptr, bool pred_guard) { +#if 0 + unsigned &data = reinterpret_cast(D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %2, 0;\n" + " mov.b32 %0, %3;\n" +#if MCTLASS_ENABLE_L2_PREFETCH + " @p ld.global.L2::128B.u32 %0, [%1];\n" +#else + " @p ld.global.u32 %0, [%1];\n" +#endif + "}\n" + : "=r"(data) + : "l"(ptr), "r"((int)pred_guard), "r"(data)); +#else + if (pred_guard) D = *(reinterpret_cast(ptr)); +#endif + } +}; + +template +struct global_load { + MCTLASS_DEVICE + global_load(AccessType &D, void const *ptr, bool pred_guard) { +#if 0 + unsigned &data = reinterpret_cast(D); + + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %2, 0;\n" + " mov.b32 %0, %3;\n" + " @p ld.global.lu.u32 %0, [%1];\n" + "}\n" + : "=r"(data) + : "l"(ptr), "r"((int)pred_guard), "r"(data)); +#else + if (pred_guard) D = *(reinterpret_cast(ptr)); +#endif + } +}; + +template +struct global_load { + MCTLASS_DEVICE + global_load(AccessType &D, void const *ptr, bool pred_guard) { +#if 0 + uint16_t &data = reinterpret_cast(D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %2, 0;\n" + " mov.b16 %0, %3;\n" +#if MCTLASS_ENABLE_L2_PREFETCH + " @p ld.global.L2::128B.u16 %0, [%1];\n" +#else + " @p ld.global.u16 %0, [%1];\n" +#endif + "}\n" + : "=h"(data) + : "l"(ptr), "r"((int)pred_guard), "h"(data)); +#else + if (pred_guard) D = *(reinterpret_cast(ptr)); +#endif + } +}; + +template +struct global_load { + MCTLASS_DEVICE + global_load(AccessType &D, void const *ptr, bool pred_guard) { +#if 0 + uint16_t &data = reinterpret_cast(D); + + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %2, 0;\n" + " mov.b16 %0, %3;\n" + " @p ld.global.lu.u16 %0, [%1];\n" + "}\n" + : "=h"(data) + : "l"(ptr), "r"((int)pred_guard), "h"(data)); +#else + if (pred_guard) D = *(reinterpret_cast(ptr)); +#endif + } +}; + +template +struct global_load { + MCTLASS_DEVICE + global_load(AccessType &D, void const *ptr, bool pred_guard) { + if (pred_guard) D = *(reinterpret_cast(ptr)); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Fragment type to store data + typename AccessType, + /// The bytes of storing + int StoreBytes + > +struct global_store; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Specializations +// +///////////////////////////////////////////////////////////////////////////////////////////////// + + +template +struct global_store { + MCTLASS_DEVICE + global_store(AccessType const &D, void *ptr, bool pred_guard) { +#if 0 + uint4 const *data = reinterpret_cast(&D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %5, 0;\n" + " @p st.global.v4.u32 [%0], {%1, %2, %3, %4};\n" + " @p st.global.v4.u32 [%6], {%7, %8, %9, %10};\n" + " @p st.global.v4.u32 [%11], {%12, %13, %14, %15};\n" + " @p st.global.v4.u32 [%16], {%17, %18, %19, %20};\n" + "}\n" + : + : "l"(ptr), "r"(data[0].x), "r"(data[0].y), "r"(data[0].z), + "r"(data[0].w), "r"((int)pred_guard), "l"(((uint8_t *)ptr) + 16), + "r"(data[1].x), "r"(data[1].y), "r"(data[1].z), "r"(data[1].w), + "l"(((uint8_t *)ptr) + 32), + "r"(data[2].x), "r"(data[2].y), "r"(data[2].z), "r"(data[2].w), + "l"(((uint8_t *)ptr) + 48), + "r"(data[3].x), "r"(data[3].y), "r"(data[3].z), "r"(data[3].w)); +#else + if (pred_guard) *(reinterpret_cast(ptr)) = D; +#endif + } +}; + + +template +struct global_store { + MCTLASS_DEVICE + global_store(AccessType const &D, void *ptr, bool pred_guard) { +#if 0 + uint4 const *data = reinterpret_cast(&D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %5, 0;\n" + " @p st.global.v4.u32 [%0], {%1, %2, %3, %4};\n" + " @p st.global.v4.u32 [%6], {%7, %8, %9, %10};\n" + "}\n" + : + : "l"(ptr), "r"(data[0].x), "r"(data[0].y), "r"(data[0].z), + "r"(data[0].w), "r"((int)pred_guard), "l"(((uint8_t *)ptr) + 16), + "r"(data[1].x), "r"(data[1].y), "r"(data[1].z), "r"(data[1].w)); +#else + if (pred_guard) *(reinterpret_cast(ptr)) = D; +#endif + } +}; + +template +struct global_store { + MCTLASS_DEVICE + global_store(AccessType const &D, void *ptr, bool pred_guard) { +#if 0 + uint4 const &data = reinterpret_cast(D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %5, 0;\n" + " @p st.global.v4.u32 [%0], {%1, %2, %3, %4};\n" + "}\n" + : + : "l"(ptr), "r"(data.x), "r"(data.y), "r"(data.z), "r"(data.w), "r"((int)pred_guard)); +#else + if (pred_guard) *(reinterpret_cast(ptr)) = D; +#endif + } +}; + +template +struct global_store { + MCTLASS_DEVICE + global_store(AccessType const &D, void *ptr, bool pred_guard) { +#if 0 + uint2 const &data = reinterpret_cast(D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %3, 0;\n" + " @p st.global.v2.u32 [%0], {%1, %2};\n" + "}\n" + : + : "l"(ptr), "r"(data.x), "r"(data.y), "r"((int)pred_guard)); +#else + if (pred_guard) *(reinterpret_cast(ptr)) = D; +#endif + } +}; + +template +struct global_store { + MCTLASS_DEVICE + global_store(AccessType const &D, void *ptr, bool pred_guard) { +#if 0 + uint32_t const &data = reinterpret_cast(D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %2, 0;\n" + " @p st.global.u32 [%0], %1;\n" + "}\n" + : + : "l"(ptr), "r"(data), "r"((int)pred_guard)); +#else + if (pred_guard) *(reinterpret_cast(ptr)) = D; +#endif + } +}; + +template +struct global_store { + MCTLASS_DEVICE + global_store(AccessType const &D, void *ptr, bool pred_guard) { +#if 0 + uint16_t const &data = reinterpret_cast(D); + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %2, 0;\n" + " @p st.global.u16 [%0], %1;\n" + "}\n" + : + : "l"(ptr), "h"(data), "r"((int)pred_guard)); +#else + if (pred_guard) *(reinterpret_cast(ptr)) = D; +#endif + } +}; + +template +struct global_store { + MCTLASS_DEVICE + global_store(AccessType const &D, void *ptr, bool pred_guard) { + if (pred_guard) *(reinterpret_cast(ptr)) = D; + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// ld.shared +template +MCTLASS_DEVICE +void shared_load(void *dst, uint32_t ptr); + +/// ld.shared - 16b +template <> +MCTLASS_DEVICE +void shared_load<2>(void *dst, uint32_t ptr) { + #if 0 + asm volatile("ld.shared.u16 %0, [%1];\n" + : "=h"(*reinterpret_cast(dst)) + : "r"(ptr)); + #endif + printf("memory.h L409 this function cannot working correctly now.\n"); +} + +/// ld.shared - 32b +template <> +MCTLASS_DEVICE +void shared_load<4>(void *dst, uint32_t ptr) { + #if 0 + asm volatile("ld.shared.u32 %0, [%1];\n" + : "=r"(*reinterpret_cast(dst)) + : "r"(ptr)); + #endif + printf("memory.h L421 this function cannot working correctly now.\n"); +} + +/// ld.shared - 64b +template <> +MCTLASS_DEVICE +void shared_load<8>(void *dst, uint32_t ptr) { + uint2 *dst_u64 = reinterpret_cast(dst); + #if 0 + asm volatile("ld.shared.v2.u32 {%0, %1}, [%2];\n" + : + "=r"(dst_u64->x), + "=r"(dst_u64->y) + : "r"(ptr)); + #endif + printf("memory.h L436 this function cannot working correctly now.\n"); +} + +/// ld.shared - 128b +template <> +MCTLASS_DEVICE +void shared_load<16>(void *dst, uint32_t ptr) { + uint4 *dst_u128 = reinterpret_cast(dst); + #if 0 + asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : + "=r"(dst_u128->x), + "=r"(dst_u128->y), + "=r"(dst_u128->z), + "=r"(dst_u128->w) + : "r"(ptr)); + #endif + printf("memory.h L453 this function cannot working correctly now.\n"); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// st.shared +template +MCTLASS_DEVICE +void shared_store(uint32_t ptr, void const *src); + +/// st.shared - 16b +template <> +MCTLASS_DEVICE +void shared_store<2>(uint32_t ptr, void const *src) { + #if 0 + asm volatile("st.shared.u16 [%0], %1;\n" + : : + "r"(ptr), + "h"(*reinterpret_cast(src)) + ); + #endif + printf("memory.h L475 this function cannot working correctly now.\n"); +} + +/// st.shared - 32b +template <> +MCTLASS_DEVICE +void shared_store<4>(uint32_t ptr, void const *src) { + #if 0 + asm volatile("st.shared.u32 [%0], %1;\n" + : : + "r"(ptr), + "r"(*reinterpret_cast(src)) + ); + #endif + printf("memory.h L489 this function cannot working correctly now.\n"); +} + +/// st.shared - 64b +template <> +MCTLASS_DEVICE +void shared_store<8>(uint32_t ptr, void const *src) { + uint2 const *dst_u64 = reinterpret_cast(src); + #if 0 + asm volatile("st.shared.v2.u32 [%0], {%1, %2};\n" + : : + "r"(ptr), + "r"(dst_u64->x), + "r"(dst_u64->y) + ); + #endif + printf("memory.h L503 this function cannot working correctly now.\n"); +} + +/// st.shared - 128b +template <> +MCTLASS_DEVICE +void shared_store<16>(uint32_t ptr, void const *src) { + uint4 const *dst_u128 = reinterpret_cast(src); + #if 0 + asm volatile("ld.shared.v4.u32 [%0], {%1, %2, %3, %4};\n" + : : + "r"(ptr), + "r"(dst_u128->x), + "r"(dst_u128->y), + "r"(dst_u128->z), + "r"(dst_u128->w) + ); + #endif + printf("memory.h L522 this function cannot working correctly now.\n"); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/arch/memory_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/arch/memory_sm75.h b/csrc/mctlass/include/mctlass/arch/memory_sm75.h new file mode 100644 index 0000000..9b6027b --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/memory_sm75.h @@ -0,0 +1,354 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "cute/arch/util.hpp" + +namespace mctlass { +namespace arch { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Layout of destination matrix (column-major implies transpose) + typename Layout, + /// .x1, .x2, or .x4 + int MatrixCount +> +inline __device__ void ldsm(Array & D, void const* ptr); + +template < + /// Layout of destination matrix (column-major implies transpose) + typename Layout, + /// .x1, .x2, or .x4 + int MatrixCount +> +inline __device__ void ldsmAtf32(Array & D, void const* ptr); + +template < + /// Layout of destination matrix (column-major implies transpose) + typename Layout, + /// .x1, .x2, or .x4 + int MatrixCount +> +inline __device__ void ldsmBtf32(Array & D, void const* ptr, int ldm); +template < + /// Layout of destination matrix (column-major implies transpose) + typename Layout, + /// .x1, .x2, or .x4 + int MatrixCount +> +inline __device__ void ldsmi8(Array & D, void const* ptr, int const* ldm); + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Determine the appropriate way to target PTX's "ldmatrix" instruction. +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// MCTLASS helper to get SMEM pointer +//Original MACA impl +inline __device__ unsigned mctlass_get_smem_pointer_maca(void *ptr) { + MCTLASS_UNUSED(ptr); + MCTLASS_NOT_IMPLEMENTED(); + return 0; +} + +inline __device__ unsigned mctlass_get_smem_pointer(void *ptr) { + return static_cast(__cvta_generic_to_shared(ptr)); +} + +/// MCTLASS helper to get SMEM pointer +inline __device__ unsigned mctlass_get_smem_pointer(void const *ptr) { + return mctlass_get_smem_pointer(const_cast(ptr)); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +inline __device__ void ldsm( + Array & D, + void const* ptr) { + + #if defined(CUTE_ARCH_LDSM_SM75_ACTIVATED) + unsigned addr = mctlass_get_smem_pointer(ptr); + + int x; + asm volatile ("ldmatrix.sync.aligned.x1.m8n8.shared.b16 {%0}, [%1];" : "=r"(x) : "r"(addr)); + reinterpret_cast(D) = x; + #else + printf("memory_sm75.h L183 this function cannot working correctly now.\n"); + MCTLASS_UNUSED(D); + MCTLASS_UNUSED(ptr); + MCTLASS_NOT_IMPLEMENTED(); + + #endif +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +inline __device__ void ldsm( + Array & D, + void const* ptr) { + + #if defined(CUTE_ARCH_LDSM_SM75_ACTIVATED) + unsigned addr = mctlass_get_smem_pointer(ptr); + + int x, y; + asm volatile ("ldmatrix.sync.aligned.x2.m8n8.shared.b16 {%0, %1}, [%2];" : "=r"(x), "=r"(y) : "r"(addr)); + reinterpret_cast(D) = make_int2(x, y); + #else + printf("memory_sm75.h L205 this function cannot working correctly now.\n"); + MCTLASS_UNUSED(D); + MCTLASS_UNUSED(ptr); + MCTLASS_NOT_IMPLEMENTED(); + + #endif +} + +template <> +inline __device__ void ldsmBtf32( + Array & D, + void const* ptr, int ldm) { + int const *p = reinterpret_cast(ptr); + int const *p1 = (p + ldm / 4); + int x, y; + x = p[0]; + y = p1[0]; + + reinterpret_cast(D) = make_int2(x, y); +} + +template <> +inline __device__ void ldsmi8( + Array & D, + void const* ptr, int const* ldm) { + int const *p = reinterpret_cast(ptr); + int const *p0 = p + ldm[0]; + int x, y; + x = p[0]; + y = p0[0]; + + reinterpret_cast(D) = make_int2(x, y); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +inline __device__ void ldsm( + Array & D, + void const* ptr) { + + #if defined(CUTE_ARCH_LDSM_SM75_ACTIVATED) + unsigned addr = mctlass_get_smem_pointer(ptr); + + int x, y, z, w; + asm volatile ("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];" : "=r"(x), "=r"(y), "=r"(z), "=r"(w) : "r"(addr)); + reinterpret_cast(D) = make_int4(x, y, z, w); + #else + printf("memory_sm75.h L253 this function cannot working correctly now.\n"); + MCTLASS_UNUSED(D); + MCTLASS_UNUSED(ptr); + MCTLASS_NOT_IMPLEMENTED(); + + #endif +} + +template <> +inline __device__ void ldsmAtf32( + Array & D, + void const* ptr) { + int const *p = reinterpret_cast(ptr); + int const *p0 = p - 1; + int x, y, z, w; + x = p[0]; + y = p0[0]; + reinterpret_cast(D) = make_int4(x, y, z, w); +} + +template <> +inline __device__ void ldsmBtf32( + Array & D, + void const* ptr, int ldm) { + + int const *p = reinterpret_cast(ptr); + + int const *p0 = p - 1; + int const *p1 = p0 + 8 * ldm; + int x, y, z, w; + x = p[0]; + y = p0[0]; + z = p1[1]; + w = p1[0]; + + reinterpret_cast(D) = make_int4(x, y, z, w); +} + +template <> +inline __device__ void ldsmi8( + Array & D, + void const* ptr, int const* ldm) { + int const *p = reinterpret_cast(ptr); + int const *p0 = p + ldm[0]; + int const *p1 = p + ldm[1]; + int const *p2 = p + ldm[2]; + int x, y, z, w; + x = p[0]; + y = p0[0]; + z = p1[0]; + w = p2[0]; + + reinterpret_cast(D) = make_int4(x, y, z, w); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Transpose on 16b granularity +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +inline __device__ void ldsm( + Array & D, + void const* ptr) { + + #if defined(CUTE_ARCH_LDSM_SM75_ACTIVATED) + unsigned addr = mctlass_get_smem_pointer(ptr); + + int x; + asm volatile ("ldmatrix.sync.aligned.x1.trans.m8n8.shared.b16 {%0}, [%1];" : "=r"(x) : "r"(addr)); + reinterpret_cast(D) = x; + #else + printf("memory_sm75.h L326 this function cannot working correctly now.\n"); + MCTLASS_UNUSED(D); + MCTLASS_UNUSED(ptr); + MCTLASS_NOT_IMPLEMENTED(); + + #endif +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +inline __device__ void ldsm( + Array & D, + void const* ptr) { + + #if defined(CUTE_ARCH_LDSM_SM75_ACTIVATED) + unsigned addr = mctlass_get_smem_pointer(ptr); + + int x, y; + asm volatile ("ldmatrix.sync.aligned.x2.trans.m8n8.shared.b16 {%0, %1}, [%2];" : "=r"(x), "=r"(y) : "r"(addr)); + reinterpret_cast(D) = make_int2(x, y); + #else + printf("memory_sm75.h L348 this function cannot working correctly now.\n"); + MCTLASS_UNUSED(D); + MCTLASS_UNUSED(ptr); + MCTLASS_NOT_IMPLEMENTED(); + + #endif +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +inline __device__ void ldsm( + Array & D, + void const* ptr) { + + #if defined(CUTE_ARCH_LDSM_SM75_ACTIVATED) + unsigned addr = mctlass_get_smem_pointer(ptr); + + int x, y, z, w; + asm volatile ("ldmatrix.sync.aligned.x4.trans.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];" : "=r"(x), "=r"(y), "=r"(z), "=r"(w) : "r"(addr)); + reinterpret_cast(D) = make_int4(x, y, z, w); + #else + printf("memory_sm75.h L370 this function cannot working correctly now.\n"); + MCTLASS_UNUSED(D); + MCTLASS_UNUSED(ptr); + MCTLASS_NOT_IMPLEMENTED(); + + #endif +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct shared_load_op { + MCTLASS_DEVICE + shared_load_op(AccessType &D, void const *ptr) { + D = *reinterpret_cast(ptr); + } +}; + +template +MCTLASS_DEVICE void shared_load(AccessType &D, void const *ptr) { + shared_load_op(D, ptr); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct shared_load_op { + MCTLASS_DEVICE + shared_load_op(AccessType &D, void const *ptr) { + unsigned addr = mctlass_get_smem_pointer(ptr); + uint4 v; + // asm volatile ("ld.shared.v4.b32 {%0, %1, %2, %3}, [%4];" : + // "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) : "r"(addr)); + printf("memory_sm75.h L403 this function cannot working correctly now.\n"); + D = reinterpret_cast(v); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct shared_load_op { + MCTLASS_DEVICE + shared_load_op(AccessType &D, void const *ptr) { + unsigned addr = mctlass_get_smem_pointer(ptr); + uint2 v; + // asm volatile ("ld.shared.v2.b32 {%0, %1}, [%2];" : + // "=r"(v.x), "=r"(v.y) : "r"(addr)); + printf("memory_sm75.h L418 this function cannot working correctly now.\n"); + D = reinterpret_cast(v); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/arch/memory_sm80.h b/csrc/mctlass/include/mctlass/arch/memory_sm80.h new file mode 100644 index 0000000..b4ab2b8 --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/memory_sm80.h @@ -0,0 +1,459 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Architecture-specific operators on memory added for SM80 +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/complex.h" +#include "mctlass/arch/memory.h" +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/arch/cache_operation.h" + +#if defined(__MACA_ARCH__) //mctlass not support cp_async now + #define CUDA_CP_ASYNC_ACTIVATED 0 +#endif + +namespace mctlass { +namespace arch { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Initiates an asynchronous copy from global memory to shared memory. +/// +/// cp.async +/// +template < + /// Size of the access in bytes + int SizeInBytes, + /// Cache operation + CacheOperation::Kind cache_op = CacheOperation::Always> +struct cp_async; + +/// Initiates an asynchronous copy from global memory to shared memory. Rather than predicate +/// the entire transfer, zeros are written to SMEM if the guard predicate is false. +/// +/// cp.async +/// +template < + /// Size of the access in bytes + int SizeInBytes, + /// Cache operation + CacheOperation::Kind cache_op = CacheOperation::Always> +struct cp_async_zfill; + +/// Initiates an asynchronous copy from global memory to shared memory. Rather than predicate +/// the entire transfer, nans (0x7eff) are written to SMEM if the guard predicate is false. +/// +/// cp.async +/// +template < + /// Size of the access in bytes + int SizeInBytes, + /// Cache operation + CacheOperation::Kind cache_op = CacheOperation::Always> +struct cp_async_nan; + +/// Either 0 or 1 are written to SMEM based on input element type +/// Used for diagonal elements of triangular matrix of BLAS3 functions +/// +/// st.shared +/// +template < + /// Type of Element + typename Element, + /// If the data is for a Hermitian matrix diagonal + bool IsHermitianData = false> +struct cp_async_diag; + +static const uint32_t OOB_NAN_F16 = 0x7eff; +static const uint32_t OOB_NAN_F16x2 = ((OOB_NAN_F16 << 16) | OOB_NAN_F16); + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization +template < + /// Size of the access in bytes + int SizeInBytes> +struct cp_async { + + /// Copy + MCTLASS_DEVICE + cp_async(void *smem_ptr, void const *global_ptr, bool pred_guard = true) { + #if CUDA_CP_ASYNC_ACTIVATED + + // Make sure the size is supported. + static_assert((SizeInBytes == 4 || SizeInBytes == 8 || SizeInBytes == 16), + "Size is not supported"); + + unsigned smem_int_ptr = mctlass_get_smem_pointer(smem_ptr); + + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" +#if MCTLASS_ENABLE_L2_PREFETCH + " @p cp.async.ca.shared.global.L2::128B [%1], [%2], %3;\n" +#else + " @p cp.async.ca.shared.global [%1], [%2], %3;\n" +#endif + "}\n" ::"r"((int)pred_guard), + "r"(smem_int_ptr), "l"(global_ptr), "n"(SizeInBytes)); + + #else + using AccessType = Array; + + if (pred_guard) { + *static_cast(smem_ptr) = *static_cast(global_ptr); + } + #endif + } +}; + +/// Partial specialization +template < + /// Size of the access in bytes + int SizeInBytes> +struct cp_async_zfill { + + /// Copy with zero fill + MCTLASS_DEVICE + cp_async_zfill(void *smem_ptr, void const *global_ptr, bool pred_guard) { + #if CUDA_CP_ASYNC_ACTIVATED + + // Make sure the size is supported. + static_assert((SizeInBytes == 4 || SizeInBytes == 8 || SizeInBytes == 16), + "Size is not supported"); + + unsigned smem_int_ptr = mctlass_get_smem_pointer(smem_ptr); + int src_in_bytes = (pred_guard ? SizeInBytes : 0); + + asm volatile( +#if MCTLASS_ENABLE_L2_PREFETCH + "cp.async.ca.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), +#else + "cp.async.ca.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), +#endif + "l"(global_ptr), "n"(SizeInBytes), "r"(src_in_bytes)); + + #else + using AccessType = Array; + + if (pred_guard) { + *static_cast(smem_ptr) = *static_cast(global_ptr); + } + else { + AccessType zeros; + zeros.clear(); + *static_cast(smem_ptr) = zeros; + } + #endif + } +}; + +/// Partial specialization +template <> +struct cp_async_nan<16, CacheOperation::Always> { + static int const kSizeInBytes = 16; + + /// Copy with nan fill + MCTLASS_DEVICE + cp_async_nan(void *smem_ptr, void const *global_ptr, bool pred_guard) { + #if CUDA_CP_ASYNC_ACTIVATED + static __constant__ uint4 OOB_NAN_F16x8 = {OOB_NAN_F16x2, OOB_NAN_F16x2, + OOB_NAN_F16x2, OOB_NAN_F16x2}; + + unsigned smem_int_ptr = mctlass_get_smem_pointer(smem_ptr); + + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" +#if MCTLASS_ENABLE_L2_PREFETCH + " @p cp.async.ca.shared.global.L2::128B [%1], [%2], %3;\n" +#else + " @p cp.async.ca.shared.global [%1], [%2], %3;\n" +#endif + " @!p st.shared.v4.u32 [%1], {%4, %5, %6, %7};\n" + "}\n" + : + : "r"((int)pred_guard), "r"(smem_int_ptr), "l"(global_ptr), + "n"(kSizeInBytes), "r"(OOB_NAN_F16x8.x), "r"(OOB_NAN_F16x8.y), "r"(OOB_NAN_F16x8.z), + "r"(OOB_NAN_F16x8.w)); + #else + + MCTLASS_UNUSED(smem_ptr); + MCTLASS_UNUSED(global_ptr); + MCTLASS_UNUSED(pred_guard); + MCTLASS_NOT_IMPLEMENTED(); + + #endif + } +}; + +/// Partial specialization to write one (1) +template +struct cp_async_diag { + using Element = Element_; + + MCTLASS_DEVICE + cp_async_diag(void *smem_ptr) { + //unsigned smem_int_ptr = mctlass_get_smem_pointer(smem_ptr); + + // if (platform::is_same>::value) { + // asm volatile("st.shared.v4.u32 [%0], {%1, %2, %3, %4};\n" + // : : + // "r"(smem_int_ptr), "r"(DIAG_DATA_DOUBLE_ONE.y), "r"(DIAG_DATA_DOUBLE_ONE.x), + // "r"(DIAG_DATA_ZERO.x), "r"(DIAG_DATA_ZERO.x)); + // } else if (platform::is_same>::value) { + // asm volatile("st.shared.v2.u32 [%0], {%1, %2};\n" + // : : + // "r"(smem_int_ptr), "r"(DIAG_DATA_FLOAT_ONE.x), "r"(DIAG_DATA_ZERO.x)); + // } else if (platform::is_same::value) { + // asm volatile("st.shared.v2.u32 [%0], {%1, %2};\n" + // : : + // "r"(smem_int_ptr), "r"(DIAG_DATA_DOUBLE_ONE.y),"r"(DIAG_DATA_DOUBLE_ONE.x)); + // } else if (platform::is_same::value) { + // asm volatile("st.shared.u32 [%0], %1;\n" + // : : + // "r"(smem_int_ptr), "r"(DIAG_DATA_FLOAT_ONE.x)); + // } else { + // MCTLASS_UNUSED(smem_int_ptr); + // MCTLASS_NOT_IMPLEMENTED(); + // } + + if ((platform::is_same::value) || (platform::is_same::value)) { + *static_cast(smem_ptr) = static_cast(1.0f); + } else if (platform::is_same>::value) { + *static_cast *>(smem_ptr) = mctlass::complex(1.0f, 0.0f); + } else if (platform::is_same>::value) { + *static_cast *>(smem_ptr) = mctlass::complex(1.0f, 0.0f); + } else { + unsigned smem_int_ptr = mctlass_get_smem_pointer(smem_ptr); + MCTLASS_UNUSED(smem_int_ptr); + MCTLASS_NOT_IMPLEMENTED(); + } + } +}; + +/// Partial specialization to write zero for the imaginary part of Hermitian data +template +struct cp_async_diag { + using Element = Element_; + + MCTLASS_DEVICE + cp_async_diag(void *smem_ptr) { + #if CUDA_CP_ASYNC_ACTIVATED + static __constant__ uint1 DIAG_DATA_ZERO = {0x00000000}; + + unsigned smem_int_ptr = mctlass_get_smem_pointer(smem_ptr); + + if (platform::is_same>::value) { + asm volatile("st.shared.v2.u32 [%0], {%1, %2};\n" + : : + "r"(smem_int_ptr), "r"(DIAG_DATA_ZERO.x), "r"(DIAG_DATA_ZERO.x)); + } else if (platform::is_same>::value) { + asm volatile("st.shared.u32 [%0], %1;\n" + : : + "r"(smem_int_ptr), "r"(DIAG_DATA_ZERO.x)); + } else { + MCTLASS_UNUSED(smem_ptr); + MCTLASS_NOT_IMPLEMENTED(); + } + #else + if (platform::is_same>::value) { + *static_cast(smem_ptr) = 0.0f; + } else if (platform::is_same>::value) { + *static_cast(smem_ptr) = 0.0f; + } else { + MCTLASS_UNUSED(smem_ptr); + MCTLASS_NOT_IMPLEMENTED(); + } +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization +template < + /// Size of the access in bytes + int SizeInBytes> +struct cp_async { + + /// Copy + MCTLASS_DEVICE + cp_async(void *smem_ptr, void const *global_ptr, bool pred_guard = true) { + #if CUDA_CP_ASYNC_ACTIVATED + + static_assert(SizeInBytes == 16, + "cp.async only supports CacheOperation::Global when access size is 16B."); + + unsigned smem_int_ptr = mctlass_get_smem_pointer(smem_ptr); + + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" +#if MCTLASS_ENABLE_L2_PREFETCH + " @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n" +#else + " @p cp.async.cg.shared.global [%1], [%2], %3;\n" +#endif + "}\n" ::"r"((int)pred_guard), + "r"(smem_int_ptr), "l"(global_ptr), "n"(SizeInBytes)); + + #else + using AccessType = Array; + + if (pred_guard) { + *static_cast(smem_ptr) = *static_cast(global_ptr); + } + #endif + } +}; + +/// Partial specialization +template < + /// Size of the access in bytes + int SizeInBytes> +struct cp_async_zfill { + + /// Copy with zero fill + MCTLASS_DEVICE + cp_async_zfill(void *smem_ptr, void const *global_ptr, bool pred_guard = true) { + #if CUDA_CP_ASYNC_ACTIVATED + + static_assert(SizeInBytes == 16, + "cp.async only supports CacheOperation::Global when access size is 16B."); + + unsigned smem_int_ptr = mctlass_get_smem_pointer(smem_ptr); + int src_in_bytes = (pred_guard ? SizeInBytes : 0); + + asm volatile( +#if MCTLASS_ENABLE_L2_PREFETCH + "cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), +#else + "cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), +#endif + "l"(global_ptr), "n"(SizeInBytes), "r"(src_in_bytes)); + + #else + using AccessType = Array; + + if (pred_guard) { + *static_cast(smem_ptr) = *static_cast(global_ptr); + } + else { + AccessType zeros; + zeros.clear(); + *static_cast(smem_ptr) = zeros; + } + #endif + } +}; + +/// Partial specialization +template <> +struct cp_async_nan<16, CacheOperation::Global> { + static int const kSizeInBytes = 16; + + /// Copy with nan fill + MCTLASS_DEVICE + cp_async_nan(void *smem_ptr, void const *global_ptr, bool pred_guard) { + #if CUDA_CP_ASYNC_ACTIVATED + static __constant__ uint4 OOB_NAN_F16x8 = {OOB_NAN_F16x2, OOB_NAN_F16x2, + OOB_NAN_F16x2, OOB_NAN_F16x2}; + + unsigned smem_int_ptr = mctlass_get_smem_pointer(smem_ptr); + + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" +#if MCTLASS_ENABLE_L2_PREFETCH + " @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n" +#else + " @p cp.async.cg.shared.global [%1], [%2], %3;\n" +#endif + " @!p st.shared.v4.u32 [%1], {%4, %5, %6, %7};\n" + "}\n" + : + : "r"((int)pred_guard), "r"(smem_int_ptr), "l"(global_ptr), + "n"(kSizeInBytes), "r"(OOB_NAN_F16x8.x), "r"(OOB_NAN_F16x8.y), "r"(OOB_NAN_F16x8.z), + "r"(OOB_NAN_F16x8.w)); + #else + + MCTLASS_UNUSED(smem_ptr); + MCTLASS_UNUSED(global_ptr); + MCTLASS_UNUSED(pred_guard); + MCTLASS_NOT_IMPLEMENTED(); + + #endif + } +}; +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Establishes an ordering w.r.t previously issued cp.async instructions. Does not block. +MCTLASS_DEVICE +void cp_async_fence() { + #if CUDA_CP_ASYNC_ACTIVATED + asm volatile("cp.async.commit_group;\n" ::); + #endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Blocks until all but previous cp.async.commit_group operations have committed. +template +MCTLASS_DEVICE void cp_async_wait() { + #if CUDA_CP_ASYNC_ACTIVATED + asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); + #endif +} + +/// Blocks until all previous cp.async.commit_group operations have committed. +template <> +MCTLASS_DEVICE void cp_async_wait<0>() { + #if CUDA_CP_ASYNC_ACTIVATED + asm volatile("cp.async.wait_all;\n" ::); + #endif +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/arch/mma.h b/csrc/mctlass/include/mctlass/arch/mma.h new file mode 100644 index 0000000..ba6c7c0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/mma.h @@ -0,0 +1,233 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates exposing architecture support for multiply-add operations +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/functional.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/arch/arch.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace arch { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tag indicating the operation implied by MMA. +struct OpMultiplyAdd {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tag indicating the result is saturated to MAX_FLOAT|MIN_FLOAT or MAX_INT|MIN_INT +struct OpMultiplyAddSaturate {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tag indicating the input is converted to a narrower type (BF16) +struct OpMultiplyAddFastBF16 {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tag indicating the input is converted to a narrower type (F16) +struct OpMultiplyAddFastF16 {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tag indicating the input is converted to 2 (big and small) TF32 components +// Perform 3xTF32 or 4xTF32 for every F32 output element +struct OpMultiplyAddFastF32 {}; + +/// Tag indicating the input is converted to 2 (big and small) TF32 components +// Perform 3xTF32 or 4xTF32 for every complex output element +struct OpMultiplyAddComplexFastF32 {}; + +/// Helper for determining whether staged accumulation should be used for a given operator +template +struct UseStagedAccumulation { + static bool const value = platform::is_same::value || + platform::is_same::value; +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tag indicating the complex multiply-add operation +struct OpMultiplyAddComplex {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tag indicating the gaussian complex multiply-add operation +struct OpMultiplyAddGaussianComplex {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tag indicating the inner product is defined by (XOR, POPC) +struct OpXorPopc {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tag classifying math operators as thread-level operations. +struct OpClassSimt {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tag classifing operators as Tensor Core operations. +struct OpClassTensorOp {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Tag classifing operators as WMMA Tensor Core operations +struct OpClassWmmaTensorOp {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template < + /// Size of the matrix product (concept: GemmShape) + typename Shape_, + /// Number of threads participating + int kThreads_, + /// Data type of A elements + typename ElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Data type of B elements + typename ElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Element type of C matrix + typename ElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Inner product operator + typename Operator +> +struct Mma; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation - specialized for 1x1x1x1 matrix multiply operation +template < + /// Data type of A elements + typename ElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Data type of B elements + typename ElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Inner product operator + typename Operator_ +> +struct Mma, 1, ElementA, LayoutA, ElementB, LayoutB, ElementC_, LayoutC, Operator_> { + + using Shape = gemm::GemmShape<1, 1, 1>; + using Operator = Operator_; + using ElementC = ElementC_; + + MCTLASS_HOST_DEVICE + void operator()( + Array &d, + Array const &a, + Array const &b, + Array const &c + ) { + + multiply_add op; + + d[0] = op(a[0], b[0], c[0]); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specifies internal data type for computation +struct SPFormatType { + enum Kind { + Thread + }; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template < + /// Size of the matrix product (concept: GemmShape) + typename Shape_, + /// Number of threads participating + int kThreads_, + /// Data type of A elements + typename ElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Data type of B elements + typename ElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Element type of C matrix + typename ElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Inner product operator + typename Operator, + /// Specifies meta data format + SPFormatType::Kind SPFormat = SPFormatType::Thread +> +struct SparseMma; + +} // namespace arch +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Specializations for each compute capability +// + +#include "mctlass/arch/mma_sm50.h" +#include "mctlass/arch/mma_sm60.h" +#include "mctlass/arch/mma_sm61.h" +#include "mctlass/arch/mma_sm70.h" +#include "mctlass/arch/mma_sm75.h" +#include "mctlass/arch/mma_sm80.h" +#include "mctlass/arch/mma_sparse_sm80.h" +#include "mctlass/arch/mma_sm90.h" +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/arch/mma_sm50.h b/csrc/mctlass/include/mctlass/arch/mma_sm50.h new file mode 100644 index 0000000..17c8c6f --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/mma_sm50.h @@ -0,0 +1,432 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Matrix multiply +*/ + +#pragma once + +#include "mctlass/arch/mma.h" +#include "mctlass/complex.h" +#include "mctlass/quaternion.h" +#include "mctlass/functional.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/gemm/gemm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace arch { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template < + /// Layout of A matrix + typename LayoutA, + /// Layout of B matrix + typename LayoutB, + /// Layout of C matrix + typename LayoutC +> +struct Mma, 1, float, LayoutA, float, LayoutB, float, LayoutC, OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 1, 1>; + using Operator = OpMultiplyAdd; + using ElementC = float; + + MCTLASS_HOST_DEVICE + void operator()( + Array &d, + Array const &a, + Array const &b, + Array const &c + ) { + d[0] = a[0] * b[0] + c[0]; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template < + /// Layout of A matrix + typename LayoutA, + /// Layout of B matrix + typename LayoutB, + /// Layout of C matrix + typename LayoutC +> +struct Mma, 1, double, LayoutA, double, LayoutB, double, LayoutC, OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 1, 1>; + using Operator = OpMultiplyAdd; + using ElementC = double; + + MCTLASS_HOST_DEVICE + void operator()( + Array &d, + Array const &a, + Array const &b, + Array const &c + ) { + + d[0] = a[0] * b[0] + c[0]; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template < + /// Layout of A matrix + typename LayoutA, + /// Layout of B matrix + typename LayoutB, + /// Layout of C matrix + typename LayoutC +> +struct Mma, 1, int, LayoutA, int, LayoutB, int, LayoutC, OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 1, 1>; + using Operator = OpMultiplyAdd; + using ElementC = int; + + MCTLASS_HOST_DEVICE + void operator()( + Array &d, + Array const &a, + Array const &b, + Array const &c + ) { + + d[0] = a[0] * b[0] + c[0]; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template < + /// Layout of A matrix + typename LayoutA, + /// Layout of B matrix + typename LayoutB, + /// Layout of C matrix + typename LayoutC +> +struct Mma< + gemm::GemmShape<1, 1, 1>, + 1, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 1, 1>; + using Operator = OpMultiplyAddComplex; + using ElementC = complex; + + MCTLASS_HOST_DEVICE + void operator()( + Array, 1> &d, + Array, 1> const &a, + Array, 1> const &b, + Array, 1> const &c + ) { + + d[0].real() = a[0].real() * b[0].real() + c[0].real(); + d[0].imag() = a[0].imag() * b[0].real() + c[0].imag(); + d[0].real() = -a[0].imag() * b[0].imag() + d[0].real(); + d[0].imag() = a[0].real() * b[0].imag() + d[0].imag(); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template < + /// Layout of A matrix + typename LayoutA, + /// Layout of B matrix + typename LayoutB, + /// Layout of C matrix + typename LayoutC +> +struct Mma< + gemm::GemmShape<1, 1, 1>, + 1, + complex, + LayoutA, + float, + LayoutB, + complex, + LayoutC, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 1, 1>; + using Operator = OpMultiplyAddComplex; + using ElementC = complex; + + MCTLASS_HOST_DEVICE + void operator()( + Array, 1> &d, + Array, 1> const &a, + Array const &b, + Array, 1> const &c + ) { + + d[0].real() = a[0].real() * b[0] + c[0].real(); + d[0].imag() = a[0].imag() * b[0] + c[0].imag(); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template < + /// Layout of A matrix + typename LayoutA, + /// Layout of B matrix + typename LayoutB, + /// Layout of C matrix + typename LayoutC +> +struct Mma< + gemm::GemmShape<1, 1, 1>, + 1, + float, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 1, 1>; + using Operator = OpMultiplyAddComplex; + using ElementC = complex; + + MCTLASS_HOST_DEVICE + void operator()( + Array, 1> &d, + Array const &a, + Array, 1> const &b, + Array, 1> const &c + ) { + + d[0].real() = a[0] * b[0].real() + c[0].real(); + d[0].imag() = a[0] * b[0].imag() + d[0].imag(); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template < + /// Layout of A matrix + typename LayoutA, + /// Layout of B matrix + typename LayoutB, + /// Layout of C matrix + typename LayoutC +> +struct Mma< + gemm::GemmShape<1, 1, 1>, + 1, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 1, 1>; + using Operator = OpMultiplyAddComplex; + using ElementC = complex; + + MCTLASS_HOST_DEVICE + void operator()( + Array, 1> &d, + Array, 1> const &a, + Array, 1> const &b, + Array, 1> const &c + ) { + + d[0].real() = a[0].real() * b[0].real() + c[0].real(); + d[0].imag() = a[0].imag() * b[0].real() + c[0].imag(); + d[0].real() = -a[0].imag() * b[0].imag() + d[0].real(); + d[0].imag() = a[0].real() * b[0].imag() + d[0].imag(); + } +}; + +/// Matrix multiply-add operation +template < + /// Layout of A matrix + typename LayoutA, + /// Layout of B matrix + typename LayoutB, + /// Layout of C matrix + typename LayoutC +> +struct Mma< + gemm::GemmShape<1, 1, 1>, + 1, + complex, + LayoutA, + double, + LayoutB, + complex, + LayoutC, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 1, 1>; + using Operator = OpMultiplyAddComplex; + using ElementC = complex; + + MCTLASS_HOST_DEVICE + void operator()( + Array, 1> &d, + Array, 1> const &a, + Array const &b, + Array, 1> const &c + ) { + + d[0].real() = a[0].real() * b[0] + c[0].real(); + d[0].imag() = a[0].imag() * b[0] + c[0].imag(); + } +}; + +/// Matrix multiply-add operation +template < + /// Layout of A matrix + typename LayoutA, + /// Layout of B matrix + typename LayoutB, + /// Layout of C matrix + typename LayoutC +> +struct Mma< + gemm::GemmShape<1, 1, 1>, + 1, + double, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 1, 1>; + using Operator = OpMultiplyAddComplex; + using ElementC = complex; + + MCTLASS_HOST_DEVICE + void operator()( + Array, 1> &d, + Array const &a, + Array, 1> const &b, + Array, 1> const &c + ) { + + d[0].real() = a[0] * b[0].real() + c[0].real(); + d[0].imag() = a[0] * b[0].imag() + d[0].imag(); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template < + /// Layout of A matrix + typename LayoutA, + /// Layout of B matrix + typename LayoutB, + /// Layout of C matrix + typename LayoutC +> +struct Mma, 1, half_t, LayoutA, half_t, LayoutB, float, LayoutC, OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 1, 1>; + using Operator = OpMultiplyAdd; + using ElementC = float; + + MCTLASS_HOST_DEVICE + void operator()( + Array &d, + Array const &a, + Array const &b, + Array const &c + ) { + d[0] = float(a[0]) * float(b[0]) + c[0]; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation for Quaternions +template < + /// Layout of A matrix + typename LayoutA, + /// Layout of B matrix + typename LayoutB, + /// Layout of C matrix + typename LayoutC +> +struct Mma, 1, Quaternion, LayoutA, Quaternion, LayoutB, Quaternion, LayoutC, OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 1, 1>; + using Operator = OpMultiplyAdd; + using Element = Quaternion; + using ElementC = Element; + + MCTLASS_HOST_DEVICE + void operator()( + Array &d, + Array const &a, + Array const &b, + Array const &c + ) { + multiply_add op; + d[0] = op(a[0], b[0], c[0]); + } + +}; + +} +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/arch/mma_sm60.h b/csrc/mctlass/include/mctlass/arch/mma_sm60.h new file mode 100644 index 0000000..53e0eca --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/mma_sm60.h @@ -0,0 +1,252 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Matrix multiply +*/ + +#pragma once + +#include + +#include "mctlass/arch/mma.h" + +#include "mctlass/layout/matrix.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace arch { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template +struct Mma< + gemm::GemmShape<2,1,1>, + 1, + half_t, + LayoutA, + half_t, + LayoutB, + half_t, + LayoutC, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<2, 1, 1>; + using Operator = OpMultiplyAdd; + using ElementC = half_t; + + MCTLASS_HOST_DEVICE + void operator()( + Array &d, + Array const &a, + Array const &b, + Array const &c + ) { + +#if defined(__MACA_ARCH__) + + __half2 const & A = reinterpret_cast<__half2 const &>(a); + __half2 B = __half2half2(reinterpret_cast<__half const &>(b)); + __half2 const & C = reinterpret_cast<__half2 const &>(c); + + __half2 D = __hfma2(A, B, C); + + d = reinterpret_cast &>(D); + +#else + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 2; ++i) { + d[i] = a[i] * b[0] + c[i]; + } +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template +struct Mma< + gemm::GemmShape<1,2,1>, + 1, + half_t, + LayoutA, + half_t, + LayoutB, + half_t, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 2, 1>; + using Operator = OpMultiplyAdd; + using ElementC = half_t; + + MCTLASS_HOST_DEVICE + void operator()( + Array &d, + Array const &a, + Array const &b, + Array const &c + ) { + +#if defined(__MACA_ARCH__) + + __half2 const & A = __half2half2(reinterpret_cast<__half const &>(a)); + __half2 B = reinterpret_cast<__half2 const &>(b); + __half2 const & C = reinterpret_cast<__half2 const &>(c); + + __half2 D = __hfma2(A, B, C); + + d = reinterpret_cast &>(D); + +#else + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 2; ++i) { + d[i] = a[0] * b[i] + c[i]; + } +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template <> +struct Mma < + gemm::GemmShape<2, 2, 1>, + 1, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<2, 2, 1>; + using Operator = OpMultiplyAdd; + using ElementC = half_t; + + MCTLASS_HOST_DEVICE + void operator()( + Array &d, + Array const &a, + Array const &b, + Array const &c + ) { + +#if defined(__MACA_ARCH__) + + __half2 const & A = reinterpret_cast<__half2 const &>(a); + __half2 Blo = __low2half2(reinterpret_cast<__half2 const &>(b)); + __half2 Bhi = __high2half2(reinterpret_cast<__half2 const &>(b)); + + __half2 const *C = reinterpret_cast<__half2 const *>(&c); + + __half2 Dlo = __hfma2(A, Blo, C[0]); + __half2 Dhi = __hfma2(A, Bhi, C[1]); + + Array * D = reinterpret_cast *>(&d); + + D[0] = reinterpret_cast const &>(Dlo); + D[1] = reinterpret_cast const &>(Dhi); + +#else + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < 2; ++j) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 2; ++i) { + d[i + 2 * j] = a[i] * b[j] + c[i + 2 * j]; + } + } +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template <> +struct Mma< + gemm::GemmShape<2, 2, 1>, + 1, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + half_t, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<2, 2, 1>; + using Operator = OpMultiplyAdd; + using ElementC = half_t; + + MCTLASS_HOST_DEVICE + void operator()( + Array &d, + Array const &a, + Array const &b, + Array const &c + ) { + +#if defined(__MACA_ARCH__) + + __half2 Alo = __low2half2(reinterpret_cast<__half2 const &>(a)); + __half2 Ahi = __high2half2(reinterpret_cast<__half2 const &>(a)); + __half2 const & B = reinterpret_cast<__half2 const &>(b); + + __half2 const *C = reinterpret_cast<__half2 const *>(&c); + + __half2 Dlo = __hfma2(Alo, B, C[0]); + __half2 Dhi = __hfma2(Ahi, B, C[0]); + + Array * D = reinterpret_cast *>(&d); + + D[0] = reinterpret_cast &>(Dlo); + D[1] = reinterpret_cast &>(Dhi); +#else + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 2; ++i) { + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < 2; ++j) { + d[i * 2 + j] = a[i] * b[j] + c[i * 2 + j]; + } + } +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} +} diff --git a/csrc/mctlass/include/mctlass/arch/mma_sm61.h b/csrc/mctlass/include/mctlass/arch/mma_sm61.h new file mode 100644 index 0000000..825716a --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/mma_sm61.h @@ -0,0 +1,144 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Matrix multiply +*/ + +#pragma once + +#include "mctlass/layout/matrix.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace arch { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template +struct Mma< + gemm::GemmShape<1,1,4>, + 1, + int8_t, + LayoutA, + int8_t, + LayoutB, + int, + LayoutC, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 1, 4>; + using Operator = OpMultiplyAdd; + using ElementC = int; + + MCTLASS_HOST_DEVICE + void operator()( + Array &d, + Array const &a, + Array const &b, + Array const &c + ) { + +//#if defined(__MACA_ARCH__) +#if 0 + + unsigned const &A = reinterpret_cast(a); + unsigned const &B = reinterpret_cast(b); + + asm volatile("dp4a.s32.s32 %0, %1, %2, %3;" + : "=r"(d[0]) + : "r"(A), "r"(B), "r"(c[0])); + +#else + + d[0] = c[0]; + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < 4; ++k) { + d[0] += a[k] * b[k]; + } + +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template +struct Mma< + gemm::GemmShape<1, 1, 2>, + 1, + int16_t, + layout::RowMajor, + int16_t, + layout::ColumnMajor, + int, + LayoutC, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<1, 1, 2>; + using Operator = OpMultiplyAdd; + using ElementC = int; + + MCTLASS_HOST_DEVICE + void operator()( + Array &d, + Array const &a, + Array const &b, + Array const &c + ) { + +//#if defined(__MACA_ARCH__) +#if 0 + + unsigned const &A = reinterpret_cast(a); + unsigned const &B = reinterpret_cast(b); + + asm volatile("dp2a.s32.s32 %0, %1, %2, %3;" + : "=r"(d[0]) + : "r"(A), "r"(B), "r"(c[0])); +#else + d[0] = c[0]; + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < 2; ++k) { + d[0] += a[k] * b[k]; + } +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} +} diff --git a/csrc/mctlass/include/mctlass/arch/mma_sm70.h b/csrc/mctlass/include/mctlass/arch/mma_sm70.h new file mode 100644 index 0000000..371a26a --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/mma_sm70.h @@ -0,0 +1,659 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Matrix multiply +*/ +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mma.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" + +#if ((__CUDACC_VER_MAJOR__ > 10) || (__CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 1)) +#define MCTLASS_ARCH_MMA_SM70_SUPPORTED +#endif + +#if (defined(__MACA_ARCH__)) + +#if ((__CUDACC_VER_MAJOR__ > 10) || (__CUDACC_VER_MAJOR__ == 10 &&__CUDACC_VER_MINOR__ >= 1)) +#define MCTLASS_ARCH_MMA_SM70_ENABLED +#endif + +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace arch { + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Matrix multiply accumulate 884 - FP16 accumulation +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F16 = F16 * F16 + F16 +template <> +struct Mma< + gemm::GemmShape<8,8,4>, + 8, + half_t, + layout::ColumnMajor, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8, 8, 4>; + + using ElementA = half_t; + using LayoutA = layout::ColumnMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = half_t; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm70; + + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) { + +#if defined(MCTLASS_ARCH_MMA_SM70_ENABLED) + + unsigned const *A = reinterpret_cast(&a); + unsigned const *B = reinterpret_cast(&b); + unsigned const *C = reinterpret_cast(&c); + unsigned *D = reinterpret_cast(&d); + + asm volatile("mma.sync.aligned.m8n8k4.col.col.f16.f16.f16.f16 {%0,%1,%2,%3}, {%4,%5}, {%6,%7}, {%8,%9,%10,%11};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]) + ); + +#else + printf("mma_sm70.h L120 this function cannot working correctly now.\n"); + assert(0); + #if defined(__MACA_ARCH__) + //asm volatile ("brkpt;\n" ::); + #endif +#endif + } +}; + +/// Matrix multiply-add operation: F16 = F16 * F16 + F16 +template <> +struct Mma< + gemm::GemmShape<8, 8, 4>, + 8, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + half_t, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8, 8, 4>; + + using ElementA = half_t; + using LayoutA = layout::ColumnMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::RowMajor; + using FragmentB = Array; + + using ElementC = half_t; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm70; + + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) { + +#if defined(MCTLASS_ARCH_MMA_SM70_ENABLED) + + unsigned const *A = reinterpret_cast(&a); + unsigned const *B = reinterpret_cast(&b); + unsigned const *C = reinterpret_cast(&c); + unsigned *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k4.col.row.f16.f16.f16.f16 {%0,%1,%2,%3}, {%4,%5}, {%6,%7}, {%8,%9,%10,%11};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]) + ); +#else + printf("mma_sm70.h L178 this function cannot working correctly now.\n"); + assert(0); + #if defined(__MACA_ARCH__) + //asm volatile ("brkpt;\n" ::); + #endif +#endif + } +}; + +/// Matrix multiply-add operation: F16 = F16 * F16 + F16 +template <> +struct Mma< + gemm::GemmShape<8, 8, 4>, + 8, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8, 8, 4>; + + using ElementA = half_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = half_t; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm70; + + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) { + +#if defined(MCTLASS_ARCH_MMA_SM70_ENABLED) + + unsigned const *A = reinterpret_cast(&a); + unsigned const *B = reinterpret_cast(&b); + unsigned const *C = reinterpret_cast(&c); + unsigned *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k4.row.col.f16.f16.f16.f16 {%0,%1,%2,%3}, {%4,%5}, {%6,%7}, {%8,%9,%10,%11};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]) + ); +#else + printf("mma_sm70.h L236 this function cannot working correctly now.\n"); + assert(0); + #if defined(__MACA_ARCH__) + // asm volatile ("brkpt;\n" ::); + #endif +#endif + } +}; + +/// Matrix multiply-add operation: F16 = F16 * F16 + F16 +template <> +struct Mma< + gemm::GemmShape<8, 8, 4>, + 8, + half_t, + layout::RowMajor, + half_t, + layout::RowMajor, + half_t, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8, 8, 4>; + + using ElementA = half_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::RowMajor; + using FragmentB = Array; + + using ElementC = half_t; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm70; + + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) { + +#if defined(MCTLASS_ARCH_MMA_SM70_ENABLED) + + unsigned const *A = reinterpret_cast(&a); + unsigned const *B = reinterpret_cast(&b); + unsigned const *C = reinterpret_cast(&c); + unsigned *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k4.row.row.f16.f16.f16.f16 {%0,%1,%2,%3}, {%4,%5}, {%6,%7}, {%8,%9,%10,%11};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]) + ); +#else + printf("mma_sm70.h L294 this function cannot working correctly now"); + assert(0); + #if defined(__MACA_ARCH__) + //asm volatile ("brkpt;\n" ::); + #endif +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Matrix multiply accumulate 884 - FP32 accumulation +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F32 = F16 * F16 + F32 +template <> +struct Mma< + gemm::GemmShape<8, 8, 4>, + 8, + half_t, + layout::ColumnMajor, + half_t, + layout::ColumnMajor, + float, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8, 8, 4>; + + using ElementA = half_t; + using LayoutA = layout::ColumnMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm70; + + /// Multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) { + +#if defined(MCTLASS_ARCH_MMA_SM70_ENABLED) + + unsigned const *A = reinterpret_cast(&a); + unsigned const *B = reinterpret_cast(&b); + float const *C = reinterpret_cast(&c); + float *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k4.col.col.f32.f16.f16.f32 {%0,%1,%2,%3,%4,%5,%6,%7}, {%8,%9}, {%10,%11}, " + "{%12,%13,%14,%15,%16,%17,%18,%19};\n" + : "=f"(D[0]), + "=f"(D[1]), + "=f"(D[2]), + "=f"(D[3]), + "=f"(D[4]), + "=f"(D[5]), + "=f"(D[6]), + "=f"(D[7]) + : "r"(A[0]), + "r"(A[1]), + "r"(B[0]), + "r"(B[1]), + "f"(C[0]), + "f"(C[1]), + "f"(C[2]), + "f"(C[3]), + "f"(C[4]), + "f"(C[5]), + "f"(C[6]), + "f"(C[7]) + ); +#else + printf("mma_sm70.h L378 this function cannot working correctly now.\n"); + assert(0); + #if defined(__MACA_ARCH__) + //asm volatile ("brkpt;\n" ::); + #endif +#endif + } +}; + +/// Matrix multiply-add operation: F32 = F16 * F16 + F32 +template <> +struct Mma< + gemm::GemmShape<8, 8, 4>, + 8, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + float, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8, 8, 4>; + + using ElementA = half_t; + using LayoutA = layout::ColumnMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::RowMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm70; + + /// Multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) { + +#if defined(MCTLASS_ARCH_MMA_SM70_ENABLED) + + unsigned const *A = reinterpret_cast(&a); + unsigned const *B = reinterpret_cast(&b); + float const *C = reinterpret_cast(&c); + float *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k4.col.row.f32.f16.f16.f32 {%0,%1,%2,%3,%4,%5,%6,%7}, {%8,%9}, {%10,%11}, " + "{%12,%13,%14,%15,%16,%17,%18,%19};\n" + : "=f"(D[0]), + "=f"(D[1]), + "=f"(D[2]), + "=f"(D[3]), + "=f"(D[4]), + "=f"(D[5]), + "=f"(D[6]), + "=f"(D[7]) + : "r"(A[0]), + "r"(A[1]), + "r"(B[0]), + "r"(B[1]), + "f"(C[0]), + "f"(C[1]), + "f"(C[2]), + "f"(C[3]), + "f"(C[4]), + "f"(C[5]), + "f"(C[6]), + "f"(C[7]) + ); +#else + printf("mma_sm70.h L456 this function cannot working correctly now.\n"); + assert(0); + #if defined(__MACA_ARCH__) + //asm volatile ("brkpt;\n" ::); + #endif +#endif + } +}; + +/// Matrix multiply-add operation: F32 = F16 * F16 + F32 +template <> +struct Mma< + gemm::GemmShape<8, 8, 4>, + 8, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + float, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8, 8, 4>; + + using ElementA = half_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm70; + + /// Multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) { + +#if defined(MCTLASS_ARCH_MMA_SM70_ENABLED) + + unsigned const *A = reinterpret_cast(&a); + unsigned const *B = reinterpret_cast(&b); + float const *C = reinterpret_cast(&c); + float *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 {%0,%1,%2,%3,%4,%5,%6,%7}, {%8,%9}, {%10,%11}, " + "{%12,%13,%14,%15,%16,%17,%18,%19};\n" + : "=f"(D[0]), + "=f"(D[1]), + "=f"(D[2]), + "=f"(D[3]), + "=f"(D[4]), + "=f"(D[5]), + "=f"(D[6]), + "=f"(D[7]) + : "r"(A[0]), + "r"(A[1]), + "r"(B[0]), + "r"(B[1]), + "f"(C[0]), + "f"(C[1]), + "f"(C[2]), + "f"(C[3]), + "f"(C[4]), + "f"(C[5]), + "f"(C[6]), + "f"(C[7]) + ); +#else + printf("mma_sm70.h L534 this function cannot working correctly now.\n"); + assert(0); + #if defined(__MACA_ARCH__) + //asm volatile ("brkpt;\n" ::); + #endif +#endif + } +}; + +/// Matrix multiply-add operation: F32 = F16 * F16 + F32 +template <> +struct Mma< + gemm::GemmShape<8, 8, 4>, + 8, + half_t, + layout::RowMajor, + half_t, + layout::RowMajor, + float, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8, 8, 4>; + + using ElementA = half_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::RowMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm70; + + /// Multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) { + +#if defined(MCTLASS_ARCH_MMA_SM70_ENABLED) + + unsigned const *A = reinterpret_cast(&a); + unsigned const *B = reinterpret_cast(&b); + float const *C = reinterpret_cast(&c); + float *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k4.row.row.f32.f16.f16.f32 {%0,%1,%2,%3,%4,%5,%6,%7}, {%8,%9}, {%10,%11}, " + "{%12,%13,%14,%15,%16,%17,%18,%19};\n" + : "=f"(D[0]), + "=f"(D[1]), + "=f"(D[2]), + "=f"(D[3]), + "=f"(D[4]), + "=f"(D[5]), + "=f"(D[6]), + "=f"(D[7]) + : "r"(A[0]), + "r"(A[1]), + "r"(B[0]), + "r"(B[1]), + "f"(C[0]), + "f"(C[1]), + "f"(C[2]), + "f"(C[3]), + "f"(C[4]), + "f"(C[5]), + "f"(C[6]), + "f"(C[7]) + ); +#else + printf("mma_sm70.h L612 this function cannot working correctly now.\n"); + assert(0); + #if defined(__MACA_ARCH__) + //asm volatile ("brkpt;\n" ::); + #endif +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation specialized for the entire warp +template < + typename LayoutA, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename Operator +> +struct Mma< + gemm::GemmShape<16, 16, 4>, + 32, + half_t, + LayoutA, + half_t, + LayoutB, + ElementC, + LayoutC, + Operator +> : + public Mma< + gemm::GemmShape<8, 8, 4>, + 8, + half_t, + LayoutA, + half_t, + LayoutB, + ElementC, + LayoutC, + Operator> { + + using Shape = gemm::GemmShape<16, 16, 4>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/arch/mma_sm75.h b/csrc/mctlass/include/mctlass/arch/mma_sm75.h new file mode 100644 index 0000000..703365e --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/mma_sm75.h @@ -0,0 +1,1285 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Matrix multiply for SM75 +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mctlass/arch/wmma.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +//#include +#include "mctlass/wmma_array.h" +#endif + +// MCTLASS includes +#include "mctlass/arch/mma.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" + +//////////////////////////////////////////////////////////////////////////////// + +#if ((__CUDACC_VER_MAJOR__ > 10) || (__CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ >= 2)) + +#define MCTLASS_ARCH_MMA_SM75_SUPPORTED 1 + +#if (defined(__MACA_ARCH__)) +#define MCTLASS_ARCH_MMA_SM75_ENABLED +#endif +#endif + +#if (defined(__MACA_ARCH__)) +#if (!defined(MCTLASS_ARCH_MMA_SM80_ENABLED)) +#define MCTLASS_ARCH_MMA_SM80_ENABLED +#endif +#endif + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace arch { + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 1688 - FP16 accumulation +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation - F16 = F16 * F16 + F16 +template <> +struct Mma< + gemm::GemmShape<16, 8, 8>, + 32, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16, 8, 8>; + + using ElementA = half_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = half_t; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm75; + + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const *A = reinterpret_cast(&a); + unsigned const *B = reinterpret_cast(&b); + unsigned const *C = reinterpret_cast(&c); + unsigned *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0,%1}, {%2,%3}, {%4}, {%5,%6};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A[0]), "r"(A[1]), "r"(B[0]), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 1688 - FP32 accumulation +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F32 = F16 * F16 + F32 +template <> +struct Mma< + gemm::GemmShape<16, 8, 8>, + 32, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + float, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16, 8, 8>; + + using ElementA = half_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c) const { + +#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +// +// unsigned const *A = reinterpret_cast(&a); +// unsigned const *B = reinterpret_cast(&b); +// float const *C = reinterpret_cast(&c); +// float *D = reinterpret_cast(&d); +// asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" +// : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) +// : +// "r"(A[0]), "r"(A[1]), +// "r"(B[0]), +// "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3]) +// ); + auto results = __builtin_mxc_mma_16x16x16f16( + {a[0].get().to_macahalf(), a[1].get().to_macahalf(), 0, 0}, + {b[0].get().to_macahalf(), b[1].get().to_macahalf(), 0, 0}, + {c[0], c[1], c[2], c[3]}); + + d[0] = results[0]; + d[1] = results[1]; + d[2] = results[2]; + d[3] = results[3]; +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Integer matrix multiply .8816 (8b) +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S8 * S8 + S32 +template <> +struct Mma< + gemm::GemmShape<8, 8, 16>, + 32, + int8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8, 8, 16>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * S8 + S32 +template <> +struct Mma< + gemm::GemmShape<8, 8, 16>, + 32, + uint8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8, 8, 16>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k16.row.col.s32.u8.s8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S8 * U8 + S32 +template <> +struct Mma< + gemm::GemmShape<8, 8, 16>, + 32, + int8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8, 8, 16>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k16.row.col.s8.u8 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); + +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * U8 + S32 +template <> +struct Mma< + gemm::GemmShape<8, 8, 16>, + 32, + uint8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8, 8, 16>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k16.row.col.s32.u8.u8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Integer matrix multiply (8b) with SATURATE +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S8 * S8 + S32 +template <> +struct Mma< + gemm::GemmShape<8,8,16>, + 32, + int8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<8,8,16>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k16.row.col.satfinite.s32.s8.s8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * S8 + S32 +template <> +struct Mma< + gemm::GemmShape<8,8,16>, + 32, + uint8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<8,8,16>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k16.row.col.satfinite.s32.u8.s8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S8 * U8 + S32 +template <> +struct Mma< + gemm::GemmShape<8,8,16>, + 32, + int8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<8,8,16>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k16.row.col.satfinite.s32.s8.u8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * U8 + S32 +template <> +struct Mma< + gemm::GemmShape<8,8,16>, + 32, + uint8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<8,8,16>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k16.row.col.satfinite.s32.u8.u8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Integer matrix multiply (4b) +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S4 * S4 + S32 +template <> +struct Mma< + gemm::GemmShape<8,8,32>, + 32, + int4b_t, + layout::RowMajor, + int4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8,8,32>; + + using ElementA = int4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k32.row.col.s32.s4.s4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U4 * S4 + S32 +template <> +struct Mma< + gemm::GemmShape<8,8,32>, + 32, + uint4b_t, + layout::RowMajor, + int4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8,8,32>; + + using ElementA = uint4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k32.row.col.s32.u4.s4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S4 * U4 + S32 +template <> +struct Mma< + gemm::GemmShape<8,8,32>, + 32, + int4b_t, + layout::RowMajor, + uint4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8,8,32>; + + using ElementA = int4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k32.row.col.s32.s4.u4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U4 * U4 + S32 +template <> +struct Mma< + gemm::GemmShape<8,8,32>, + 32, + uint4b_t, + layout::RowMajor, + uint4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8,8,32>; + + using ElementA = uint4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k32.row.col.s32.u4.u4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Integer matrix multiply (4b) - SATURATE +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S4 * S4 + S32 +template <> +struct Mma< + gemm::GemmShape<8,8,32>, + 32, + int4b_t, + layout::RowMajor, + int4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<8,8,32>; + + using ElementA = int4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k32.row.col.satfinite.s32.s4.s4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U4 * S4 + S32 +template <> +struct Mma< + gemm::GemmShape<8,8,32>, + 32, + uint4b_t, + layout::RowMajor, + int4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<8,8,32>; + + using ElementA = uint4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k32.row.col.satfinite.s32.u4.s4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S4 * U4 + S32 +template <> +struct Mma< + gemm::GemmShape<8,8,32>, + 32, + int4b_t, + layout::RowMajor, + uint4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<8,8,32>; + + using ElementA = int4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k32.row.col.satfinite.s32.s4.u4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U4 * U4 + S32 +template <> +struct Mma< + gemm::GemmShape<8,8,32>, + 32, + uint4b_t, + layout::RowMajor, + uint4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<8,8,32>; + + using ElementA = uint4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) + + unsigned const & A = reinterpret_cast(a); + unsigned const & B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile("mma.sync.aligned.m8n8k32.row.col.satfinite.s32.u4.u4.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A), "r"(B), "r"(C[0]), "r"(C[1])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// b1 ^ b1 + s32 => s32 +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation +template <> +struct Mma< + gemm::GemmShape<8,8,128>, + 32, + uint1b_t, + layout::RowMajor, + uint1b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpXorPopc> { + + using Shape = gemm::GemmShape<8,8,128>; + + using ElementA = uint1b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint1b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpXorPopc; + using ArchTag = arch::Sm75; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM75_ENABLED) +#if 0 //mcTlass not support int1b_t +#if defined(MCTLASS_ARCH_WMMA_ENABLED) + using WmmaFragmentA = mxmaca::wmma::fragment< + mxmaca::wmma::matrix_a, + Shape::kM, + Shape::kN, + Shape::kK, + mxmaca::wmma::experimental::precision::b1, + mxmaca::wmma::row_major>; + + using WmmaFragmentB = mxmaca::wmma::fragment< + mxmaca::wmma::matrix_b, + Shape::kM, + Shape::kN, + Shape::kK, + mxmaca::wmma::experimental::precision::b1, + mxmaca::wmma::col_major>; + + using WmmaFragmentC = mxmaca::wmma::fragment< + mxmaca::wmma::accumulator, + Shape::kM, + Shape::kN, + Shape::kK, + int>; + + WmmaFragmentA const & A = reinterpret_cast(a); + WmmaFragmentB const & B = reinterpret_cast(b); + + WmmaFragmentC const & C = reinterpret_cast(c); + WmmaFragmentC & D = reinterpret_cast(d); + + mxmaca::wmma::bmma_sync(D, A, B, C, mxmaca::wmma::experimental::bmmaBitOpXOR, + mxmaca::wmma::experimental::bmmaAccumulateOpPOPC); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); // WMMA must be supported to issue binary matrix multiply-accumulate instructions. + +#endif // defined(MCTLASS_ARCH_WMMA_ENABLED) + +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/arch/mma_sm80.h b/csrc/mctlass/include/mctlass/arch/mma_sm80.h new file mode 100644 index 0000000..605fed4 --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/mma_sm80.h @@ -0,0 +1,2339 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Matrix multiply +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mctlass/mctlass.h" +#include "mma.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" + + +//////////////////////////////////////////////////////////////////////////////// + +//#if ((__CUDACC_VER_MAJOR__ > 11) || (__CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ >= 0)) +#if 1 +#define MCTLASS_ARCH_MMA_SM80_SUPPORTED 1 + +#if defined(__MACA_ARCH__) +#define MCTLASS_ARCH_MMA_SM80_ENABLED +#include "__clang_maca_mma_functions.h" +#endif +#endif + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace arch { + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 1688 - Float BF16, FP32 accumulation +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation - F32 = bf16 * bf16 + F32 +template <> +struct Mma< + gemm::GemmShape<16, 8, 8>, + 32, + bfloat16_t, + layout::RowMajor, + bfloat16_t, + layout::ColumnMajor, + float, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16, 8, 8>; + + using ElementA = bfloat16_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = bfloat16_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + MCTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + float const *C = reinterpret_cast(&c); + float *D = reinterpret_cast(&d); + asm( + "mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" + : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) + : + "r"(A[0]), "r"(A[1]), + "r"(B[0]), + "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3]) + ); +#else + + MCTLASS_UNUSED(d); + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_NOT_IMPLEMENTED(); + +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 1684 - Float TF32 +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F32 = tf32 * tf32 + F32 +template <> +struct Mma< + gemm::GemmShape<16, 8, 4>, + 32, + tfloat32_t, + layout::RowMajor, + tfloat32_t, + layout::ColumnMajor, + float, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16, 8, 4>; + + using ElementA = tfloat32_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = tfloat32_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + + // uint32_t const *A = reinterpret_cast(&a); + // uint32_t const *B = reinterpret_cast(&b); + // float const *C = reinterpret_cast(&c); + // float *D = reinterpret_cast(&d); + // asm volatile( + // "mma.sync.aligned.m16n8k4.row.col.f32.tf32.tf32.f32 {%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" + // : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) + // : + // "r"(A[0]), "r"(A[1]), + // "r"(B[0]), + // "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3]) + // ); +#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) + float const *A = reinterpret_cast(&a); + float const *B = reinterpret_cast(&b); + float const *C = reinterpret_cast(&c); + float *D = reinterpret_cast(&d); + + auto temp = __builtin_mxc_mma_16x16x8tf32({A[0], 0}, {B[0], 0}, {C[0], C[1], C[2], C[3]}); + + D[0] = temp[0]; + D[1] = temp[1]; + D[2] = temp[2]; + D[3] = temp[3]; +#else + + MCTLASS_UNUSED(d); + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_NOT_IMPLEMENTED(); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 1688 - Float TF32 +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F32 = tf32 * tf32 + F32 +template <> +struct Mma, 32, tfloat32_t, layout::RowMajor, + tfloat32_t, layout::ColumnMajor, float, layout::RowMajor, + OpMultiplyAdd> { + using Shape = gemm::GemmShape<16, 8, 8>; + + using ElementA = tfloat32_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = tfloat32_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + MCTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c) const { + +// #if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) + +// uint32_t const *A = reinterpret_cast(&a); +// uint32_t const *B = reinterpret_cast(&b); +// float const *C = reinterpret_cast(&c); +// float *D = reinterpret_cast(&d); +// asm volatile( +// "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 " +// "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" +// : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) +// : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), +// "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3])); +// #endif +#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) + float const *fA = reinterpret_cast(&a); + float const *fB = reinterpret_cast(&b); + float const *C = reinterpret_cast(&c); + float *D = reinterpret_cast(&d); + + auto temp = __builtin_mxc_mma_16x16x8tf32({fA[0], fA[1]}, {fB[0], fB[1]}, {C[0], C[1], C[2], C[3]}); + + D[0] = temp[0]; + D[1] = temp[1]; + D[2] = temp[2]; + D[3] = temp[3]; +#else + + MCTLASS_UNUSED(d); + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_NOT_IMPLEMENTED(); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 16816 +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F16 = F16 * F16 + F16 +template <> +struct Mma< + gemm::GemmShape<16, 8, 16>, + 32, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16, 8, 16>; + + using ElementA = half_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = half_t; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c) const { + +#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +// +// uint32_t const *A = reinterpret_cast(&a); +// uint32_t const *B = reinterpret_cast(&b); +// uint32_t const *C = reinterpret_cast(&c); +// uint32_t *D = reinterpret_cast(&d); +// asm volatile("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n" +// : "=r"(D[0]), "=r"(D[1]) +// : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), +// "r"(B[0]), "r"(B[1]), +// "r"(C[0]), "r"(C[1]) +// ); + auto results = __builtin_mxc_mma_16x16x16f16( + {a[0].get().to_macahalf(), a[1].get().to_macahalf(), + a[2].get().to_macahalf(), a[3].get().to_macahalf()}, + {b[0].get().to_macahalf(), b[1].get().to_macahalf(), + b[2].get().to_macahalf(), b[3].get().to_macahalf()}, + {c[0].get().to_macahalf(), c[1].get().to_macahalf(), + c[2].get().to_macahalf(), c[3].get().to_macahalf()}); + + d[0] = mctlass::half_t(results[0]); + d[1] = mctlass::half_t(results[1]); + d[2] = mctlass::half_t(results[2]); + d[3] = mctlass::half_t(results[3]); +#else + + MCTLASS_UNUSED(d); + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_NOT_IMPLEMENTED(); + +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F32 = bf16 * bf16 + F32 +template <> +struct Mma< + gemm::GemmShape<16, 8, 16>, + 32, + bfloat16_t, + layout::RowMajor, + bfloat16_t, + layout::ColumnMajor, + float, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16, 8, 16>; + + using ElementA = bfloat16_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = bfloat16_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + float const *C = reinterpret_cast(&c); + float *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3])); +#else + + MCTLASS_UNUSED(d); + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_NOT_IMPLEMENTED(); + +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F32 = F16 * F16 + F32 +template <> +struct Mma< + gemm::GemmShape<16, 8, 16>, + 32, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + float, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16, 8, 16>; + + using ElementA = half_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) + + // uint32_t const *A = reinterpret_cast(&a); + // uint32_t const *B = reinterpret_cast(&b); + // float const *C = reinterpret_cast(&c); + // float *D = reinterpret_cast(&d); + // asm volatile( + // "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + // "{%10,%11,%12,%13};\n" + // : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) + // : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + // "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3])); + + //We haven't a implicit type conversion which from mctlass::half_t to __fp16(maca), + //so we convert mctlass:half_t to float firstly. + auto results = __builtin_mxc_mma_16x16x16f16( + {a[0].get().to_macahalf(), a[1].get().to_macahalf(), + a[2].get().to_macahalf(), a[3].get().to_macahalf()}, + {b[0].get().to_macahalf(), b[1].get().to_macahalf(), + b[2].get().to_macahalf(), b[3].get().to_macahalf()}, + {c[0], c[1], c[2], c[3]}); + + d[0] = results[0]; + d[1] = results[1]; + d[2] = results[2]; + d[3] = results[3]; + +#else + + MCTLASS_UNUSED(d); + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_NOT_IMPLEMENTED(); + +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 884 - F64 +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F64 = F64 * F64 + F64 +template <> +struct Mma< + gemm::GemmShape<8,8,4>, + 32, + double, + layout::RowMajor, + double, + layout::ColumnMajor, + double, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<8,8,4>; + + using ElementA = double; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = double; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = double; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + + using ArchTag = arch::Sm80; + + MCTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c) const { + +#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) + + double const & A = reinterpret_cast(a); + double const & B = reinterpret_cast(b); + + double const *C = reinterpret_cast(&c); + double *D = reinterpret_cast(&d); + +__syncthreads(); //We need remove this someday + auto temp = __builtin_mxc_mma_16x16x4f64(A, B, + {C[0], C[1], D[0], D[1]}); +__syncthreads();//We need remove this someday + D[0] = temp[0]; + D[1] = temp[1]; + +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 16816 - S8 input, S32 accumulation +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S8 * S8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,16>, + 32, + int8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16,8,16>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + uint32_t const *A = reinterpret_cast(&a); + uint32_t const &B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5}, {%6}, " + "{%7,%8,%9,%10};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(B), "r"(C[0]), "r"(C[1]), "r"(C[2]), + "r"(C[3])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * S8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,16>, + 32, + uint8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16,8,16>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + uint32_t const *A = reinterpret_cast(&a); + uint32_t const &B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.u8.s8.s32 {%0,%1,%2,%3}, {%4,%5}, {%6}, " + "{%7,%8,%9,%10};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(B), "r"(C[0]), "r"(C[1]), "r"(C[2]), + "r"(C[3])); +#else + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S8 * U8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,16>, + 32, + int8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16,8,16>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const &B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.s8.u8.s32 {%0,%1,%2,%3}, {%4,%5}, {%6}, " + "{%7,%8,%9,%10};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(B), "r"(C[0]), "r"(C[1]), "r"(C[2]), + "r"(C[3])); +#else + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * U8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,16>, + 32, + uint8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16,8,16>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const &B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.u8.u8.s32 {%0,%1,%2,%3}, {%4,%5}, {%6}, " + "{%7,%8,%9,%10};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(B), "r"(C[0]), "r"(C[1]), "r"(C[2]), + "r"(C[3])); + +#else + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 16816 - S8 input, S32 accumulation - SATURATE +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S8 * S8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,16>, + 32, + int8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<16,8,16>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const &B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + //asm volatile( + // "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite {%0,%1,%2,%3}, {%4,%5}, " + // "{%6}, {%7,%8,%9,%10};\n" + // : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + // : "r"(A[0]), "r"(A[1]), "r"(B), "r"(C[0]), "r"(C[1]), "r"(C[2]), + // "r"(C[3])); + auto temp = __builtin_mxc_mma_16x16x16i8(A[0], B, {C[0], C[1], C[2], C[3]}); + D[0] = temp[0]; + D[1] = temp[1]; + D[2] = temp[2]; + D[3] = temp[3]; +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * S8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,16>, + 32, + uint8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<16,8,16>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const &B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.u8.s8.s32.satfinite {%0,%1,%2,%3}, {%4,%5}, " + "{%6}, {%7,%8,%9,%10};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(B), "r"(C[0]), "r"(C[1]), "r"(C[2]), + "r"(C[3])); +#else + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S8 * U8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,16>, + 32, + int8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<16,8,16>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const &B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.s8.u8.s32.satfinite {%0,%1,%2,%3}, {%4,%5}, " + "{%6}, {%7,%8,%9,%10};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(B), "r"(C[0]), "r"(C[1]), "r"(C[2]), + "r"(C[3])); +#else + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * U8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,16>, + 32, + uint8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<16,8,16>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const &B = reinterpret_cast(b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.s32.u8.u8.s32.satfinite {%0,%1,%2,%3}, {%4,%5}, " + "{%6}, {%7,%8,%9,%10};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(B), "r"(C[0]), "r"(C[1]), "r"(C[2]), + "r"(C[3])); +#else + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 16832 - S8 input, S32 accumulation +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S8 * S8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,32>, + 32, + int8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16,8,32>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * S8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,32>, + 32, + uint8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16,8,32>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.u8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S8 * U8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,32>, + 32, + int8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16,8,32>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.u8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * U8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,32>, + 32, + uint8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16,8,32>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.u8.u8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 16832 - S8 input, S32 accumulation - SATURATE +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S8 * S8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,32>, + 32, + int8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<16,8,32>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) + + uint32_t const * A = reinterpret_cast(&a); + uint32_t const * B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + // asm volatile( + // "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite {%0,%1,%2,%3}, " + // "{%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + // : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + // : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + // "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); + auto temp = __builtin_mxc_mma_16x16x16i8(A[0], B[0], {C[0], C[1], C[2], C[3]}); + D[0] = temp[0]; + D[1] = temp[1]; + D[2] = temp[2]; + D[3] = temp[3]; + + temp= __builtin_mxc_mma_16x16x16i8(A[1], B[1], {D[0], D[1], D[2], D[3]}); + D[0] = temp[0]; + D[1] = temp[1]; + D[2] = temp[2]; + D[3] = temp[3]; + D[4] = temp[4]; + D[5] = temp[5]; + D[6] = temp[6]; + D[7] = temp[7]; +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * S8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,32>, + 32, + uint8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<16,8,32>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.u8.s8.s32.satfinite {%0,%1,%2,%3}, " + "{%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S8 * U8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,32>, + 32, + int8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<16,8,32>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.u8.s32.satfinite {%0,%1,%2,%3}, " + "{%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * U8 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,32>, + 32, + uint8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<16,8,32>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.u8.u8.s32.satfinite {%0,%1,%2,%3}, " + "{%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 16864 - S4 input, S32 accumulation +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S4 * S4 + S32 +template <> +struct Mma< + gemm::GemmShape<16, 8, 64>, + 32, + mctlass::int4b_t, + layout::RowMajor, + mctlass::int4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16, 8, 64>; + + using ElementA = mctlass::int4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::int4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U4 * S4 + S32 +template <> +struct Mma< + gemm::GemmShape<16, 8, 64>, + 32, + mctlass::uint4b_t, + layout::RowMajor, + mctlass::int4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16, 8, 64>; + + using ElementA = mctlass::uint4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::int4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.u4.s4.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S4 * U4 + S32 +template <> +struct Mma< + gemm::GemmShape<16, 8, 64>, + 32, + mctlass::int4b_t, + layout::RowMajor, + mctlass::uint4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16, 8, 64>; + + using ElementA = mctlass::int4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::uint4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.u4.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U4 * U4 + S32 +template <> +struct Mma< + gemm::GemmShape<16, 8, 64>, + 32, + mctlass::uint4b_t, + layout::RowMajor, + mctlass::uint4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16, 8, 64>; + + using ElementA = mctlass::uint4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::uint4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.u4.u4.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 16864 - S4 input, S32 accumulation - SATURATE +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S4 * S4 + S32 +template <> +struct Mma< + gemm::GemmShape<16, 8, 64>, + 32, + mctlass::int4b_t, + layout::RowMajor, + mctlass::int4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<16, 8, 64>; + + using ElementA = mctlass::int4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::int4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const * A = reinterpret_cast(&a); + uint32_t const * B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32.satfinite {%0,%1,%2,%3}, " + "{%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U4 * S4 + S32 +template <> +struct Mma< + gemm::GemmShape<16, 8, 64>, + 32, + mctlass::uint4b_t, + layout::RowMajor, + mctlass::int4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<16, 8, 64>; + + using ElementA = mctlass::uint4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::int4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.u4.s4.s32.satfinite {%0,%1,%2,%3}, " + "{%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S4 * U4 + S32 +template <> +struct Mma< + gemm::GemmShape<16, 8, 64>, + 32, + mctlass::int4b_t, + layout::RowMajor, + mctlass::uint4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<16, 8, 64>; + + using ElementA = mctlass::int4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::uint4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.u4.s32.satfinite {%0,%1,%2,%3}, " + "{%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U4 * U4 + S32 +template <> +struct Mma< + gemm::GemmShape<16, 8, 64>, + 32, + mctlass::uint4b_t, + layout::RowMajor, + mctlass::uint4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate> { + + using Shape = gemm::GemmShape<16, 8, 64>; + + using ElementA = mctlass::uint4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::uint4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAddSaturate; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.u4.u4.s32.satfinite {%0,%1,%2,%3}, " + "{%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = B1 & B1 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,256>, + 32, + mctlass::uint1b_t, + layout::RowMajor, + mctlass::uint1b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16,8,256>; + + using ElementA = mctlass::uint1b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::uint1b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int32_t; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k256.row.col.s32.b1.b1.s32.and.popc {%0,%1,%2,%3}, " + "{%4,%5,%6,%7}, " + "{%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 168256 - B1 input, S32 accumulation - XOR,POPC +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = B1 & B1 + S32 +template <> +struct Mma< + gemm::GemmShape<16,8,256>, + 32, + mctlass::uint1b_t, + layout::RowMajor, + mctlass::uint1b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpXorPopc> { + + using Shape = gemm::GemmShape<16,8,256>; + + using ElementA = mctlass::uint1b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::uint1b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpXorPopc; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +//#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) +#if 0 //mctlass not support this data-type and shape now + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + asm volatile( + "mma.sync.aligned.m16n8k256.row.col.s32.b1.b1.s32.xor.popc {%0,%1,%2,%3}, " + "{%4,%5,%6,%7}, " + "{%8,%9}, {%10,%11,%12,%13};\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3])); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); + +#endif // defined(MCTLASS_ARCH_MMA_SM80_ENABLED) + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Matrix Multiply 161616 +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F16 = F16 * F16 + F16 +template <> +struct Mma< + gemm::GemmShape<16, 16, 16>, + 32, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16, 16, 16>; + + using ElementA = half_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = half_t; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c) const { + +#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) + + auto results = __builtin_mxc_mma_16x16x16f16( + {a[0].get().to_macahalf(), a[1].get().to_macahalf(), + a[2].get().to_macahalf(), a[3].get().to_macahalf()}, + {b[0].get().to_macahalf(), b[1].get().to_macahalf(), + b[2].get().to_macahalf(), b[3].get().to_macahalf()}, + {c[0].get().to_macahalf(), c[1].get().to_macahalf(), + c[2].get().to_macahalf(), c[3].get().to_macahalf()}); + + d[0] = mctlass::half_t(results[0]); + d[1] = mctlass::half_t(results[1]); + d[2] = mctlass::half_t(results[2]); + d[3] = mctlass::half_t(results[3]); +#else + + MCTLASS_UNUSED(d); + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_NOT_IMPLEMENTED(); + +#endif + } +}; + +/// Matrix multiply-add operation: F32 = F16 * F16 + F32 +template <> +struct Mma< + gemm::GemmShape<16, 16, 16>, + 32, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + float, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16, 16, 16>; + + using ElementA = half_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c + ) const { + +#if defined(MCTLASS_ARCH_MMA_SM80_ENABLED) + + auto results = __builtin_mxc_mma_16x16x16f16( + {a[0].get().to_macahalf(), a[1].get().to_macahalf(), + a[2].get().to_macahalf(), a[3].get().to_macahalf()}, + {b[0].get().to_macahalf(), b[1].get().to_macahalf(), + b[2].get().to_macahalf(), b[3].get().to_macahalf()}, + {c[0], c[1], c[2], c[3]}); + + d[0] = results[0]; + d[1] = results[1]; + d[2] = results[2]; + d[3] = results[3]; + +#else + + MCTLASS_UNUSED(d); + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_NOT_IMPLEMENTED(); + +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/arch/mma_sm90.h b/csrc/mctlass/include/mctlass/arch/mma_sm90.h new file mode 100644 index 0000000..b6164fb --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/mma_sm90.h @@ -0,0 +1,266 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Matrix multiply +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mma.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" + +//////////////////////////////////////////////////////////////////////////////// + +// #if ((__CUDACC_VER_MAJOR__ > 11) || (__CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ >= 8)) +#if 0 + #define MCTLASS_ARCH_MMA_SM90_F64_MMA_SUPPORTED + #if (!defined(MCTLASS_ARCH_MMA_SM90_F64_MMA_ENABLED)) + #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + #define MCTLASS_ARCH_MMA_SM90_F64_MMA_ENABLED + #endif + #endif +#endif + +// #if (__CUDACC_VER_MAJOR__ >= 12) +#if 0 + #define MCTLASS_ARCH_MMA_SM90_SUPPORTED + #if (!defined(MCTLASS_ARCH_MMA_SM90_ENABLED)) + #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + #define MCTLASS_ARCH_MMA_SM90_ENABLED + #endif + #endif +#endif + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace arch { + +//////////////////////////////////////////////////////////////////////////////// +/// Matrix Multiply-Add 16x8x4 fp64 +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F64 = F64 * F64 + F64 +template <> +struct Mma< + gemm::GemmShape<16,8,4>, + 32, + double, + layout::RowMajor, + double, + layout::ColumnMajor, + double, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16,8,4>; + + using ElementA = double; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = double; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = double; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + + using ArchTag = arch::Sm90; + + MCTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c) const { + +#if defined(MCTLASS_ARCH_MMA_SM90_F64_MMA_ENABLED) + + double const *A = reinterpret_cast(&a); + double const *B = reinterpret_cast(&b); + + double const *C = reinterpret_cast(&c); + double *D = reinterpret_cast(&d); + + asm volatile("mma.sync.aligned.m16n8k4.row.col.f64.f64.f64.f64.rn {%0, %1, %2, %3}, {%4, %5}, {%6}, {%7, %8, %9, %10};\n" + : "=d"(D[0]), "=d"(D[1]), "=d"(D[2]), "=d"(D[3]) + : "d"(A[0]), "d"(A[1]), + "d"(B[0]), + "d"(C[0]), "d"(C[1]), "d"(C[2]), "d"(C[3])); + +#else + MCTLASS_UNUSED(d); + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_NOT_IMPLEMENTED(); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +/// Matrix Multiply-Add 16x8x8 fp64 +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F64 = F64 * F64 + F64 +template <> +struct Mma< + gemm::GemmShape<16,8,8>, + 32, + double, + layout::RowMajor, + double, + layout::ColumnMajor, + double, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16,8,8>; + + using ElementA = double; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = double; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = double; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + + using ArchTag = arch::Sm90; + + MCTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c) const { + +#if defined(MCTLASS_ARCH_MMA_SM90_F64_MMA_ENABLED) + + double const *A = reinterpret_cast(&a); + double const *B = reinterpret_cast(&b); + + double const *C = reinterpret_cast(&c); + double *D = reinterpret_cast(&d); + + asm volatile("mma.sync.aligned.m16n8k8.row.col.f64.f64.f64.f64 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};\n" + : "=d"(D[0]), "=d"(d[1]), "=d"(d[2]), "=d"(d[3]) + : "d"(A[0]), "d"(A[1]), "d"(A[2]), "d"(A[3]), + "d"(B[0]), "d"(B[1]), + "d"(C[0]), "d"(C[1]), "d"(C[2]), "d"(C[3])); + +#else + + MCTLASS_UNUSED(d); + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_NOT_IMPLEMENTED(); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +/// Matrix Multiply-Add 16x8x16 fp64 +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F64 = F64 * F64 + F64 +template <> +struct Mma< + gemm::GemmShape<16,8,16>, + 32, + double, + layout::RowMajor, + double, + layout::ColumnMajor, + double, + layout::RowMajor, + OpMultiplyAdd> { + + using Shape = gemm::GemmShape<16,8,16>; + + using ElementA = double; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = double; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = double; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using Operator = OpMultiplyAdd; + + using ArchTag = arch::Sm90; + + MCTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c) const { + +#if defined(MCTLASS_ARCH_MMA_SM90_F64_MMA_ENABLED) + + double const *A = reinterpret_cast(&a); + double const *B = reinterpret_cast(&b); + + double const *C = reinterpret_cast(&c); + double *D = reinterpret_cast(&d); + + asm volatile("mma.sync.aligned.m16n8k16.row.col.f64.f64.f64.f64 {%0, %1, %2, %3}, {%4, %5, %6, %7, %8, %9, %10, %11}, {%12, %13, %14, %15}, {%16, %17, %18, %19};\n" + : "=d"(D[0]), "=d"(D[1]), "=d"(D[2]), "=d"(D[3]) + : "d"(A[0]), "d"(A[2]), "d"(A[2]), "d"(A[3]), "d"(A[4]), "d"(A[5]), "d"(A[6]), "d"(A[7]), + "d"(B[0]), "d"(B[1]), "d"(B[2]), "d"(B[3]), + "d"(C[0]), "d"(C[1]), "d"(C[2]), "d"(C[3])); + +#else + MCTLASS_NOT_IMPLEMENTED(); +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/arch/mma_sparse_sm80.h b/csrc/mctlass/include/mctlass/arch/mma_sparse_sm80.h new file mode 100644 index 0000000..4b39b1b --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/mma_sparse_sm80.h @@ -0,0 +1,1646 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Sparse matrix multiply accumulate for SM80 +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mma.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +//#if ((__CUDACC_VER_MAJOR__ > 11) || (__CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ >= 1)) +#if 0 //mctlass not support sparse mma +#define MCTLASS_ARCH_SPARSE_MMA_SM80_SUPPORTED 1 + +#if defined(__MACA_ARCH__) +#define MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED +#endif +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace arch { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// +// +// Sparse Matrix Multiply 16832 +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F16 = F16 * F16 + F16 +template <> +struct SparseMma< + gemm::GemmShape<16, 8, 32>, + 32, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + OpMultiplyAdd, + SPFormatType::Thread +> { + + using Shape = gemm::GemmShape<16, 8, 32>; + + using ElementA = half_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = half_t; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 2; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c, uint32_t const &E, int const id2) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + uint32_t const *C = reinterpret_cast(&c); + uint32_t *D = reinterpret_cast(&d); + if (id2 == 0) { + asm volatile( + "mma.sp.sync.aligned.m16n8k32.row.col.f16.f16.f16.f16 {%0,%1}, " + "{%2,%3,%4,%5}, {%6,%7,%8,%9}, {%10,%11}, %12, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(B[2]), "r"(B[3]), "r"(C[0]), "r"(C[1]), "r"(E)); + } + else if (id2 == 1) { + asm volatile( + "mma.sp.sync.aligned.m16n8k32.row.col.f16.f16.f16.f16 {%0,%1}, " + "{%2,%3,%4,%5}, {%6,%7,%8,%9}, {%10,%11}, %12, 0x1;\n" + : "=r"(D[0]), "=r"(D[1]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(B[2]), "r"(B[3]), "r"(C[0]), "r"(C[1]), "r"(E)); + } + else { + assert(0); + } +#else + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F32 = F16 * F16 + F32 +template <> +struct SparseMma< + gemm::GemmShape<16, 8, 32>, + 32, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + float, + layout::RowMajor, + OpMultiplyAdd, + SPFormatType::Thread + > { + + using Shape = gemm::GemmShape<16, 8, 32>; + + using ElementA = half_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = half_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 2; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c, uint32_t const &E, int const id2) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + float const *C = reinterpret_cast(&c); + float *D = reinterpret_cast(&d); + if (id2 == 0) { + asm volatile( + "mma.sp.sync.aligned.m16n8k32.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, " + "{%4,%5,%6,%7}, {%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(B[2]), "r"(B[3]), "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3]), + "r"(E)); + } + else if (id2 == 1) { + asm volatile( + "mma.sp.sync.aligned.m16n8k32.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, " + "{%4,%5,%6,%7}, {%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x1;\n" + : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "r"(B[2]), "r"(B[3]), "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3]), + "r"(E)); + } + else { + assert(0); + } +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Sparse Matrix Multiply 16832 - Float BF16, FP32 accumulation +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F32 = bf16 * bf16 + F32 +template <> +struct SparseMma, 32, bfloat16_t, layout::RowMajor, + bfloat16_t, layout::ColumnMajor, float, layout::RowMajor, + OpMultiplyAdd, SPFormatType::Thread> { + using Shape = gemm::GemmShape<16, 8, 32>; + + using ElementA = bfloat16_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = bfloat16_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 2; + + MCTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c, uint32_t const &E, int const id2) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + float const *C = reinterpret_cast(&c); + float *D = reinterpret_cast(&d); + if (id2 == 0) { + asm volatile( + "mma.sp.sync.aligned.m16n8k32.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3]), "r"(E)); + } else if (id2 == 1) { + asm volatile( + "mma.sp.sync.aligned.m16n8k32.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x1;\n" + : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3]), "r"(E)); + } else { + assert(0); + } +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Sparse Matrix Multiply 16816 - Float TF32 +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: F32 = tf32 * tf32 + F32 +template <> +struct SparseMma, 32, tfloat32_t, layout::RowMajor, + tfloat32_t, layout::ColumnMajor, float, layout::RowMajor, + OpMultiplyAdd, SPFormatType::Thread> { + using Shape = gemm::GemmShape<16, 8, 16>; + + using ElementA = tfloat32_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = tfloat32_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = float; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 4; + + static int const kMaxID2 = 2; + + MCTLASS_HOST_DEVICE + void operator()(FragmentC &d, FragmentA const &a, FragmentB const &b, + FragmentC const &c, uint32_t const &E, int const id2) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + float const *C = reinterpret_cast(&c); + float *D = reinterpret_cast(&d); + if (id2 == 0) { + asm volatile( + "mma.sp.sync.aligned.m16n8k16.row.col.f32.tf32.tf32.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3]), "r"(E)); + } else if (id2 == 1) { + asm volatile( + "mma.sp.sync.aligned.m16n8k16.row.col.f32.tf32.tf32.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x1;\n" + : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "f"(C[0]), "f"(C[1]), "f"(C[2]), "f"(C[3]), "r"(E)); + } else { + assert(0); + } +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Sparse Matrix Multiply 16864 - S8 input, S32 accumulation +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S8 * S8 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,64>, + 32, + int8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,64>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k64.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S8 * U8 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,64>, + 32, + int8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,64>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k64.row.col.s32.s8.u8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * S8 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,64>, + 32, + uint8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,64>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k64.row.col.s32.u8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * U8 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,64>, + 32, + uint8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,64>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k64.row.col.s32.u8.u8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Sparse Matrix Multiply 16864 - S8 input, S32 accumulation - SATURATE +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S8 * S8 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,64>, + 32, + int8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,64>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k64.row.col.s32.s8.s8.s32.satfinite {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S8 * U8 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,64>, + 32, + int8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,64>; + + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k64.row.col.s32.s8.u8.s32.satfinite {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * S8 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,64>, + 32, + uint8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,64>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k64.row.col.s32.u8.s8.s32.satfinite {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U8 * U8 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,64>, + 32, + uint8_t, + layout::RowMajor, + uint8_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,64>; + + using ElementA = uint8_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = uint8_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k64.row.col.s32.u8.u8.s32.satfinite {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Sparse Matrix Multiply 168128 - S4 input, S32 accumulation +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S4 * S4 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,128>, + 32, + mctlass::int4b_t, + layout::RowMajor, + mctlass::int4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,128>; + + using ElementA = mctlass::int4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::int4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k128.row.col.s32.s4.s4.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S4 * U4 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,128>, + 32, + mctlass::int4b_t, + layout::RowMajor, + mctlass::uint4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,128>; + + using ElementA = mctlass::int4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::uint4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k128.row.col.s32.s4.u4.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U4 * S4 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,128>, + 32, + mctlass::uint4b_t, + layout::RowMajor, + mctlass::int4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,128>; + + using ElementA = mctlass::uint4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::int4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k128.row.col.s32.u4.s4.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U4 * U4 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,128>, + 32, + mctlass::uint4b_t, + layout::RowMajor, + mctlass::uint4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAdd, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,128>; + + using ElementA = mctlass::uint4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::uint4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k128.row.col.s32.u4.u4.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// Sparse Matrix Multiply 168128 - S4 input, S32 accumulation - SATURATE +// +//////////////////////////////////////////////////////////////////////////////// + +/// Matrix multiply-add operation: S32 = S4 * S4 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,128>, + 32, + mctlass::int4b_t, + layout::RowMajor, + mctlass::int4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,128>; + + using ElementA = mctlass::int4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::int4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k128.row.col.s32.s4.s4.s32.satfinite {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = S4 * U4 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,128>, + 32, + mctlass::int4b_t, + layout::RowMajor, + mctlass::uint4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,128>; + + using ElementA = mctlass::int4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::uint4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k128.row.col.s32.s4.u4.s32.satfinite {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U4 * S4 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,128>, + 32, + mctlass::uint4b_t, + layout::RowMajor, + mctlass::int4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,128>; + + using ElementA = mctlass::uint4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::int4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k128.row.col.s32.u4.s4.s32.satfinite {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +/// Matrix multiply-add operation: S32 = U4 * U4 + S32 +template <> +struct SparseMma< + gemm::GemmShape<16,8,128>, + 32, + mctlass::uint4b_t, + layout::RowMajor, + mctlass::uint4b_t, + layout::ColumnMajor, + int, + layout::RowMajor, + OpMultiplyAddSaturate, + SPFormatType::Thread> { + + using Shape = gemm::GemmShape<16,8,128>; + + using ElementA = mctlass::uint4b_t; + using LayoutA = layout::RowMajor; + using FragmentA = Array; + + using ElementB = mctlass::uint4b_t; + using LayoutB = layout::ColumnMajor; + using FragmentB = Array; + + using ElementC = int; + using LayoutC = layout::RowMajor; + using FragmentC = Array; + + using FragmentE = uint32_t; + + using Operator = OpMultiplyAdd; + using ArchTag = arch::Sm80; + + static int const kSparse = 2; + + static int const kMetaSizeInBits = 2; + + static int const kMaxID2 = 1; + + /// Computes multiply-add + MCTLASS_HOST_DEVICE + void operator()( + FragmentC &d, + FragmentA const &a, + FragmentB const &b, + FragmentC const &c, + uint32_t const &E, + int const id2 + ) const { + +#if defined(MCTLASS_ARCH_SPARSE_MMA_SM80_ENABLED) + + uint32_t const *A = reinterpret_cast(&a); + uint32_t const *B = reinterpret_cast(&b); + + int const *C = reinterpret_cast(&c); + int *D = reinterpret_cast(&d); + if (id2 == 0) + asm volatile( + "mma.sp.sync.aligned.m16n8k128.row.col.s32.u4.u4.s32.satfinite {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9,%10,%11}, {%12,%13,%14,%15}, %16, 0x0;\n" + : "=r"(D[0]), "=r"(D[1]), "=r"(D[2]), "=r"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(B[2]), "r"(B[3]), + "r"(C[0]), "r"(C[1]), "r"(C[2]), "r"(C[3]), "r"(E)); + else + assert(0); +#else + + MCTLASS_UNUSED(a); + MCTLASS_UNUSED(b); + MCTLASS_UNUSED(c); + MCTLASS_UNUSED(d); + assert(0); +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/arch/reg_reconfig.h b/csrc/mctlass/include/mctlass/arch/reg_reconfig.h new file mode 100644 index 0000000..72f3795 --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/reg_reconfig.h @@ -0,0 +1,68 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief PTX for CTA Reconfiguration +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDACC_VER_MAJOR__ >= 12)) + #if (defined(__CUDA_ARCH_FEAT_SM90_ALL)) + #define CUDA_CTA_RECONFIG_ACTIVATED 1 + #endif +#else + #define CUDA_CTA_RECONFIG_ACTIVATED 0 +#endif + +namespace mctlass { +namespace arch { + +template +MCTLASS_DEVICE +void warpgroup_reg_alloc(){ +#if CUDA_CTA_RECONFIG_ACTIVATED + asm volatile( "setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(RegCount) ); +#endif +} + +template +MCTLASS_DEVICE +void warpgroup_reg_dealloc(){ +#if CUDA_CTA_RECONFIG_ACTIVATED + asm volatile( "setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(RegCount) ); +#endif +} + +} // namespace arch +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/arch/simd.h b/csrc/mctlass/include/mctlass/arch/simd.h new file mode 100644 index 0000000..8628df1 --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/simd.h @@ -0,0 +1,125 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates exposing SIMD operators +*/ + +#pragma once + +#include "../array.h" +#include "../numeric_types.h" + +namespace mctlass { +namespace arch { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Element-wise operators +// + +MCTLASS_HOST_DEVICE +template +Array operator*(Array const &a, Array const &b) { + Array d; + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + d[i] = a[i] * b[i]; + } + return d; +} + +MCTLASS_HOST_DEVICE +template +Array operator+(Array const &a, Array const &b) { + Array d; + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + d[i] = a[i] + b[i]; + } + return d; +} + +MCTLASS_HOST_DEVICE +template +Array operator-(Array const &a, Array const &b) { + Array d; + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + d[i] = a[i] - b[i]; + } + return d; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Multiply-accumulate operators +// + +MCTLASS_HOST_DEVICE +template +Array mac(Array const &a, Array const &b, Array const &c) { + Array d; + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + d[i] = a[i] * b[i] + c[i]; + } + return d; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Dot product operator +// + +MCTLASS_HOST_DEVICE +template +Accumulator dot(Array const &a, Array const &b, Accumulator accum) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + accum += a[i] * b[i]; + } + return accum; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "simd_sm60.h" +#include "simd_sm61.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/arch/simd_sm60.h b/csrc/mctlass/include/mctlass/arch/simd_sm60.h new file mode 100644 index 0000000..9a1d499 --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/simd_sm60.h @@ -0,0 +1,110 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates exposing SIMD operators for SM60 +*/ + +#pragma once + +#include "simd.h" + +namespace mctlass { +namespace arch { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Element-wise operators - specialized for half_t x 2 +// + +MCTLASS_HOST_DEVICE +template <> +Array operator*(Array const &a, Array const &b) { + Array d; + + + return d; +} + +MCTLASS_HOST_DEVICE +template <> +Array operator+(AArray const &a, Array const &b) { + Array d; + + + return d; +} + +MCTLASS_HOST_DEVICE +template <> +Array operator-(Array const &a, Array const &b) { + Array d; + + + return d; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Multiply-accumulate operators - specialized for half_t x 2 +MCTLASS_HOST_DEVICE +template <> +Array mac(Array const &a, Array const &b, Array const &c) { + Array d; + + + return d; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Dot product operator - specialized for half_t <- (half_t * half_t) x 2 + half_t +MCTLASS_HOST_DEVICE +template <> +half_t dot(Array const &a, Array const &b, half_t accum) { + + + return accum; +} + +/// Dot product operator - specialized for float <- (half_t * half_t) x 2 + float +MCTLASS_HOST_DEVICE +template <> +float dot(Array const &a, Array const &b, float accum) { + + + return accum; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/arch/simd_sm61.h b/csrc/mctlass/include/mctlass/arch/simd_sm61.h new file mode 100644 index 0000000..e1a4b05 --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/simd_sm61.h @@ -0,0 +1,147 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates exposing SIMD operators for SM61 +*/ + +#pragma once + +#include "simd.h" + +namespace mctlass { +namespace arch { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Dot product operator - specialized for int32_t <- (int8_t * int8_t) x 4 + int32_t +MCTLASS_HOST_DEVICE +template <> +int32_t dot(Array const &a, Array const &b, int32_t accum) { + + return accum; +} + +/// Dot product operator - specialized for int32_t <- (uint8_t * int8_t) x 4 + int32_t +MCTLASS_HOST_DEVICE +template <> +int32_t dot(Array const &a, Array const &b, int32_t accum) { + + return accum; +} + +/// Dot product operator - specialized for int32_t <- (int8_t * uint8_t) x 4 + int32_t +MCTLASS_HOST_DEVICE +template <> +int32_t dot(Array const &a, Array const &b, int32_t accum) { + + return accum; +} + +/// Dot product operator - specialized for int32_t <- (uint8_t * uint8_t) x 4 + int32_t +MCTLASS_HOST_DEVICE +template <> +int32_t dot(Array const &a, Array const &b, int32_t accum) { + + return accum; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Dot product operator - specialized for int32_t <- (int16_t * int8_t) x 2 + int32_t +MCTLASS_HOST_DEVICE +template <> +int32_t dot(Array const &a, Array const &b, int32_t accum) { + + return accum; +} + +/// Dot product operator - specialized for int32_t <- (uint16_t * int8_t) x 2 + int32_t +MCTLASS_HOST_DEVICE +template <> +int32_t dot(Array const &a, Array const &b, int32_t accum) { + + return accum; +} + +/// Dot product operator - specialized for int32_t <- (int16_t * int8_t) x 2 + int32_t +MCTLASS_HOST_DEVICE +template <> +int32_t dot(Array const &a, Array const &b, int32_t accum) { + + return accum; +} + +/// Dot product operator - specialized for int32_t <- (uint16_t * int8_t) x 2 + int32_t +MCTLASS_HOST_DEVICE +template <> +int32_t dot(Array const &a, Array const &b, int32_t accum) { + + return accum; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Dot product operator - specialized for int32_t <- (int16_t * int16_t) x 2 + int32_t +MCTLASS_HOST_DEVICE +template <> +int32_t dot(Array const &a, Array const &b, int32_t accum) { + + return accum; +} + +/// Dot product operator - specialized for int32_t <- (uint16_t * int16_t) x 2 + int32_t +MCTLASS_HOST_DEVICE +template <> +int32_t dot(Array const &a, Array const &b, int32_t accum) { + + return accum; +} + +/// Dot product operator - specialized for int32_t <- (int16_t * int16_t) x 2 + int32_t +MCTLASS_HOST_DEVICE +template <> +int32_t dot(Array const &a, Array const &b, int32_t accum) { + + return accum; +} + +/// Dot product operator - specialized for int32_t <- (uint16_t * int16_t) x 2 + int32_t +MCTLASS_HOST_DEVICE +template <> +int32_t dot(Array const &a, Array const &b, int32_t accum) { + + return accum; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/arch/wmma.h b/csrc/mctlass/include/mctlass/arch/wmma.h new file mode 100644 index 0000000..2bf33ef --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/wmma.h @@ -0,0 +1,230 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates exposing architecture support for warp matrix multiply-add (WMMA) operations +*/ + +#pragma once + +// MCTLASS WMMA does not support clang at present. +//#if !(defined(__clang__) && defined(__MACA__)) +#if 1 + +//#if (__CUDACC_VER_MAJOR__ >= 9) +#if 1 +#if defined(__MACA_ARCH__) + #define MCTLASS_ARCH_WMMA_ENABLED + #define MCTLASS_ARCH_WMMA_SM70_ENABLED + #include "__clang_maca_mma_functions.h" +#endif +#endif + +//#if (__CUDACC_VER_MAJOR__ >= 10) +#if 0 +#if defined(__MACA_ARCH__) +#define MCTLASS_ARCH_INTEGER_MATRIX_MULTIPLY_ENABLED +#define MCTLASS_ARCH_WMMA_SM72_ENABLED +#endif +#endif + +//#if (__CUDACC_VER_MAJOR__ >= 10) +#if 0 +#if defined(__MACA_ARCH__) +#define MCTLASS_SUBBYTE_INTEGER_MATRIX_MULTIPLY_ENABLED +#define MCTLASS_ARCH_WMMA_SM75_ENABLED +#endif +#endif + +#endif //!(defined(__clang__) && defined(__MACA__)) + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) + +// #include +#include "mctlass/arch/mma.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/gemm/gemm.h" + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace arch { + +//////////////////////////////////////////////////////////////////////////////////////////////// +/// Statically maps mctlass data types => mxmaca::wmma data types +///////////////////////////////////////////////////////////////////////////////////////////////// +template +struct MctlassToWmmaDataType{ + using Type = Type_; +}; + +/// Statically maps mctlass::half_t => __half +template<> +struct MctlassToWmmaDataType { + using Type = __half; +}; + +#if defined(__MACA_ARCH__) && (__CUDACC_VER_MAJOR__ >= 11) +template<> +struct MctlassToWmmaDataType { + using Type = maca_bfloat16; +}; +#endif + +/// Statically maps int8_t => char +template<> +struct MctlassToWmmaDataType { + using Type = signed char; +}; + +/// Statically maps uint8_t => char +template<> +struct MctlassToWmmaDataType { + using Type = unsigned char; +}; + +/// Statically maps int32_t => int +template<> +struct MctlassToWmmaDataType { + using Type = int; +}; + +//#if defined(MCTLASS_SUBBYTE_INTEGER_MATRIX_MULTIPLY_ENABLED) +#if 0 +/// Statically maps mctlass::int4b_t => experimental::precision::s4 +template<> +struct MctlassToWmmaDataType { + using Type = mxmaca::wmma::experimental::precision::s4; +}; + +/// Statically maps mctlass::uint4b_t => experimental::precision::s4 +template<> +struct MctlassToWmmaDataType { + using Type = mxmaca::wmma::experimental::precision::u4; +}; + +/// Statically maps mctlass::uint1b_t => experimental::precision::b1 +template<> +struct MctlassToWmmaDataType { + using Type = mxmaca::wmma::experimental::precision::b1; +}; +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////// +/// Statically maps mctlass::layout => mxmaca::wmma layout tags +//////////////////////////////////////////////////////////////////////////////////////////////// +template +struct MctlassToWmmaLayout { +}; + +/// Statically maps mctlass::layout::RowMajor => wmma::row_major layout tags +template <> +struct MctlassToWmmaLayout { + using Layout = mxmaca::wmma::row_major; + static mxmaca::wmma::layout_t const value = mxmaca::wmma::layout_t::mem_row_major; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////// +/// Statically maps mctlass::layout::RowMajor => wmma::row_major layout tags +//////////////////////////////////////////////////////////////////////////////////////////////// +template <> +struct MctlassToWmmaLayout { + using Layout = mxmaca::wmma::col_major; + static mxmaca::wmma::layout_t const value = mxmaca::wmma::layout_t::mem_col_major; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////////////////////// +/// Statically maps mxmaca::wmma data types => mctlass data types +///////////////////////////////////////////////////////////////////////////////////////////////// +template +struct WmmaToMctlassDataType{ + using Type = Type_; +}; + +/// Statically maps __half => mctlass::half_t +template<> +struct WmmaToMctlassDataType<__half> { + using Type = mctlass::half_t; +}; + +#if defined(__MACA_ARCH__) && (__CUDACC_VER_MAJOR__ >= 11) +template<> +struct WmmaToMctlassDataType { + using Type = mctlass::bfloat16_t; +}; +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// +// WMMA template structure defines mxmaca::wmma::fragments and static assertion chaeks +// for a specific template paramterized data type (Element[A|B|C]), layout (Layout[A|B|C]), +// and native wmma size (Shape) +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + typename Shape_, ///< Size of the matrix product (concept: GemmShape) + typename ElementA_, ///< Data type of A elements + typename LayoutA_, ///< Layout of A matrix (concept: MatrixLayout) + typename ElementB_, ///< Data type of B elements + typename LayoutB_, ///< Layout of B matrix (concept: MatrixLayout) + typename ElementC_, ///< Element type of C matrix + typename LayoutC_, /// Layout of C matrix (concept: MatrixLayout) + typename Operator_ = mctlass::arch::OpMultiplyAdd ///< Inner product operator (multiply-add, xor.popc) +> +struct Wmma; +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace arch +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Specializations for each compute capability +// +#ifdef MCTLASS_ARCH_WMMA_SM70_ENABLED +#include "mctlass/arch/wmma_sm70.h" +#endif + +#ifdef MCTLASS_ARCH_WMMA_SM72_ENABLED +#include "mctlass/arch/wmma_sm72.h" +#endif + +#ifdef MCTLASS_ARCH_WMMA_SM75_ENABLED +#include "mctlass/arch/wmma_sm75.h" +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#endif //MCTLASS_ARCH_WMMA_ENABLED diff --git a/csrc/mctlass/include/mctlass/arch/wmma_sm70.h b/csrc/mctlass/include/mctlass/arch/wmma_sm70.h new file mode 100644 index 0000000..2a55c6e --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/wmma_sm70.h @@ -0,0 +1,137 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Matrix multiply +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif +#include "mctlass/layout/matrix.h" + +//////////////////////////////////////////////////////////////////////////////// +namespace mctlass { +namespace arch { + + +//////////////////////////////////////////////////////////////////////////////// +// +// WMMA template structure defines mxmaca::wmma::fragments and static assert for +// wmma native instruction sizes supported for half +// +//////////////////////////////////////////////////////////////////////////////// +template < +typename Shape_, +typename LayoutA_, +typename LayoutB_, +typename ElementC_, +typename LayoutC_> +struct Wmma< + Shape_, ///< Size of the matrix product (concept: GemmShape) + mctlass::half_t, ///< ElementA + LayoutA_, ///< LayoutA + mctlass::half_t, ///< ElementB + LayoutB_, ///< LayoutB + ElementC_, ///< ElementC + LayoutC_, ///< LayoutC + mctlass::arch::OpMultiplyAdd ///< Operator (multiply-add, xor.popc) +> { + +#if defined(MCTLASS_ARCH_WMMA_SM70_ENABLED) + using Shape = Shape_; + using ElementA = mctlass::half_t; + using LayoutA = LayoutA_; + using ElementB = mctlass::half_t; + using LayoutB = LayoutB_; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using Operator = mctlass::arch::OpMultiplyAdd; + using ArchTag = arch::Sm70; + + // check supported wmma shape for the given multiplicand data types + static_assert( + platform::is_same, Shape>::value || + platform::is_same, Shape>::value || + platform::is_same, Shape>::value || + platform::is_same, Shape>::value, + "Supported list of wmma operator shape for f16 multiplicands are: 16x16x16, 8x32x16, 32x8x16 and 16x8x16"); + + // check supported wmma output data type for the given multiplicand data types + static_assert( + platform::is_same::value || platform::is_same::value, + "Supported of wmma output data type for f16 multiplicands are: f16 and f32"); + + // Wmma Fragment + using FragmentA = mxmaca::wmma::fragment< + mxmaca::wmma::matrix_a, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type, + typename MctlassToWmmaLayout::Layout>; + + using FragmentB = mxmaca::wmma::fragment< + mxmaca::wmma::matrix_b, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type, + typename MctlassToWmmaLayout::Layout>; + + using FragmentC = mxmaca::wmma::fragment< + mxmaca::wmma::accumulator, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type>; + + /// Performs a mxmaca::wmma matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + FragmentA const &A, + FragmentB const &B, + FragmentC const &C) const { + + mxmaca::wmma::mma_sync(D, A, B, C); + } +#else + static_assert(false, "wmma.mma.sync for floating point multiplicands is avialable only for SM70 and beyond"); +#endif + +}; + +} // namespace arch +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/arch/wmma_sm72.h b/csrc/mctlass/include/mctlass/arch/wmma_sm72.h new file mode 100644 index 0000000..3402b4d --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/wmma_sm72.h @@ -0,0 +1,210 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Matrix multiply +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif +#include "mctlass/layout/matrix.h" + +//////////////////////////////////////////////////////////////////////////////// +namespace mctlass { +namespace arch { + +//////////////////////////////////////////////////////////////////////////////// +// +// WMMA template structure defines mxmaca::wmma::fragments and static assert for +// wmma native instruction sizes supported for int8_t +// +//////////////////////////////////////////////////////////////////////////////// +template < +typename Shape_, +typename LayoutA_, +typename LayoutB_, +typename LayoutC_> +struct Wmma< + Shape_, ///< Size of the matrix product (concept: GemmShape) + int8_t, ///< ElementA + LayoutA_, ///< LayoutA + int8_t, ///< ElementB + LayoutB_, ///< LayoutB + int32_t, ///< ElementC + LayoutC_, ///< LayoutC + mctlass::arch::OpMultiplyAdd ///< Operator (multiply-add, xor.popc) +> { +#if defined(MCTLASS_ARCH_WMMA_SM72_ENABLED) + using Shape = Shape_; + using ElementA = int8_t; + using LayoutA = LayoutA_; + using ElementB = int8_t; + using LayoutB = LayoutB_; + using ElementC = int32_t; + using LayoutC = LayoutC_; + using Operator = mctlass::arch::OpMultiplyAdd; + using ArchTag = arch::Sm72; + + // check supported wmma shape for the given multiplicand data types + static_assert( + platform::is_same, Shape>::value || + platform::is_same, Shape>::value || + platform::is_same, Shape>::value, + "Supported list of wmma operator shape for s8 multiplicands are: 16x16x16, 8x32x16, and 32x8x16"); + + + // Wmma Fragment + using FragmentA = mxmaca::wmma::fragment< + mxmaca::wmma::matrix_a, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type, + typename MctlassToWmmaLayout::Layout>; + + using FragmentB = mxmaca::wmma::fragment< + mxmaca::wmma::matrix_b, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type, + typename MctlassToWmmaLayout::Layout>; + + using FragmentC = mxmaca::wmma::fragment< + mxmaca::wmma::accumulator, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type>; + + /// Performs a mxmaca::wmma matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + FragmentA const &A, + FragmentB const &B, + FragmentC const &C) const { + + mxmaca::wmma::mma_sync(D, A, B, C); + } + +#else + static_assert(false, "wmma.mma.sync interger type multiplicands is avialable only for SM72 and beyond"); +#endif + +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// WMMA template structure defines mxmaca::wmma::fragments and static assert for +// wmma native instruction sizes supported for uint8_t +// +//////////////////////////////////////////////////////////////////////////////// +template < +typename Shape_, +typename LayoutA_, +typename LayoutB_, +typename LayoutC_> +struct Wmma< + Shape_, ///< Size of the matrix product (concept: GemmShape) + uint8_t, ///< ElementA + LayoutA_, ///< LayoutA + uint8_t, ///< ElementB + LayoutB_, ///< LayoutB + int32_t, ///< ElementC + LayoutC_, ///< LayoutC + mctlass::arch::OpMultiplyAdd ///< Operator (multiply-add, xor.popc) +> { +#if defined(MCTLASS_ARCH_WMMA_SM72_ENABLED) + using Shape = Shape_; + using ElementA = uint8_t; + using LayoutA = LayoutA_; + using ElementB = uint8_t; + using LayoutB = LayoutB_; + using ElementC = int32_t; + using LayoutC = LayoutC_; + using Operator = mctlass::arch::OpMultiplyAdd; + using ArchTag = arch::Sm72; + + // check supported wmma shape for the given multiplicand data types + static_assert( + platform::is_same, Shape>::value || + platform::is_same, Shape>::value || + platform::is_same, Shape>::value, + "Supported list of wmma operator shape for u8 multiplicands are: 16x16x16, 8x32x16, and 32x8x16"); + + // Wmma Fragment + using FragmentA = mxmaca::wmma::fragment< + mxmaca::wmma::matrix_a, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type, + typename MctlassToWmmaLayout::Layout>; + + using FragmentB = mxmaca::wmma::fragment< + mxmaca::wmma::matrix_b, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type, + typename MctlassToWmmaLayout::Layout>; + + using FragmentC = mxmaca::wmma::fragment< + mxmaca::wmma::accumulator, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type>; + + /// Performs a mxmaca::wmma matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + FragmentA const &A, + FragmentB const &B, + FragmentC const &C) const { + + mxmaca::wmma::mma_sync(D, A, B, C); + } + +#else + static_assert(false, "wmma.mma.sync interger type multiplicands is avialable only for SM72 and beyond"); +#endif + +}; + +} // namespace arch +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/arch/wmma_sm75.h b/csrc/mctlass/include/mctlass/arch/wmma_sm75.h new file mode 100644 index 0000000..9861f3e --- /dev/null +++ b/csrc/mctlass/include/mctlass/arch/wmma_sm75.h @@ -0,0 +1,209 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Matrix multiply +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif +#include "mctlass/layout/matrix.h" + +//////////////////////////////////////////////////////////////////////////////// +namespace mctlass { +namespace arch { + +//////////////////////////////////////////////////////////////////////////////// +// +// WMMA template structure defines mxmaca::wmma::fragments and static assert for +// wmma native instruction sizes supported for mctlass::int4b_t (experimental::s4). +// +//////////////////////////////////////////////////////////////////////////////// +template < +typename Shape_, +typename LayoutA_, +typename LayoutB_, +typename LayoutC_> +struct Wmma< + Shape_, ///< Size of the matrix product (concept: GemmShape) + mctlass::int4b_t, ///< ElementA + LayoutA_, ///< LayoutA + mctlass::int4b_t, ///< ElementB + LayoutB_, ///< LayoutB + int32_t, ///< ElementC + LayoutC_, ///< LayoutC + mctlass::arch::OpMultiplyAdd ///< Operator (multiply-add, xor.popc) +> { +//#if defined(MCTLASS_ARCH_WMMA_SM75_ENABLED) +#if 0 //mcTlass not support int4b_t + using Shape = Shape_; + using ElementA = mctlass::int4b_t; + using LayoutA = LayoutA_; + using ElementB = mctlass::int4b_t; + using LayoutB = LayoutB_; + using ElementC = int32_t; + using LayoutC = LayoutC_; + using Operator = mctlass::arch::OpMultiplyAdd; + using ArchTag = arch::Sm75; + + // check supported wmma shape for the given multiplicand data types + static_assert( + platform::is_same, Shape>::value, + "Supported list of wmma operator shape for s8 multiplicands is: 8x8x32"); + + + // Wmma Fragment + using FragmentA = mxmaca::wmma::fragment< + mxmaca::wmma::matrix_a, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type, + typename MctlassToWmmaLayout::Layout>; + + using FragmentB = mxmaca::wmma::fragment< + mxmaca::wmma::matrix_b, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type, + typename MctlassToWmmaLayout::Layout>; + + using FragmentC = mxmaca::wmma::fragment< + mxmaca::wmma::accumulator, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type>; + + /// Performs a mxmaca::wmma matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + FragmentA const &A, + FragmentB const &B, + FragmentC const &C) const { + mxmaca::wmma::mma_sync(D, A, B, C); + + } + +#else + static_assert(false, "wmma.mma.sync interger type multiplicands is avialable only for SM75 and beyond"); +#endif + +}; + +//////////////////////////////////////////////////////////////////////////////// +// +// WMMA template structure defines mxmaca::wmma::fragments and static assert for +// wmma native instruction sizes supported for mctlass::uint1b_t (experimental::b1). +// +//////////////////////////////////////////////////////////////////////////////// +template < +typename Shape_, +typename LayoutA_, +typename LayoutB_, +typename LayoutC_> +struct Wmma< + Shape_, ///< Size of the matrix product (concept: GemmShape) + mctlass::uint1b_t, ///< ElementA + LayoutA_, ///< LayoutA + mctlass::uint1b_t, ///< ElementB + LayoutB_, ///< LayoutB + int32_t, ///< ElementC + LayoutC_, ///< LayoutC + mctlass::arch::OpXorPopc ///< Operator (multiply-add, xor.popc) +> { +//#if defined(MCTLASS_ARCH_WMMA_SM75_ENABLED) +#if 0 //mctlass not support uint1b_t + using Shape = Shape_; + using ElementA = mctlass::uint1b_t; + using LayoutA = LayoutA_; + using ElementB = mctlass::uint1b_t; + using LayoutB = LayoutB_; + using ElementC = int32_t; + using LayoutC = LayoutC_; + using Operator = mctlass::arch::OpXorPopc; + using ArchTag = arch::Sm75; + + // check supported wmma shape for the given multiplicand data types + static_assert( + platform::is_same, Shape>::value, + "Supported list of wmma operator shape for b1 multiplicands is: 8x8x128"); + + + // Wmma Fragment + using FragmentA = mxmaca::wmma::fragment< + mxmaca::wmma::matrix_a, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type, + typename MctlassToWmmaLayout::Layout>; + + using FragmentB = mxmaca::wmma::fragment< + mxmaca::wmma::matrix_b, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type, + typename MctlassToWmmaLayout::Layout>; + + using FragmentC = mxmaca::wmma::fragment< + mxmaca::wmma::accumulator, + Shape::kM, + Shape::kN, + Shape::kK, + typename MctlassToWmmaDataType::Type>; + + /// Performs a mxmaca::wmma matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + FragmentA const &A, + FragmentB const &B, + FragmentC const &C) const { + mxmaca::wmma::bmma_sync(D, A, B, C, mxmaca::wmma::experimental::bmmaBitOpXOR, + mxmaca::wmma::experimental::bmmaAccumulateOpPOPC); + } + +#else + static_assert(false, "wmma.mma.sync interger type multiplicands is avialable only for SM75 and beyond"); +#endif + +}; + +} // namespace arch +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/array.h b/csrc/mctlass/include/mctlass/array.h new file mode 100644 index 0000000..85b4392 --- /dev/null +++ b/csrc/mctlass/include/mctlass/array.h @@ -0,0 +1,2461 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Statically sized array of elements that accommodates all mctlass-supported numeric types + and is safe to use in a union. +*/ + +#pragma once +#include "mctlass/mctlass.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_types.h" +#include "mctlass/half.h" + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Statically sized array for any data type +template < + typename T, + int N, + bool RegisterSized = sizeof_bits::value >= 32 +> +class Array; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines the size of an Array<> in bits +template +struct sizeof_bits > { + static int const value = + int(sizeof(typename Array::Storage)) * 8 * int(Array::kStorageElements); +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Returns true if the argument is a power of 2 +MCTLASS_HOST_DEVICE +constexpr bool ispow2(unsigned x) { + return x && (!(x & (x - 1))); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Returns the largest power of two not greater than the argument. +MCTLASS_HOST_DEVICE +constexpr unsigned floor_pow_2(unsigned x) { + return (x == 0 || ispow2(x)) ? x : ((floor_pow_2(x >> 1)) << 1); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Statically sized array for any data type +template < + typename T, + int N +> +class Array { +public: + + /// Storage type + using Storage = T; + + /// Element type + using Element = T; + + /// Number of storage elements + //static std::size_t const kStorageElements = N; + static size_t const kStorageElements = N; + + /// Number of logical elements + static size_t const kElements = N; + + // + // C++ standard members + // + + typedef T value_type; + typedef size_t size_type; + typedef ptrdiff_t difference_type; + typedef value_type &reference; + typedef value_type const & const_reference; + typedef value_type *pointer; + typedef value_type const * const_pointer; + + // + // Iterators + // + + /// Bidirectional iterator over elements + class iterator { + + /// Pointer to object + T *ptr_; + + public: + + MCTLASS_HOST_DEVICE + iterator(): ptr_(nullptr) { } + + MCTLASS_HOST_DEVICE + iterator(T *_ptr): ptr_(_ptr) { } + + MCTLASS_HOST_DEVICE + iterator &operator++() { + ++ptr_; + return *this; + } + + MCTLASS_HOST_DEVICE + iterator &operator--() { + --ptr_; + return *this; + } + + MCTLASS_HOST_DEVICE + iterator operator++(int) { + iterator ret(*this); + ++ptr_; + return ret; + } + + MCTLASS_HOST_DEVICE + iterator operator--(int) { + iterator ret(*this); + --ptr_; + return ret; + } + + MCTLASS_HOST_DEVICE + T &operator*() const { + return *ptr_; + } + + MCTLASS_HOST_DEVICE + bool operator==(iterator const &other) const { + return ptr_ == other.ptr_; + } + + MCTLASS_HOST_DEVICE + bool operator!=(iterator const &other) const { + return ptr_ != other.ptr_; + } + }; + + /// Bidirectional constant iterator over elements + class const_iterator { + + /// Pointer to object + const T *ptr_; + + public: + + MCTLASS_HOST_DEVICE + const_iterator(): ptr_(nullptr) { } + + MCTLASS_HOST_DEVICE + const_iterator(T const *_ptr): ptr_(_ptr) { } + + MCTLASS_HOST_DEVICE + const_iterator &operator++() { + ++ptr_; + return *this; + } + + MCTLASS_HOST_DEVICE + const_iterator &operator--() { + --ptr_; + return *this; + } + + MCTLASS_HOST_DEVICE + const_iterator operator++(int) { + const_iterator ret(*this); + ++ptr_; + return ret; + } + + MCTLASS_HOST_DEVICE + const_iterator operator--(int) { + const_iterator ret(*this); + --ptr_; + return ret; + } + + MCTLASS_HOST_DEVICE + T const &operator*() const { + return *ptr_; + } + + MCTLASS_HOST_DEVICE + bool operator==(const_iterator const &other) const { + return ptr_ == other.ptr_; + } + + MCTLASS_HOST_DEVICE + bool operator!=(const_iterator const &other) const { + return ptr_ != other.ptr_; + } + }; + + /// Bidirectional iterator over elements + class reverse_iterator { + + /// Pointer to object + T *ptr_; + + public: + + MCTLASS_HOST_DEVICE + reverse_iterator(): ptr_(nullptr) { } + + MCTLASS_HOST_DEVICE + reverse_iterator(T *_ptr): ptr_(_ptr) { } + + MCTLASS_HOST_DEVICE + reverse_iterator &operator++() { + --ptr_; + return *this; + } + + MCTLASS_HOST_DEVICE + reverse_iterator &operator--() { + ++ptr_; + return *this; + } + + MCTLASS_HOST_DEVICE + reverse_iterator operator++(int) { + iterator ret(*this); + --ptr_; + return ret; + } + + MCTLASS_HOST_DEVICE + reverse_iterator operator--(int) { + iterator ret(*this); + ++ptr_; + return ret; + } + + MCTLASS_HOST_DEVICE + T &operator*() const { + return *(ptr_ - 1); + } + + MCTLASS_HOST_DEVICE + bool operator==(reverse_iterator const &other) const { + return ptr_ == other.ptr_; + } + + MCTLASS_HOST_DEVICE + bool operator!=(reverse_iterator const &other) const { + return ptr_ != other.ptr_; + } + }; + + /// Bidirectional constant iterator over elements + class const_reverse_iterator { + + /// Pointer to object + T const *ptr_; + + public: + + MCTLASS_HOST_DEVICE + const_reverse_iterator(): ptr_(nullptr) { } + + MCTLASS_HOST_DEVICE + const_reverse_iterator(T const *_ptr): ptr_(_ptr) { } + + MCTLASS_HOST_DEVICE + const_reverse_iterator &operator++() { + --ptr_; + return *this; + } + + MCTLASS_HOST_DEVICE + const_reverse_iterator &operator--() { + ++ptr_; + return *this; + } + + MCTLASS_HOST_DEVICE + const_reverse_iterator operator++(int) { + const_reverse_iterator ret(*this); + --ptr_; + return ret; + } + + MCTLASS_HOST_DEVICE + const_reverse_iterator operator--(int) { + const_reverse_iterator ret(*this); + ++ptr_; + return ret; + } + + MCTLASS_HOST_DEVICE + T const &operator*() const { + return *(ptr_ - 1); + } + + MCTLASS_HOST_DEVICE + bool operator==(const_iterator const &other) const { + return ptr_ == other.ptr_; + } + + MCTLASS_HOST_DEVICE + bool operator!=(const_iterator const &other) const { + return ptr_ != other.ptr_; + } + }; + +private: + + /// Internal storage + Storage storage[kElements]; + +public: + + #if 0 + MCTLASS_HOST_DEVICE + Array() { } + + MCTLASS_HOST_DEVICE + Array(Array const &x) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kElements; ++i) { + storage[i] = x.storage[i]; + } + } + #endif + + /// Efficient clear method + MCTLASS_HOST_DEVICE + void clear() { + fill(T(0)); + } + + MCTLASS_HOST_DEVICE + reference at(size_type pos) { + return reinterpret_cast(storage[pos]); + } + + MCTLASS_HOST_DEVICE + const_reference at(size_type pos) const { + return reinterpret_cast(storage[pos]); + } + + MCTLASS_HOST_DEVICE + reference operator[](size_type pos) { + return reinterpret_cast(storage[pos]); + } + + MCTLASS_HOST_DEVICE + const_reference operator[](size_type pos) const { + return reinterpret_cast(storage[pos]); + } + + MCTLASS_HOST_DEVICE + reference front() { + return reinterpret_cast(storage[0]); + } + + MCTLASS_HOST_DEVICE + const_reference front() const { + return reinterpret_cast(storage[0]); + } + + MCTLASS_HOST_DEVICE + reference back() { + return reinterpret_cast(storage[kStorageElements - 1]); + } + + MCTLASS_HOST_DEVICE + const_reference back() const { + return reinterpret_cast(storage[kStorageElements - 1]); + } + + MCTLASS_HOST_DEVICE + pointer data() { + return reinterpret_cast(storage); + } + + MCTLASS_HOST_DEVICE + const_pointer data() const { + return reinterpret_cast(storage); + } + + MCTLASS_HOST_DEVICE + pointer raw_data() { + return reinterpret_cast(storage); + } + + MCTLASS_HOST_DEVICE + const_pointer raw_data() const { + return reinterpret_cast(storage); + } + + + MCTLASS_HOST_DEVICE + constexpr bool empty() const { + return !kElements; + } + + MCTLASS_HOST_DEVICE + constexpr size_type size() const { + return kElements; + } + + MCTLASS_HOST_DEVICE + constexpr size_type max_size() const { + return kElements; + } + + MCTLASS_HOST_DEVICE + void fill(T const &value) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kElements; ++i) { + storage[i] = static_cast(value); + } + } + + MCTLASS_HOST_DEVICE + iterator begin() { + return iterator(storage); + } + + MCTLASS_HOST_DEVICE + const_iterator begin() const { + return cbegin(); + } + + MCTLASS_HOST_DEVICE + const_iterator cbegin() const { + return const_iterator(storage); + } + + MCTLASS_HOST_DEVICE + iterator end() { + return iterator(reinterpret_cast(storage + kStorageElements)); + } + + MCTLASS_HOST_DEVICE + const_iterator end() const { + return cend(); + } + + MCTLASS_HOST_DEVICE + const_iterator cend() const { + return const_iterator(reinterpret_cast(storage + kStorageElements)); + } + + MCTLASS_HOST_DEVICE + reverse_iterator rbegin() { + return reverse_iterator(reinterpret_cast(storage + kStorageElements)); + } + + MCTLASS_HOST_DEVICE + const_reverse_iterator rbegin() const { + return crbegin(); + } + + MCTLASS_HOST_DEVICE + const_reverse_iterator crbegin() const { + return const_reverse_iterator(reinterpret_cast(storage + kStorageElements)); + } + + MCTLASS_HOST_DEVICE + reverse_iterator rend() { + return reverse_iterator(reinterpret_cast(storage)); + } + + MCTLASS_HOST_DEVICE + const_reverse_iterator rend() const { + return crend(); + } + + MCTLASS_HOST_DEVICE + const_reverse_iterator crend() const { + return const_reverse_iterator(reinterpret_cast(storage)); + } + + // + // Comparison operators + // + +}; + + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Factories +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +MCTLASS_HOST_DEVICE +Array make_Array(Element x) { + Array m; + m[0] = x; + return m; +} + +template +MCTLASS_HOST_DEVICE +Array make_Array(Element x, Element y) { + Array m; + m[0] = x; + m[1] = y; + return m; +} + +template +MCTLASS_HOST_DEVICE +Array make_Array(Element x, Element y, Element z) { + Array m; + m[0] = x; + m[1] = y; + m[2] = z; + return m; +} + +template +MCTLASS_HOST_DEVICE +Array make_Array(Element x, Element y, Element z, Element w) { + Array m; + m[0] = x; + m[1] = y; + m[2] = z; + m[3] = w; + return m; +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// +// functional.h numeric specializations +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct absolute_value_op< Array > { + + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs) const { + + Array result; + absolute_value_op scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i]); + } + + return result; + } +}; + +template +struct plus> { + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs, Array const &rhs) const { + + Array result; + plus scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i], rhs[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs, T const &scalar) const { + + Array result; + plus scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i], scalar); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( T const &scalar, Array const &rhs) const { + + Array result; + plus scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(scalar, rhs[i]); + } + + return result; + } +}; +template +struct minus> { + + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs, Array const &rhs) const { + + Array result; + minus scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i], rhs[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs, T const &scalar) const { + + Array result; + minus scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i], scalar); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( T const &scalar, Array const &rhs) const { + + Array result; + minus scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(scalar, rhs[i]); + } + + return result; + } +}; + +template +struct multiplies> { + + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs, Array const &rhs) const { + + Array result; + multiplies scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i], rhs[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs, T const &scalar) const { + + Array result; + multiplies scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i], scalar); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( T const &scalar, Array const &rhs) const { + + Array result; + multiplies scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(scalar, rhs[i]); + } + + return result; + } +}; + +template +struct divides> { + + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs, Array const &rhs) const { + + Array result; + divides scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i], rhs[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs, T const &scalar) const { + + Array result; + divides scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i], scalar); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( T const &scalar, Array const &rhs) const { + + Array result; + divides scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(scalar, rhs[i]); + } + + return result; + } +}; + +template +struct maximum> { + + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs, Array const &rhs) const { + + Array result; + maximum scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i], rhs[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs, T const &scalar) const { + + Array result; + maximum scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i], scalar); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( T const &scalar, Array const &rhs) const { + + Array result; + maximum scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(scalar, rhs[i]); + } + + return result; + } +}; + +template +struct minimum> { + + MCTLASS_HOST_DEVICE + static T scalar_op(T const &lhs, T const &rhs) { + return (rhs < lhs ? rhs : lhs); + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs, Array const &rhs) const { + + Array result; + minimum scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i], rhs[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs, T const &scalar) const { + + Array result; + minimum scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i], scalar); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( T const &scalar, Array const &rhs) const { + + Array result; + minimum scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(scalar, rhs[i]); + } + + return result; + } +}; + +template +struct negate> { + + MCTLASS_HOST_DEVICE + Array operator()(Array const &lhs) const { + + Array result; + negate scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(lhs[i]); + } + + return result; + } +}; + +/// Fused multiply-add +template +struct multiply_add, Array, Array> { + + MCTLASS_HOST_DEVICE + Array operator()(Array const &a, Array const &b, Array const &c) const { + + Array result; + multiply_add scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(a[i], b[i], c[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const &a, T const &scalar, Array const &c) const { + + Array result; + multiply_add scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(a[i], scalar, c[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(T const &scalar, Array const &b, Array const &c) const { + + Array result; + multiply_add scalar_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = scalar_op(scalar, b[i], c[i]); + } + + return result; + } +}; + +/// Fused multiply-add-relu0 +template +struct multiply_add_relu0, Array, Array> { + + MCTLASS_HOST_DEVICE + Array operator()(Array const &a, Array const &b, Array const &c) const { + + Array result; + multiply_add scalar_op; + maximum mx; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = mx(scalar_op(a[i], b[i], c[i]), T(0)); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const &a, T const &scalar, Array const &c) const { + + Array result; + multiply_add scalar_op; + maximum mx; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = mx(scalar_op(a[i], scalar, c[i]), T(0)); + } + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(T const &scalar, Array const &b, Array const &c) const { + + Array result; + multiply_add scalar_op; + maximum mx; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = mx(scalar_op(scalar, b[i], c[i]), T(0)); + } + + return result; + } +}; + + +template +struct conjugate > { + MCTLASS_HOST_DEVICE + Array operator()(Array const &a) const { + + conjugate conj_op; + + Array ca; + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + ca[i] = conj_op(a[i]); + } + return ca; + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// +// functional.h numeric specializations targeting SIMD instructions in device code. +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct plus> { + MCTLASS_HOST_DEVICE + Array operator()(Array const & lhs, Array const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs); + __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hadd2(lhs_ptr[i], rhs_ptr[i]); + } + + if (N % 2) { + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs); + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs); + __half d_residual = __hadd(a_residual_ptr[N - 1], b_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = lhs[i] + rhs[i]; + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(half_t const & lhs, Array const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 lhs_pair = __half2half2(reinterpret_cast<__half const &>(lhs)); + __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hadd2(lhs_pair, rhs_ptr[i]); + } + + if (N % 2) { + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs); + __half d_residual = __hadd(reinterpret_cast<__half const &>(lhs), b_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = lhs + rhs[i]; + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const & lhs, half_t const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs); + __half2 rhs_pair = __half2half2(reinterpret_cast<__half const &>(rhs)); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hadd2(lhs_ptr[i], rhs_pair); + } + + if (N % 2) { + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs); + __half d_residual = __hadd(a_residual_ptr[N - 1], reinterpret_cast<__half const &>(rhs)); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = lhs[i] + rhs; + } + #endif + + return result; + } +}; + +template +struct minus> { + MCTLASS_HOST_DEVICE + Array operator()(Array const & lhs, Array const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs); + __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hsub2(lhs_ptr[i], rhs_ptr[i]); + } + + if (N % 2) { + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs); + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs); + __half d_residual = __hsub(a_residual_ptr[N - 1], b_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = lhs[i] - rhs[i]; + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(half_t const & lhs, Array const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 lhs_pair = __half2half2(reinterpret_cast<__half const &>(lhs)); + __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hsub2(lhs_pair, rhs_ptr[i]); + } + + if (N % 2) { + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs); + __half d_residual = __hsub(reinterpret_cast<__half const &>(lhs), b_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = lhs - rhs[i]; + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const & lhs, half_t const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs); + __half2 rhs_pair = __half2half2(reinterpret_cast<__half const &>(rhs)); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hsub2(lhs_ptr[i], rhs_pair); + } + + if (N % 2) { + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs); + __half d_residual = __hsub(a_residual_ptr[N - 1], reinterpret_cast<__half const &>(rhs)); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = lhs[i] - rhs; + } + #endif + + return result; + } +}; + +template +struct multiplies> { + MCTLASS_HOST_DEVICE + Array operator()(Array const & lhs, Array const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs); + __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hmul2(lhs_ptr[i], rhs_ptr[i]); + } + + if (N % 2) { + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs); + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs); + __half d_residual = __hmul(a_residual_ptr[N - 1], b_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = lhs[i] * rhs[i]; + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(half_t const & lhs, Array const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 lhs_pair = __half2half2(reinterpret_cast<__half const &>(lhs)); + __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hmul2(lhs_pair, rhs_ptr[i]); + } + + if (N % 2) { + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs); + + __half d_residual = __hmul( + reinterpret_cast<__half const &>(lhs), + b_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = lhs * rhs[i]; + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const & lhs, half_t const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs); + __half2 rhs_pair = __half2half2(reinterpret_cast<__half const &>(rhs)); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hmul2(lhs_ptr[i], rhs_pair); + } + + if (N % 2) { + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs); + + __half d_residual = __hmul( + a_residual_ptr[N - 1], + reinterpret_cast<__half const &>(rhs)); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = lhs[i] * rhs; + } + #endif + + return result; + } +}; + +template +struct divides> { + MCTLASS_HOST_DEVICE + Array operator()(Array const & lhs, Array const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs); + __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __h2div(lhs_ptr[i], rhs_ptr[i]); + } + + if (N % 2) { + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs); + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs); + + __half d_residual = __hdiv( + a_residual_ptr[N - 1], + b_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = lhs[i] / rhs[i]; + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(half_t const & lhs, Array const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 lhs_pair = __half2half2(reinterpret_cast<__half const &>(lhs)); + __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __h2div(lhs_pair, rhs_ptr[i]); + } + + if (N % 2) { + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs); + + __half d_residual = __hdiv( + reinterpret_cast<__half const &>(lhs), + b_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = lhs / rhs[i]; + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const & lhs, half_t const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs); + __half2 rhs_pair = __half2half2(reinterpret_cast<__half const &>(rhs)); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __h2div(lhs_ptr[i], rhs_pair); + } + + if (N % 2) { + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs); + + __half d_residual = __hdiv( + a_residual_ptr[N - 1], + reinterpret_cast<__half const &>(rhs)); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = lhs[i] / rhs; + } + #endif + + return result; + } +}; + +template +struct negate> { + MCTLASS_HOST_DEVICE + Array operator()(Array const & lhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *source_ptr = reinterpret_cast<__half2 const *>(&lhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hneg2(source_ptr[i]); + } + + if (N % 2) { + half_t x = lhs[N - 1]; + __half lhs_val = -reinterpret_cast<__half const &>(x); + result[N - 1] = reinterpret_cast(lhs_val); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = -lhs[i]; + } + #endif + + return result; + } +}; + +/// Fused multiply-add +template +struct multiply_add, Array, Array> { + + MCTLASS_HOST_DEVICE + Array operator()( + Array const &a, + Array const &b, + Array const &c) const { + + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *a_ptr = reinterpret_cast<__half2 const *>(&a); + __half2 const *b_ptr = reinterpret_cast<__half2 const *>(&b); + __half2 const *c_ptr = reinterpret_cast<__half2 const *>(&c); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hfma2(a_ptr[i], b_ptr[i], c_ptr[i]); + } + + if (N % 2) { + + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&a); + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&b); + __half const *c_residual_ptr = reinterpret_cast<__half const *>(&c); + + __half d_residual = __hfma( + a_residual_ptr[N - 1], + b_residual_ptr[N - 1], + c_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + multiply_add op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = op(a[i], b[i], c[i]); + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( + half_t const &a, + Array const &b, + Array const &c) const { + + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 a_pair = __half2half2(reinterpret_cast<__half const &>(a)); + __half2 const *b_ptr = reinterpret_cast<__half2 const *>(&b); + __half2 const *c_ptr = reinterpret_cast<__half2 const *>(&c); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hfma2(a_pair, b_ptr[i], c_ptr[i]); + } + + if (N % 2) { + + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&b); + __half const *c_residual_ptr = reinterpret_cast<__half const *>(&c); + __half d_residual = __hfma( + reinterpret_cast<__half const &>(a), + b_residual_ptr[N - 1], + c_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + multiply_add op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = op(a, b[i], c[i]); + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( + Array const &a, + half_t const &b, + Array const &c) const { + + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *a_ptr = reinterpret_cast<__half2 const *>(&a); + __half2 b_pair = __half2half2(reinterpret_cast<__half const &>(b)); + __half2 const *c_ptr = reinterpret_cast<__half2 const *>(&c); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hfma2(a_ptr[i], b_pair, c_ptr[i]); + } + + if (N % 2) { + + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&a); + __half const *c_residual_ptr = reinterpret_cast<__half const *>(&c); + + __half d_residual = __hfma( + a_residual_ptr[N - 1], + reinterpret_cast<__half const &>(b), + c_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + multiply_add op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = op(a[i], b, c[i]); + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( + Array const &a, + Array const &b, + half_t const &c) const { + + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *a_ptr = reinterpret_cast<__half2 const *>(&a); + __half2 const *b_ptr = reinterpret_cast<__half2 const *>(&b); + __half2 c_pair = __half2half2(reinterpret_cast<__half const &>(c)); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hfma2(a_ptr[i], b_ptr[i], c_pair); + } + + if (N % 2) { + + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&a); + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&b); + + __half d_residual = __hfma( + a_residual_ptr[N - 1], + b_residual_ptr[N - 1], + reinterpret_cast<__half const &>(c)); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + multiply_add op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = op(a[i], b[i], c); + } + #endif + + return result; + } +}; + +/// Fused multiply-add-relu0 +template +struct multiply_add_relu0, Array, Array> { + + MCTLASS_HOST_DEVICE + Array operator()( + Array const &a, + Array const &b, + Array const &c) const { + + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *a_ptr = reinterpret_cast<__half2 const *>(&a); + __half2 const *b_ptr = reinterpret_cast<__half2 const *>(&b); + __half2 const *c_ptr = reinterpret_cast<__half2 const *>(&c); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hfma2_relu(a_ptr[i], b_ptr[i], c_ptr[i]); + } + + if (N % 2) { + + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&a); + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&b); + __half const *c_residual_ptr = reinterpret_cast<__half const *>(&c); + + __half d_residual = __hfma_relu( + a_residual_ptr[N - 1], + b_residual_ptr[N - 1], + c_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + multiply_add op; + maximum mx; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = mx(op(a[i], b[i], c[i]), (half_t)0); + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( + half_t const &a, + Array const &b, + Array const &c) const { + + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 a_pair = __half2half2(reinterpret_cast<__half const &>(a)); + __half2 const *b_ptr = reinterpret_cast<__half2 const *>(&b); + __half2 const *c_ptr = reinterpret_cast<__half2 const *>(&c); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hfma2_relu(a_pair, b_ptr[i], c_ptr[i]); + } + + if (N % 2) { + + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&b); + __half const *c_residual_ptr = reinterpret_cast<__half const *>(&c); + __half d_residual = __hfma_relu( + reinterpret_cast<__half const &>(a), + b_residual_ptr[N - 1], + c_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + multiply_add op; + maximum mx; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = mx(op(a, b[i], c[i]), half_t(0)); + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( + Array const &a, + half_t const &b, + Array const &c) const { + + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *a_ptr = reinterpret_cast<__half2 const *>(&a); + __half2 b_pair = __half2half2(reinterpret_cast<__half const &>(b)); + __half2 const *c_ptr = reinterpret_cast<__half2 const *>(&c); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hfma2_relu(a_ptr[i], b_pair, c_ptr[i]); + } + + if (N % 2) { + + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&a); + __half const *c_residual_ptr = reinterpret_cast<__half const *>(&c); + + __half d_residual = __hfma_relu( + a_residual_ptr[N - 1], + reinterpret_cast<__half const &>(b), + c_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + multiply_add op; + maximum mx; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = mx(op(a[i], b, c[i]), half_t(0)); + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( + Array const &a, + Array const &b, + half_t const &c) const { + + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *a_ptr = reinterpret_cast<__half2 const *>(&a); + __half2 const *b_ptr = reinterpret_cast<__half2 const *>(&b); + __half2 c_pair = __half2half2(reinterpret_cast<__half const &>(c)); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hfma2_relu(a_ptr[i], b_ptr[i], c_pair); + } + + if (N % 2) { + + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&a); + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&b); + + __half d_residual = __hfma_relu( + a_residual_ptr[N - 1], + b_residual_ptr[N - 1], + reinterpret_cast<__half const &>(c)); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + multiply_add op; + maximum mx; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = mx(op(a[i], b[i], c), half_t(0)); + } + #endif + + return result; + } +}; + +template +struct minimum> { + MCTLASS_HOST_DEVICE + Array operator()(Array const & lhs, Array const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs); + __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hmin2(lhs_ptr[i], rhs_ptr[i]); + } + + if (N % 2) { + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs); + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs); + + __half d_residual = __hmin( + a_residual_ptr[N - 1], + b_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = (rhs[i] < lhs[i] ? rhs[i] : lhs[i]); + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(half_t const & lhs, Array const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 lhs_pair = __half2half2(reinterpret_cast<__half const &>(lhs)); + __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hmin2(lhs_pair, rhs_ptr[i]); + } + + if (N % 2) { + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs); + + __half d_residual = __hmin( + reinterpret_cast<__half const &>(lhs), + b_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = (rhs[i] < lhs ? rhs[i] : lhs); + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const & lhs, half_t const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs); + __half2 rhs_pair = __half2half2(reinterpret_cast<__half const &>(rhs)); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hmin2(lhs_ptr[i], rhs_pair); + } + + if (N % 2) { + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs); + + __half d_residual = __hmin( + a_residual_ptr[N - 1], + reinterpret_cast<__half const &>(rhs)); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = (rhs < lhs[i] ? rhs : lhs[i]); + } + #endif + + return result; + } +}; + +template +struct maximum> { + MCTLASS_HOST_DEVICE + Array operator()(Array const & lhs, Array const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs); + __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hmax2(lhs_ptr[i], rhs_ptr[i]); + } + + if (N % 2) { + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs); + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs); + + __half d_residual = __hmax( + a_residual_ptr[N - 1], + b_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = (lhs[i] < rhs[i] ? rhs[i] : lhs[i]); + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(half_t const & lhs, Array const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 lhs_pair = __half2half2(reinterpret_cast<__half const &>(lhs)); + __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hmax2(lhs_pair, rhs_ptr[i]); + } + + if (N % 2) { + __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs); + + __half d_residual = __hmax( + reinterpret_cast<__half const &>(lhs), + b_residual_ptr[N - 1]); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = (lhs < rhs[i] ? rhs[i] : lhs); + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const & lhs, half_t const &rhs) const { + Array result; + #if defined(__MACA_ARCH__) + + __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs); + __half2 rhs_pair = __half2half2(reinterpret_cast<__half const &>(rhs)); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = __hmax2(lhs_ptr[i], rhs_pair); + } + + if (N % 2) { + __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs); + + __half d_residual = __hmax( + a_residual_ptr[N - 1], + reinterpret_cast<__half const &>(rhs)); + + result[N - 1] = reinterpret_cast(d_residual); + } + + #else + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = (lhs[i] < rhs ? rhs : lhs[i]); + } + #endif + + return result; + } +}; + +/// Fused multiply-add +template +struct multiply_add, Array, Array> { + + MCTLASS_HOST_DEVICE + Array operator()( + Array const &a, + Array const &b, + Array const &c) const { + + Array result; + // #if defined(__MACA_ARCH__) + #if 0 + + unsigned *result_ptr = reinterpret_cast(&result); + unsigned const *a_ptr = reinterpret_cast(&a); + unsigned const *b_ptr = reinterpret_cast(&b); + unsigned const *c_ptr = reinterpret_cast(&c); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + asm ("fma.rn.bf16x2 %0, %1, %2, %3;\n" + : "=r"(result_ptr[i]) + : "r"(a_ptr[i]), "r"(b_ptr[i]), "r"(c_ptr[i]) + ); + } + + if (N % 2) { + + uint16_t *result_ptr = reinterpret_cast(&result); + uint16_t const *a_residual_ptr = reinterpret_cast(&a); + uint16_t const *b_residual_ptr = reinterpret_cast(&b); + uint16_t const *c_residual_ptr = reinterpret_cast(&c); + + asm ("fma.rn.bf16 %0, %1, %2, %3;\n" + : "=h"(result_ptr[N - 1]) + : "h"(a_residual_ptr[N - 1]), "h"(b_residual_ptr[N - 1]), "h"(c_residual_ptr[N - 1]) + ); + } + + #else + + multiply_add op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = op(a[i], b[i], c[i]); + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( + bfloat16_t const &a, + Array const &b, + Array const &c) const { + + Array result; + // #if defined(__MACA_ARCH__) + #if 0 + + unsigned *result_ptr = reinterpret_cast(&result); + + unsigned const *b_ptr = reinterpret_cast(&b); + unsigned const *c_ptr = reinterpret_cast(&c); + + unsigned a_packed = static_cast(a.raw()); + a_packed = (a_packed | (a_packed << 16)); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + asm ("fma.rn.bf16x2 %0, %1, %2, %3;\n" + : "=r"(result_ptr[i]) + : "r"(a_packed), "r"(b_ptr[i]), "r"(c_ptr[i]) + ); + } + + if (N % 2) { + + uint16_t *result_ptr = reinterpret_cast(&result); + uint16_t const *a_residual_ptr = reinterpret_cast(&a); + uint16_t const *b_residual_ptr = reinterpret_cast(&b); + uint16_t const *c_residual_ptr = reinterpret_cast(&c); + + asm ("fma.rn.bf16 %0, %1, %2, %3;\n" + : "=h"(result_ptr[N - 1]) + : "h"(a_residual_ptr[0]), "h"(b_residual_ptr[N - 1]), "h"(c_residual_ptr[N - 1]) + ); + } + + #else + + multiply_add op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = op(a, b[i], c[i]); + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( + Array const &a, + bfloat16_t const &b, + Array const &c) const { + + Array result; + // #if defined(__MACA_ARCH__) + #if 0 + + unsigned *result_ptr = reinterpret_cast(&result); + + unsigned const *a_ptr = reinterpret_cast(&a); + unsigned const *c_ptr = reinterpret_cast(&c); + + unsigned b_packed = static_cast(b.raw()); + b_packed = (b_packed | (b_packed << 16)); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + asm ("fma.rn.bf16x2 %0, %1, %2, %3;\n" + : "=r"(result_ptr[i]) + : "r"(a_ptr[i]), "r"(b_packed), "r"(c_ptr[i]) + ); + } + + if (N % 2) { + + uint16_t *result_ptr = reinterpret_cast(&result); + uint16_t const *a_residual_ptr = reinterpret_cast(&a); + uint16_t const *b_residual_ptr = reinterpret_cast(&b); + uint16_t const *c_residual_ptr = reinterpret_cast(&c); + + asm ("fma.rn.bf16 %0, %1, %2, %3;\n" + : "=h"(result_ptr[N - 1]) + : "h"(a_residual_ptr[N - 1]), "h"(b_residual_ptr[0]), "h"(c_residual_ptr[N - 1]) + ); + } + + #else + + multiply_add op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = op(a[i], b, c[i]); + } + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + Array operator()( + Array const &a, + Array const &b, + bfloat16_t const &c) const { + + Array result; + // #if defined(__MACA_ARCH__) + #if 0 + + unsigned *result_ptr = reinterpret_cast(&result); + + unsigned const *a_ptr = reinterpret_cast(&a); + unsigned const *b_ptr = reinterpret_cast(&b); + + unsigned c_packed = static_cast(c.raw()); + c_packed = (c_packed | (c_packed << 16)); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + asm ("fma.rn.bf16x2 %0, %1, %2, %3;\n" + : "=r"(result_ptr[i]) + : "r"(a_ptr[i]), "r"(b_ptr[i]), "r"(c_packed) + ); + } + + if (N % 2) { + + uint16_t *result_ptr = reinterpret_cast(&result); + uint16_t const *a_residual_ptr = reinterpret_cast(&a); + uint16_t const *b_residual_ptr = reinterpret_cast(&b); + uint16_t const *c_residual_ptr = reinterpret_cast(&c); + + asm ("fma.rn.bf16 %0, %1, %2, %3;\n" + : "=h"(result_ptr[N - 1]) + : "h"(a_residual_ptr[N - 1]), "h"(b_residual_ptr[N - 1]), "h"(c_residual_ptr[0]) + ); + } + + #else + + multiply_add op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = op(a[i], b[i], c); + } + #endif + + return result; + } +}; + + +/// bit_and +template +struct bit_and> { + MCTLASS_HOST_DEVICE + Array operator()(Array const &a, Array const &b) const { + using ArrayType = Array; + using Storage = typename ArrayType::Storage; + ArrayType result; + + Storage *result_data = result.raw_data(); + Storage const *a_data = a.raw_data(); + Storage const *b_data = b.raw_data(); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < ArrayType::kStorageElements; ++i) { + result_data[i] = (a_data[i] & b_data[i]); + } + + return result; + } +}; + + +/// bit_or +template +struct bit_or> { + MCTLASS_HOST_DEVICE + Array operator()(Array const &a, Array const &b) const { + using ArrayType = Array; + using Storage = typename ArrayType::Storage; + ArrayType result; + + Storage *result_data = result.raw_data(); + Storage const *a_data = a.raw_data(); + Storage const *b_data = b.raw_data(); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < ArrayType::kStorageElements; ++i) { + result_data[i] = (a_data[i] | b_data[i]); + } + + return result; + } +}; + + +/// bit_not +template +struct bit_not> { + MCTLASS_HOST_DEVICE + Array operator()(Array const &a) const { + using ArrayType = Array; + using Storage = typename ArrayType::Storage; + ArrayType result; + + Storage *result_data = result.raw_data(); + Storage const *a_data = a.raw_data(); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < ArrayType::kStorageElements; ++i) { + result_data[i] = (~a_data[i]); + } + + return result; + } +}; + + +/// bit_xor +template +struct bit_xor> { + MCTLASS_HOST_DEVICE + Array operator()(Array const &a, Array const &b) const { + using ArrayType = Array; + using Storage = typename ArrayType::Storage; + ArrayType result; + + Storage *result_data = result.raw_data(); + Storage const *a_data = a.raw_data(); + Storage const *b_data = b.raw_data(); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < ArrayType::kStorageElements; ++i) { + result_data[i] = (a_data[i] ^ b_data[i]); + } + + return result; + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Operator overloads +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +MCTLASS_HOST_DEVICE +Array operator+(Array const &lhs, Array const &rhs) { + plus> op; + return op(lhs, rhs); +} + +template +MCTLASS_HOST_DEVICE +Array operator-(Array const &lhs, Array const &rhs) { + minus> op; + return op(lhs, rhs); +} + +template +MCTLASS_HOST_DEVICE +Array operator-(Array const &lhs) { + negate> op; + return op(lhs); +} + +template +MCTLASS_HOST_DEVICE +Array operator*(Array const &lhs, Array const &rhs) { + multiplies> op; + return op(lhs, rhs); +} + +template +MCTLASS_HOST_DEVICE +Array operator*(T lhs, Array const &rhs) { + multiplies> op; + return op(lhs, rhs); +} + +template +MCTLASS_HOST_DEVICE +Array operator*(Array const &lhs, T rhs) { + multiplies> op; + return op(lhs, rhs); +} + +template +MCTLASS_HOST_DEVICE +Array operator/(Array const &lhs, Array const &rhs) { + divides> op; + return op(lhs, rhs); +} + +template +MCTLASS_HOST_DEVICE +Array fma(Array const &a, Array const &b, Array const &c) { + multiply_add> op; + return op(a, b, c); +} + +template +MCTLASS_HOST_DEVICE +Array fma(T a, Array const &b, Array const &c) { + multiply_add> op; + return op(a, b, c); +} + +template +MCTLASS_HOST_DEVICE +Array fma(Array const &a, T b, Array const &c) { + multiply_add> op; + return op(a, b, c); +} + +template +MCTLASS_HOST_DEVICE +Array fma(Array const &a, Array const &b, T c) { + multiply_add> op; + return op(a, b, c); +} + + +//////////////////////////////////////////////////////////////////////////////////////////////////// + + + +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#include "mctlass/array_subbyte.h" + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// AlignedArray +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Aligned array type +template < + /// Element type + typename T, + /// Number of elements in the array + int N, + /// Alignment requirement in bytes + int Alignment = sizeof_bits::value * N / 8 +> +class alignas(Alignment) AlignedArray: public Array { +public: + +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/array_planar_complex.h b/csrc/mctlass/include/mctlass/array_planar_complex.h new file mode 100644 index 0000000..f8d7b2e --- /dev/null +++ b/csrc/mctlass/include/mctlass/array_planar_complex.h @@ -0,0 +1,103 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Array holding planar complex elements +template +struct ArrayPlanarComplex { + + /// Underlying real element + using Element = Element_; + + /// Number of logical elements + static size_t const kElements = N; + + /// Underlying Fragment of real-valued elemenets + using ArrayReal = Array; + +public: + + /// Fragment of real-valued elements representing the real part + ArrayReal real; + + /// Fragment of real-valued elements representing the imaginary part + ArrayReal imag; + +public: + + /// Ctor + MCTLASS_HOST_DEVICE + ArrayPlanarComplex() { } + + /// Ctor + MCTLASS_HOST_DEVICE + ArrayPlanarComplex( + ArrayReal const &real_, + ArrayReal const &imag_ + ): + real(real_), imag(imag_) { } + + /// Sets the array to zero efficiently + MCTLASS_HOST_DEVICE + void clear() { + real.clear(); + imag.clear(); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Helper to deduce template arguments +template +MCTLASS_HOST_DEVICE +ArrayPlanarComplex +make_ArrayPlanarComplex(Array const &real, Array const &imag) { + return ArrayPlanarComplex(real, imag); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/array_subbyte.h b/csrc/mctlass/include/mctlass/array_subbyte.h new file mode 100644 index 0000000..4976f9d --- /dev/null +++ b/csrc/mctlass/include/mctlass/array_subbyte.h @@ -0,0 +1,564 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Statically sized array of elements that accommodates all MCTLASS-supported numeric types + and is safe to use in a union. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/platform/platform.h" + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Statically sized array for any data type +template < + typename T, + int N +> +class Array { +public: + + static int const kSizeBits = sizeof_bits::value * N; + + /// Storage type + using Storage = typename platform::conditional< + ((kSizeBits % 32) != 0), + typename platform::conditional< + ((kSizeBits % 16) != 0), + uint8_t, + uint16_t + >::type, + uint32_t + >::type; + + /// Element type + using Element = T; + + /// Number of logical elements per stored object + static int const kElementsPerStoredItem = int(sizeof(Storage) * 8) / sizeof_bits::value; + + /// Number of storage elements + static size_t const kStorageElements = N / kElementsPerStoredItem; + + /// Number of logical elements + static size_t const kElements = N; + + /// Bitmask for covering one item + static Storage const kMask = ((Storage(1) << sizeof_bits::value) - 1); + + // + // C++ standard members with pointer types removed + // + + typedef T value_type; + typedef size_t size_type; + typedef ptrdiff_t difference_type; + typedef value_type *pointer; + typedef value_type const *const_pointer; + + // + // References + // + + /// Reference object inserts or extracts sub-byte items + class reference { + /// Pointer to storage element + Storage *ptr_; + + /// Index into elements packed into Storage object + int idx_; + + public: + + /// Default ctor + MCTLASS_HOST_DEVICE + reference(): ptr_(nullptr), idx_(0) { } + + /// Ctor + MCTLASS_HOST_DEVICE + reference(Storage *ptr, int idx = 0): ptr_(ptr), idx_(idx) { } + + /// Assignment + MCTLASS_HOST_DEVICE + reference &operator=(T x) { + Storage item = (reinterpret_cast(x) & kMask); + + Storage kUpdateMask = Storage(~(kMask << (idx_ * sizeof_bits::value))); + *ptr_ = Storage(((*ptr_ & kUpdateMask) | (item << idx_ * sizeof_bits::value))); + + return *this; + } + + MCTLASS_HOST_DEVICE + T get() const { + Storage item = Storage((*ptr_ >> (idx_ * sizeof_bits::value)) & kMask); + return reinterpret_cast(item); + } + + /// Extract + MCTLASS_HOST_DEVICE + operator T() const { + return get(); + } + + /// Explicit cast to int + MCTLASS_HOST_DEVICE + explicit operator int() const { + return int(get()); + } + + /// Explicit cast to float + MCTLASS_HOST_DEVICE + explicit operator float() const { + return float(get()); + } + }; + + /// Reference object extracts sub-byte items + class const_reference { + + /// Pointer to storage element + Storage const *ptr_; + + /// Index into elements packed into Storage object + int idx_; + + public: + + /// Default ctor + MCTLASS_HOST_DEVICE + const_reference(): ptr_(nullptr), idx_(0) { } + + /// Ctor + MCTLASS_HOST_DEVICE + const_reference(Storage const *ptr, int idx = 0): ptr_(ptr), idx_(idx) { } + + MCTLASS_HOST_DEVICE + const T get() const { + Storage item = (*ptr_ >> (idx_ * sizeof_bits::value)) & kMask; + return reinterpret_cast(item); + } + + /// Extract + MCTLASS_HOST_DEVICE + operator T() const { + Storage item = Storage(Storage(*ptr_ >> Storage(idx_ * sizeof_bits::value)) & kMask); + return reinterpret_cast(item); + } + + /// Explicit cast to int + MCTLASS_HOST_DEVICE + explicit operator int() const { + return int(get()); + } + + /// Explicit cast to float + MCTLASS_HOST_DEVICE + explicit operator float() const { + return float(get()); + } + }; + + // + // Iterators + // + + /// Bidirectional iterator over elements + class iterator { + + /// Pointer to storage element + Storage *ptr_; + + /// Index into elements packed into Storage object + int idx_; + + public: + + MCTLASS_HOST_DEVICE + iterator(): ptr_(nullptr), idx_(0) { } + + MCTLASS_HOST_DEVICE + iterator(Storage *ptr, int idx = 0): ptr_(ptr), idx_(idx) { } + + MCTLASS_HOST_DEVICE + iterator &operator++() { + ++idx_; + if (idx_ == kElementsPerStoredItem) { + ++ptr_; + idx_ = 0; + } + return *this; + } + + MCTLASS_HOST_DEVICE + iterator &operator--() { + if (!idx_) { + --ptr_; + idx_ = kElementsPerStoredItem - 1; + } + else { + --idx_; + } + return *this; + } + + MCTLASS_HOST_DEVICE + iterator operator++(int) { + iterator ret(*this); + ++idx_; + if (idx_ == kElementsPerStoredItem) { + ++ptr_; + idx_ = 0; + } + return ret; + } + + MCTLASS_HOST_DEVICE + iterator operator--(int) { + iterator ret(*this); + if (!idx_) { + --ptr_; + idx_ = kElementsPerStoredItem - 1; + } + else { + --idx_; + } + return ret; + } + + MCTLASS_HOST_DEVICE + reference operator*() const { + return reference(ptr_, idx_); + } + + MCTLASS_HOST_DEVICE + bool operator==(iterator const &other) const { + return ptr_ == other.ptr_ && idx_ == other.idx_; + } + + MCTLASS_HOST_DEVICE + bool operator!=(iterator const &other) const { + return !(*this == other); + } + }; + + /// Bidirectional constant iterator over elements + class const_iterator { + + /// Pointer to storage element + Storage const *ptr_; + + /// Index into elements packed into Storage object + int idx_; + + public: + + MCTLASS_HOST_DEVICE + const_iterator(): ptr_(nullptr), idx_(0) { } + + MCTLASS_HOST_DEVICE + const_iterator(Storage const *ptr, int idx = 0): ptr_(ptr), idx_(idx) { } + + MCTLASS_HOST_DEVICE + iterator &operator++() { + ++idx_; + if (idx_ == kElementsPerStoredItem) { + ++ptr_; + idx_ = 0; + } + return *this; + } + + MCTLASS_HOST_DEVICE + iterator &operator--() { + if (!idx_) { + --ptr_; + idx_ = kElementsPerStoredItem - 1; + } + else { + --idx_; + } + return *this; + } + + MCTLASS_HOST_DEVICE + iterator operator++(int) { + iterator ret(*this); + ++idx_; + if (idx_ == kElementsPerStoredItem) { + ++ptr_; + idx_ = 0; + } + return ret; + } + + MCTLASS_HOST_DEVICE + iterator operator--(int) { + iterator ret(*this); + if (!idx_) { + --ptr_; + idx_ = kElementsPerStoredItem - 1; + } + else { + --idx_; + } + return ret; + } + + MCTLASS_HOST_DEVICE + const_reference operator*() const { + return const_reference(ptr_, idx_); + } + + MCTLASS_HOST_DEVICE + bool operator==(iterator const &other) const { + return ptr_ == other.ptr_ && idx_ == other.idx_; + } + + MCTLASS_HOST_DEVICE + bool operator!=(iterator const &other) const { + return !(*this == other); + } + }; + + /// Bidirectional iterator over elements + class reverse_iterator { + + /// Pointer to storage element + Storage *ptr_; + + /// Index into elements packed into Storage object + int idx_; + + public: + + MCTLASS_HOST_DEVICE + reverse_iterator(): ptr_(nullptr), idx_(0) { } + + MCTLASS_HOST_DEVICE + reverse_iterator(Storage *ptr, int idx = 0): ptr_(ptr), idx_(idx) { } + }; + + /// Bidirectional constant iterator over elements + class const_reverse_iterator { + + /// Pointer to storage element + Storage const *ptr_; + + /// Index into elements packed into Storage object + int idx_; + + public: + + MCTLASS_HOST_DEVICE + const_reverse_iterator(): ptr_(nullptr), idx_(0) { } + + MCTLASS_HOST_DEVICE + const_reverse_iterator(Storage const *ptr, int idx = 0): ptr_(ptr), idx_(idx) { } + }; + +private: + + /// Internal storage + Storage storage[kStorageElements]; + +public: + + #if 0 + MCTLASS_HOST_DEVICE + Array() { } + + MCTLASS_HOST_DEVICE + Array(Array const &x) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < int(kStorageElements); ++i) { + storage[i] = x.storage[i]; + } + } + #endif + + /// Efficient clear method + MCTLASS_HOST_DEVICE + void clear() { + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < int(kStorageElements); ++i) { + storage[i] = Storage(0); + } + } + + MCTLASS_HOST_DEVICE + reference at(size_type pos) { + return reference(storage + pos / kElementsPerStoredItem, pos % kElementsPerStoredItem); + } + + MCTLASS_HOST_DEVICE + const_reference at(size_type pos) const { + return const_reference(storage + pos / kElementsPerStoredItem, pos % kElementsPerStoredItem); + } + + MCTLASS_HOST_DEVICE + reference operator[](size_type pos) { + return at(pos); + } + + MCTLASS_HOST_DEVICE + const_reference operator[](size_type pos) const { + return at(pos); + } + + MCTLASS_HOST_DEVICE + reference front() { + return at(0); + } + + MCTLASS_HOST_DEVICE + const_reference front() const { + return at(0); + } + + MCTLASS_HOST_DEVICE + reference back() { + return reference(storage + kStorageElements - 1, kElementsPerStoredItem - 1); + } + + MCTLASS_HOST_DEVICE + const_reference back() const { + return const_reference(storage + kStorageElements - 1, kElementsPerStoredItem - 1); + } + + MCTLASS_HOST_DEVICE + pointer data() { + return reinterpret_cast(storage); + } + + MCTLASS_HOST_DEVICE + const_pointer data() const { + return reinterpret_cast(storage); + } + + MCTLASS_HOST_DEVICE + Storage * raw_data() { + return storage; + } + + MCTLASS_HOST_DEVICE + Storage const * raw_data() const { + return storage; + } + + + MCTLASS_HOST_DEVICE + constexpr bool empty() const { + return !kElements; + } + + MCTLASS_HOST_DEVICE + constexpr size_type size() const { + return kElements; + } + + MCTLASS_HOST_DEVICE + constexpr size_type max_size() const { + return kElements; + } + + MCTLASS_HOST_DEVICE + void fill(T const &value) { + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kElementsPerStoredItem; ++i) { + reference ref(storage, i); + ref = value; + } + + MCTLASS_PRAGMA_UNROLL + for (int i = 1; i < kStorageElements; ++i) { + storage[i] = storage[0]; + } + } + + MCTLASS_HOST_DEVICE + iterator begin() { + return iterator(storage); + } + + MCTLASS_HOST_DEVICE + const_iterator cbegin() const { + return const_iterator(storage); + } + + MCTLASS_HOST_DEVICE + iterator end() { + return iterator(storage + kStorageElements); + } + + MCTLASS_HOST_DEVICE + const_iterator cend() const { + return const_iterator(storage + kStorageElements); + } + + MCTLASS_HOST_DEVICE + reverse_iterator rbegin() { + return reverse_iterator(storage + kStorageElements); + } + + MCTLASS_HOST_DEVICE + const_reverse_iterator crbegin() const { + return const_reverse_iterator(storage + kStorageElements); + } + + MCTLASS_HOST_DEVICE + reverse_iterator rend() { + return reverse_iterator(storage); + } + + MCTLASS_HOST_DEVICE + const_reverse_iterator crend() const { + return const_reverse_iterator(storage); + } + + // + // Comparison operators + // + +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/barrier.h b/csrc/mctlass/include/mctlass/barrier.h new file mode 100644 index 0000000..011d885 --- /dev/null +++ b/csrc/mctlass/include/mctlass/barrier.h @@ -0,0 +1,185 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Implementation of a CTA-wide barrier for inter-CTA synchronization. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// CTA-wide semaphore for inter-CTA synchronization. +struct Barrier +{ + +public: + + /// Flag type + using T = int; + + /// Initial flag value + static const T INIT = 0; + + +protected: + + /// Load flag, as a strong acquire operation (int specialization) + MCTLASS_DEVICE + static int ld_acquire(int *ptr) + { + int state = 0; + +#if defined(__MACA_ARCH__) + /// SM70 and newer use memory consistency qualifiers + + // Acquire pattern using acquire modifier + // asm volatile ("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(ptr)); + +#else + // asm volatile ("ld.cg.global.b32 %0, [%1];\n" : "=r"(state) : "l"(ptr)); +#endif // (__MACA_ARCH__) + + return state; + } + + + /// Reduce into flag, with release pattern (int specialization) + MCTLASS_DEVICE + static void red_release(int *ptr, int val) + { +#if defined(__MACA_ARCH__) + /// SM70 and newer use memory consistency qualifiers + + // Release pattern using acq_rel fence + relaxed modifier. (The fence also releases data + // that was weakly-written by other threads prior to the last syncthreads) + // asm volatile ("fence.acq_rel.gpu;\n"); + // asm volatile ("red.relaxed.gpu.global.add.s32 [%0], %1;\n" : : "l"(ptr), "r"(val)); + +#else + __threadfence(); + atomicAdd(ptr, val); +#endif // (__MACA_ARCH__) + } + + +public: + + /// Uses thread[0] to wait for at least the specified count of signals on the given flag counter + MCTLASS_DEVICE + static void wait_lt(void *lock_ptr, int thread_idx, int flag_idx, int count) + { + T *flag_ptr = reinterpret_cast(lock_ptr) + flag_idx; + + if (thread_idx == 0) + { + // Spin-loop + #pragma unroll 1 + while(ld_acquire(flag_ptr) < count) {} + } + + __syncthreads(); + } + + /// Uses thread[0] to wait for at least the specified count of signals on the given flag counter + MCTLASS_DEVICE + static void wait_eq(void *lock_ptr, int thread_idx, int flag_idx, T val = 1) + { + T *flag_ptr = reinterpret_cast(lock_ptr) + flag_idx; + + if (thread_idx == 0) + { + // Spin-loop + #pragma unroll 1 + while(ld_acquire(flag_ptr) != val) {} + } + __syncthreads(); + } + + /// Uses thread[0] to wait for the specified count of signals on the given flag counter + MCTLASS_DEVICE + static void wait_eq_reset(void *lock_ptr, int thread_idx, int flag_idx, T val = 1) { + T *flag_ptr = reinterpret_cast(lock_ptr) + flag_idx; + + if (thread_idx == 0) + { + // Spin-loop + #pragma unroll 1 + while(atomicCAS(flag_ptr, val, 0) != val) {} + } + + __syncthreads(); + } + + /// Increment the arrival count for a flag + MCTLASS_DEVICE + static void arrive_inc(void *lock_ptr, int thread_idx, int flag_idx) + { + T* flag_ptr = reinterpret_cast(lock_ptr) + flag_idx; + + __syncthreads(); + + if (thread_idx == 0) + { + red_release(flag_ptr, 1); + } + } + + + /// Increment the arrival counts for a range of flags + MCTLASS_DEVICE + static void arrive_range_inc(void *lock_ptr, int thread_idx, int first_flag_idx, int count = 1) + { + int flag_idx = first_flag_idx + thread_idx; + T* flag_ptr = reinterpret_cast(lock_ptr) + flag_idx; + + // Barrier to make sure all other threads in block have written their data + __syncthreads(); + + // Select threads increment their flags + if (thread_idx < count) { + red_release(flag_ptr, 1); + } + } +}; + + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/bfloat16.h b/csrc/mctlass/include/mctlass/bfloat16.h new file mode 100644 index 0000000..c5242f0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/bfloat16.h @@ -0,0 +1,514 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Defines a proxy class for storing non-standard 16-bit floating point values with + 8 bits of exponent and 7 bit of mantissa. +*/ +#pragma once + +#if defined(__CUDACC_RTC__) +#include "mctlass/floating_point_nvrtc.h" +#else +#include +#include +#include +#include +#endif + +#include "mctlass/mctlass.h" + +namespace mctlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Floating-point type with 8 bits of exponent and 7 bits of mantissa. +struct alignas(2) bfloat16_t { + + // + // Data members + // + + /// Storage type + uint16_t storage; + + // + // Methods + // + + /// Constructs from an unsigned short + MCTLASS_HOST_DEVICE + static bfloat16_t bitcast(uint16_t x) { + bfloat16_t h; + h.storage = x; + return h; + } + + /// Default constructor + bfloat16_t() = default; + + /// Floating-point conversion - round toward nearest + MCTLASS_HOST_DEVICE + explicit bfloat16_t(float x) { + + uint32_t bits; + + #if defined(__MACA_ARCH__) + bits = reinterpret_cast(x); + #else + std::memcpy(&bits, &x, sizeof(bits)); + #endif + + #if defined(__FAST_HALF_CVT__) + // Reference from maca_bfloat16.hpp and we remove + // the checking of NaNs. + + //bits += uint32_t(0x7fff + ((bits >> 16) & 1)); + /** + * When x = 0b xxxx xxxx xxxx xxx0 1000 0000 0000 0000, + * the result of the calculation above will be 0b xxxx xxxx xxxx xxx0, + * the result of the calculation below will be 0x xxxx xxxx xxxx xxx1. + * For example, + * x = 0011 1111 1000 0000 1000 0000 0000 0000 (f32=1.0039062), + * the above result will be (bits += uint32_t(0x7fff + ((bits >> 16) & 1));) + * result = 0011 1111 1000 0000 (bf16=1.00) + * the below result will be + * result = 0011 1111 1000 0001 (bf16=1.0078125) + */ + bits += uint32_t(0x8000); + storage = uint16_t(bits >> 16); + #else + if ((bits & 0x7f800000) != 0x7f800000) { + + bool mantissa_bit = ((bits & (1 << 16)) != 0); + bool round_bit = ((bits & (1 << 15)) != 0); + bool sticky_bit = ((bits & ((1 << 15) - 1)) != 0); + + if ((round_bit && sticky_bit) || (round_bit && mantissa_bit)) { + bits += uint32_t(1 << 16); + } + } + else if (bits & ~0xff800000) { + bits = 0x7fffffff; + } + + storage = uint16_t((bits >> 16) & 0xffff); + #endif + } + + /// Floating-point conversion - round toward nearest + MCTLASS_HOST_DEVICE + explicit bfloat16_t(double x): bfloat16_t(float(x)) { + + } + + /// Integer conversion - round toward nearest + MCTLASS_HOST_DEVICE + explicit bfloat16_t(int x) { + float flt = static_cast(x); + uint32_t bits; + + #if defined(__MACA_ARCH__) + bits = reinterpret_cast(flt); + #else + std::memcpy(&bits, &flt, sizeof(bits)); + #endif + + storage = uint16_t(bits >> 16); + } + + /// Converts to float + MCTLASS_HOST_DEVICE + operator float() const { + unsigned bits = (unsigned(storage) << 16); + #if defined(__MACA_ARCH__) + return reinterpret_cast(bits); + #else + float flt; + std::memcpy(&flt, &bits, sizeof(flt)); + return flt; + #endif + } + + /// Converts to float + MCTLASS_HOST_DEVICE + explicit operator double() const { + return double(float(*this)); + } + + /// Converts to int + MCTLASS_HOST_DEVICE + explicit operator int() const { + return int(float(*this)); + } + + /// Casts to bool + MCTLASS_HOST_DEVICE + explicit operator bool() const { + return (float(*this) != 0.0f); + } + + /// Obtains raw bits + MCTLASS_HOST_DEVICE + uint16_t raw() const { + return storage; + } + /// Returns the sign bit + MCTLASS_HOST_DEVICE + bool signbit() const { + return ((raw() & 0x8000) != 0); + } + + /// Returns the biased exponent + MCTLASS_HOST_DEVICE + int exponent_biased() const { + return int((raw() >> 7) & 0x0ff); + } + + /// Returns the unbiased exponent + MCTLASS_HOST_DEVICE + int exponent() const { + return exponent_biased() - 127; + } + + /// Returns the mantissa + MCTLASS_HOST_DEVICE + int mantissa() const { + return int(raw() & 0x7f); + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +MCTLASS_HOST_DEVICE +bool signbit(mctlass::bfloat16_t const& h) { + return h.signbit(); +} + +MCTLASS_HOST_DEVICE +mctlass::bfloat16_t abs(mctlass::bfloat16_t const& h) { + return mctlass::bfloat16_t::bitcast(h.raw() & 0x7fffffff); +} + +MCTLASS_HOST_DEVICE +bool isnan(mctlass::bfloat16_t const& h) { + return (h.exponent_biased() == 0x0ff) && h.mantissa(); +} + +MCTLASS_HOST_DEVICE +bool isfinite(mctlass::bfloat16_t const& h) { + return (h.exponent_biased() != 0x0ff); +} + +MCTLASS_HOST_DEVICE +mctlass::bfloat16_t nan_bf16(const char*) { + // NVIDIA canonical NaN + return mctlass::bfloat16_t::bitcast(0x7fff); +} + +MCTLASS_HOST_DEVICE +bool isinf(mctlass::bfloat16_t const& h) { + return (h.exponent_biased() == 0x0ff) && !h.mantissa(); +} + +MCTLASS_HOST_DEVICE +bool isnormal(mctlass::bfloat16_t const& h) { + return h.exponent_biased() && h.exponent_biased() != 0x0ff; +} + +MCTLASS_HOST_DEVICE +int fpclassify(mctlass::bfloat16_t const& h) { + int exp = h.exponent_biased(); + int mantissa = h.mantissa(); + if (exp == 0x0ff) { + if (mantissa) { + return FP_NAN; + } + else { + return FP_INFINITE; + } + } + else if (!exp) { + if (mantissa) { + return FP_SUBNORMAL; + } + else { + return FP_ZERO; + } + } + return FP_NORMAL; +} + +MCTLASS_HOST_DEVICE +mctlass::bfloat16_t sqrt(mctlass::bfloat16_t const& h) { +#if defined(__MACACC_RTC__) + return mctlass::bfloat16_t(sqrtf(float(h))); +#else + return mctlass::bfloat16_t(std::sqrt(float(h))); +#endif +} + +MCTLASS_HOST_DEVICE +bfloat16_t copysign(bfloat16_t const& a, bfloat16_t const& b) { + + uint16_t a_bits; + uint16_t b_bits; + + #if defined(__MACA_ARCH__) + a_bits = reinterpret_cast(a); + b_bits = reinterpret_cast(b); + #else + std::memcpy(&a_bits, &a, sizeof(a_bits)); + std::memcpy(&b_bits, &b, sizeof(b_bits)); + #endif + + uint16_t a_mag = (a_bits & 0x7fff); + uint16_t b_sign = (b_bits & 0x8000); + uint16_t result = (a_mag | b_sign); + + return bfloat16_t::bitcast(result); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Standard Library operations and definitions +// +/////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace std { + +#if !defined(__MACACC_RTC__) +/// Numeric limits +template <> +struct numeric_limits { + static bool const is_specialized = true; + static bool const is_signed = true; + static bool const is_integer = false; + static bool const is_exact = false; + static bool const has_infinity = true; + static bool const has_quiet_NaN = true; + static bool const has_signaling_NaN = false; + static std::float_denorm_style const has_denorm = std::denorm_present; + static bool const has_denorm_loss = true; + static std::float_round_style const round_style = std::round_to_nearest; + static bool const is_iec559 = false; + static bool const is_bounded = true; + static bool const is_modulo = false; + static int const digits = 7; + + /// Least positive value + MCTLASS_HOST_DEVICE + static mctlass::bfloat16_t min() { return mctlass::bfloat16_t::bitcast(0x01); } + + /// Minimum finite value + MCTLASS_HOST_DEVICE + static mctlass::bfloat16_t lowest() { return mctlass::bfloat16_t::bitcast(0xff7f); } + + /// Maximum finite value + MCTLASS_HOST_DEVICE + static mctlass::bfloat16_t max() { return mctlass::bfloat16_t::bitcast(0x7f7f); } + + /// Returns smallest finite value + MCTLASS_HOST_DEVICE + static mctlass::bfloat16_t epsilon() { return mctlass::bfloat16_t::bitcast(0x1000); } + + /// Returns smallest finite value + MCTLASS_HOST_DEVICE + static mctlass::bfloat16_t round_error() { return mctlass::bfloat16_t(0.5f); } + + /// Returns smallest finite value + MCTLASS_HOST_DEVICE + static mctlass::bfloat16_t infinity() { return mctlass::bfloat16_t::bitcast(0x7f80); } + + /// Returns smallest finite value + MCTLASS_HOST_DEVICE + static mctlass::bfloat16_t quiet_NaN() { return mctlass::bfloat16_t::bitcast(0x7fff); } + + /// Returns smallest finite value + MCTLASS_HOST_DEVICE + static mctlass::bfloat16_t signaling_NaN() { return mctlass::bfloat16_t::bitcast(0x7fff); } + + /// Returns smallest finite value + MCTLASS_HOST_DEVICE + static mctlass::bfloat16_t denorm_min() { return mctlass::bfloat16_t::bitcast(0x1); } +}; +#endif + +} // namespace std + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Arithmetic operators +// +/////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +MCTLASS_HOST_DEVICE +bool operator==(bfloat16_t const& lhs, bfloat16_t const& rhs) { + return float(lhs) == float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator!=(bfloat16_t const& lhs, bfloat16_t const& rhs) { + return float(lhs) != float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator<(bfloat16_t const& lhs, bfloat16_t const& rhs) { + return float(lhs) < float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator<=(bfloat16_t const& lhs, bfloat16_t const& rhs) { + return float(lhs) <= float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator>(bfloat16_t const& lhs, bfloat16_t const& rhs) { + return float(lhs) > float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator>=(bfloat16_t const& lhs, bfloat16_t const& rhs) { + return float(lhs) >= float(rhs); +} + +MCTLASS_HOST_DEVICE +bfloat16_t operator+(bfloat16_t const& lhs, bfloat16_t const& rhs) { + return bfloat16_t(float(lhs) + float(rhs)); +} + +MCTLASS_HOST_DEVICE +bfloat16_t operator-(bfloat16_t const& lhs) { + return bfloat16_t(-float(lhs)); +} + +MCTLASS_HOST_DEVICE +bfloat16_t operator-(bfloat16_t const& lhs, bfloat16_t const& rhs) { + return bfloat16_t(float(lhs) - float(rhs)); +} + +MCTLASS_HOST_DEVICE +bfloat16_t operator*(bfloat16_t const& lhs, bfloat16_t const& rhs) { + return bfloat16_t(float(lhs) * float(rhs)); +} + +MCTLASS_HOST_DEVICE +bfloat16_t operator/(bfloat16_t const& lhs, bfloat16_t const& rhs) { + return bfloat16_t(float(lhs) / float(rhs)); +} + +MCTLASS_HOST_DEVICE +bfloat16_t& operator+=(bfloat16_t & lhs, bfloat16_t const& rhs) { + lhs = bfloat16_t(float(lhs) + float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +bfloat16_t& operator-=(bfloat16_t & lhs, bfloat16_t const& rhs) { + lhs = bfloat16_t(float(lhs) - float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +bfloat16_t& operator*=(bfloat16_t & lhs, bfloat16_t const& rhs) { + lhs = bfloat16_t(float(lhs) * float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +bfloat16_t& operator/=(bfloat16_t & lhs, bfloat16_t const& rhs) { + lhs = bfloat16_t(float(lhs) / float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +bfloat16_t& operator++(bfloat16_t & lhs) { + float tmp(lhs); + ++tmp; + lhs = bfloat16_t(tmp); + return lhs; +} + +MCTLASS_HOST_DEVICE +bfloat16_t& operator--(bfloat16_t & lhs) { + float tmp(lhs); + --tmp; + lhs = bfloat16_t(tmp); + return lhs; +} + +MCTLASS_HOST_DEVICE +bfloat16_t operator++(bfloat16_t & lhs, int) { + bfloat16_t ret(lhs); + float tmp(lhs); + tmp++; + lhs = bfloat16_t(tmp); + return ret; +} + +MCTLASS_HOST_DEVICE +bfloat16_t operator--(bfloat16_t & lhs, int) { + bfloat16_t ret(lhs); + float tmp(lhs); + tmp--; + lhs = bfloat16_t(tmp); + return ret; +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// +// User-defined literals +// + +MCTLASS_HOST_DEVICE +mctlass::bfloat16_t operator "" _bf16(long double x) { + return mctlass::bfloat16_t(float(x)); +} + +MCTLASS_HOST_DEVICE +mctlass::bfloat16_t operator "" _bf16(unsigned long long int x) { + return mctlass::bfloat16_t(int(x)); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/blas3.h b/csrc/mctlass/include/mctlass/blas3.h new file mode 100644 index 0000000..ebbdb3e --- /dev/null +++ b/csrc/mctlass/include/mctlass/blas3.h @@ -0,0 +1,176 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Basic include for MCTLASS BLAS3/HPC code. + + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/complex.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_types.h" + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Enumerated type describing the type of kernel (based on input or output matrices). +enum class BlasMode { + kGemm, + kSymmetric, + kHermitian, + kTriangular, + kInvalid +}; + +/// Enumerated type describing the fill mode for matrices for BLAS functions. +enum class FillMode { + kFull, /// The entire tensor is covered. + kLower, /// The 'lower' part of a tensor is covered including diagonal + kUpper, /// The 'upper' part of a tensor is covered including diaognal + kDiagonal, /// Only diagonal elements are covered. + kNone, /// No element is covered. + kInvalid +}; + +/// Enumerated type describing the diagonal property of matrices for BLAS functions. +enum class DiagType { + kNonUnit, + kUnit, + kZero, // Only used internally for computing SYMM/HEMM + kInvalid +}; + +/// Enumerated type describing the side dense matrix is in matrix equation for BLAS functions. +enum class SideMode { + kLeft, + kRight, + kInvalid +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines FillMode inversions +template +struct InvertFillMode; + +/// Invert FillMode lower to upper +template <> +struct InvertFillMode { + static FillMode const mode = FillMode::kUpper; +}; + +/// Invert FillMode upper to lower +template <> +struct InvertFillMode { + static FillMode const mode = FillMode::kLower; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines SideMode inversions +template +struct InvertSideMode; + +/// Invert SideMode left to right +template <> +struct InvertSideMode { + static SideMode const mode = SideMode::kRight; +}; + +/// Invert SideMode right to left +template <> +struct InvertSideMode { + static SideMode const mode = SideMode::kLeft; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines correct compare operation for Triangular matrix boundary +template +struct TrMatrixCompareOp { + using Index = int32_t; + using Type = typename platform::conditional< + (kFillMode == FillMode::kLower), + greater_equal, + less_equal>::type; +}; + +template +struct TrMatrixCompareOp { + using Index = int32_t; + using Type = typename platform::conditional< + (kFillMode == FillMode::kLower), + greater_equal, + less_equal>::type; +}; + +template +struct TrMatrixCompareOp { + using Index = int32_t; + using Type = typename platform::conditional< + (kFillMode == FillMode::kLower), + greater, + less>::type; +}; +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Returns precision in terms of bits (based on datatype) to fill tensors with. +// Defaults to 5 bits of mantissa for TF32 and FP32 (with implicit round-offs). +// Also defines acceptable mantissa result variance/error. +template +struct MantissaInBits { + static int constexpr bits = 5; + static double constexpr error = 1.0e-7; +}; + +// Full precision is supported for FP64 +template <> +struct MantissaInBits { + static int constexpr bits = 30; + static double constexpr error = 1.0e-15; +}; + +template <> +struct MantissaInBits> { + static int constexpr bits = 30; + static double constexpr error = 1.0e-15; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/block_striped.h b/csrc/mctlass/include/mctlass/block_striped.h new file mode 100644 index 0000000..bed3050 --- /dev/null +++ b/csrc/mctlass/include/mctlass/block_striped.h @@ -0,0 +1,267 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Utilities for performing block-striped access (load, store, reduce) of trivially-copyable, + statically-sized array types to global memory. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/wmma_array.h" +#include "mctlass/functional.h" +#include "mctlass/complex.h" + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// +// AccessWidth +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Computes the maximal power-of-two that evenly divides the size of T, capped at Limit +template < + typename T, + int Limit> +struct AccessWidth +{ + // Inductive case + template < + int ObjectBytes, /// Size of T in bytes + int AlignBytes, /// Template induction variable + bool IsAligned = /// Whether ObjectBytes is an even multiple of AlignBytes + ((AlignBytes <= Limit) && (ObjectBytes % AlignBytes == 0))> + struct Detail + { + static const int value = Detail::value; + }; + + // Base case (ObjectBytes is not an even multiple of AlignBytes) + template < + int ObjectBytes, /// Size of T in bytes + int AlignBytes> /// Template induction variable + struct Detail + { + static const int value = AlignBytes / 2; + }; + + /// The maximal power-of-two that evenly divides the size of T + static const int value = Detail< + (int) sizeof(T), + 1>::value; +}; + + + +///////////////////////////////////////////////////////////////////////////////////////////////// +// StripedAccessType +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// ReinterpretCast type for striping a trivially-copyable type in global memory +/// (Default specialization. Striping granularity is type T.) +template < + typename T, /// Data type + int TransferBytes = /// Data access width (16 byte max for global memory access on current architectures) + AccessWidth::value> +struct alignas(TransferBytes) StripedAccessType : public T +{}; + + +/// ReinterpretCast type for striping a trivially-copyable type in global memory +/// (Specialization for mctlass::Array. Striping granularity is a multiple of T.) +template < + typename T, /// Array element type + int N, /// Number of elements in array + bool RegisterSized, /// T is register-sized + int TransferBytes> /// Data access width +struct StripedAccessType< + Array, + TransferBytes> +: public AlignedArray< + T, // Element type of StripedAccessType + __NV_STD_MAX(1, TransferBytes / (int) sizeof(T)), // Number of elements T in StripedAccessType + TransferBytes> // Alignment of StripedAccessType +{}; + + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) + +/// ReinterpretCast type for striping a trivially-copyable type in global memory +/// (Specialization for mctlass::WmmaFragmentArray. Striping granularity is a multiple of T.) +template< + typename Use, + int m, + int n, + int k, + typename ElementT, + typename Layout, + int kFragments, + int TransferBytes> +struct StripedAccessType< + WmmaFragmentArray, kFragments>, + TransferBytes> +: public AlignedArray< + ElementT, + __NV_STD_MAX(1, TransferBytes / (int) sizeof(ElementT)), + TransferBytes> +{}; + +#endif // if defined(MCTLASS_ARCH_WMMA_ENABLED) + + +///////////////////////////////////////////////////////////////////////////////////////////////// +// BlockStriped +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Utility for performing block-striped access (load, store) of trivially-copyable, +/// statically-sized array types to global memory +template < + int BlockThreads, + typename ArrayT, + typename AccessT = StripedAccessType > +struct BlockStriped +{ + /// Number of striped accesses + static const int kStripes = int(sizeof(ArrayT) / sizeof(AccessT)); + static_assert(kStripes > 0, "AccessT type must be smaller than or equal to ArrayT type"); + + /// Load + MCTLASS_DEVICE + static void load(ArrayT &data, ArrayT *ptr, int thread_idx) + { + AccessT *access_input = reinterpret_cast(ptr); + AccessT *access_data = reinterpret_cast(&data); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kStripes; ++i) { + access_data[i] = access_input[(BlockThreads * i) + thread_idx]; + } + } + + /// Load & Add + MCTLASS_DEVICE + static void load_add(ArrayT &data, ArrayT *ptr, int thread_idx) + { + AccessT *access_input = reinterpret_cast(ptr); + AccessT *access_data = reinterpret_cast(&data); + + plus add; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kStripes; ++i) + { + access_data[i] = add(access_data[i], access_input[(BlockThreads * i) + thread_idx]); + } + } + + /// Store + MCTLASS_DEVICE + static void store(ArrayT *ptr, const ArrayT &data, int thread_idx) + { + AccessT *access_output = reinterpret_cast(ptr); + const AccessT *access_data = reinterpret_cast(&data); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kStripes; ++i) { + access_output[(BlockThreads * i) + thread_idx] = access_data[i]; + } + } + +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// +// BlockStripedReduce +///////////////////////////////////////////////////////////////////////////////////////////////// + + +/// Utility for performing block-striped access (load, store, reduce) of trivially-copyable, +/// statically-sized array types to global memory. +/// (Default specialization) +template < + int BlockThreads, + typename ArrayT, + typename ElementT = typename StripedAccessType::Element> +struct BlockStripedReduce : + BlockStriped< + BlockThreads, + ArrayT, + ElementT> +{ + /// Reduce + MCTLASS_DEVICE + static void reduce(ArrayT *ptr, const ArrayT &data, int thread_idx) + { + mctlass::red reduce; + ElementT *access_output = reinterpret_cast(ptr); + const ElementT *access_data = reinterpret_cast(&data); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < BlockStripedReduce::kStripes; ++i) { + reduce(access_output + (BlockThreads * i) + thread_idx, access_data[i]); + } + } +}; + + +/// Utility for performing block-striped access (load, store, reduce) of trivially-copyable, +/// statically-sized array types to global memory. +/// (Specialization for half_t. Uses half2 vectorized-reduction.) +template < + int BlockThreads, + typename ArrayT> +struct BlockStripedReduce : + BlockStriped< + BlockThreads, + ArrayT, + half2> +{ + static_assert(BlockStripedReduce::kStripes % 2 == 0, "Array of half must be even number in length"); + + /// Reduce + MCTLASS_DEVICE + static void reduce(ArrayT *ptr, const ArrayT &data, int thread_idx) + { + mctlass::red reduce; + half2 *access_output = reinterpret_cast(ptr); + const half2 *access_data = reinterpret_cast(&data); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < BlockStripedReduce::kStripes; ++i) + { + reduce(access_output + (BlockThreads * i) + thread_idx, access_data[i]); + } + } +}; + + +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/cluster_launch.hpp b/csrc/mctlass/include/mctlass/cluster_launch.hpp new file mode 100644 index 0000000..347e5b2 --- /dev/null +++ b/csrc/mctlass/include/mctlass/cluster_launch.hpp @@ -0,0 +1,235 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief PTX for TMA Tensor Memory Access operators on memory added for SM90 +*/ + +#pragma once + +#include +#include "mctlass/mctlass.h" +#include "mctlass/trace.h" + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#if 0 +# define MCTLASS_SM90_CLUSTER_LAUNCH_ENABLED +#endif + +namespace mctlass { + +#ifndef NDEBUG +#define Return_Status(mcError_t_status) \ + if (mcError_t_status != mcSuccess) { \ + fprintf(stderr, \ + "[ ERROR: CUDA Runtime ] %s:%d: %s\n", \ + __FILE__, \ + __LINE__, \ + mcGetErrorString(mcError_t_status)); \ + return Status::kInvalid; \ + } else { \ + return Status::kSuccess; \ + } +#else +#define Return_Status(mcError_t_status) \ + if (mcError_t_status != mcSuccess) { \ + return Status::kInvalid; \ + } else { \ + return Status::kSuccess; \ + } +#endif + +struct ClusterLauncher { + constexpr static int MaxClusterSize = 32; + + // Check for hardware compatibility + static inline __host__ + Status check_cluster_dims(dim3 grid, dim3 cluster) { + if (((cluster.x * cluster.y * cluster.z) <= MaxClusterSize) && + (grid.x % cluster.x == 0) && (grid.y % cluster.y == 0) && (grid.z % cluster.z == 0)) { + return Status::kSuccess; + } + else { + MCTLASS_TRACE_HOST("ClusterLauncher: Invalid cluster configuration -- aborting launch."); + return Status::kInvalid; + } + } + + static inline __host__ + Status +#if defined(MCTLASS_SM90_CLUSTER_LAUNCH_ENABLED) + init(void const* kernel_function) +#else + init(void const* /* kernel_function */) +#endif + { +#if defined(MCTLASS_SM90_CLUSTER_LAUNCH_ENABLED) + // This attribute was added in CUDA 11.8. + mcError_t status = + mcFuncSetAttribute( + kernel_function, mcFuncAttributeNonPortableClusterSizeAllowed, 1); + Return_Status(status); +#else + return Status::kInvalid; +#endif + } + + // This is the method we expect to use going forward + static inline __host__ + Status launch( + dim3 const grid_dims, + dim3 const cluster_dims, + dim3 const block_dims, + size_t const smem_size, + mcStream_t mc_stream, + void const* kernel, + void** kernel_params) { +#if defined(MCTLASS_SM90_CLUSTER_LAUNCH_ENABLED) + if (check_cluster_dims(grid_dims, cluster_dims) != Status::kSuccess) { + MCTLASS_TRACE_HOST("ClusterLauncher: check_cluster_dims() failed. Aborting."); + return Status::kInvalid; + } + + auto init_status = init(kernel); + if (init_status != Status::kSuccess) { + MCTLASS_TRACE_HOST("ClusterLauncher: init(kernel) failed with status " << int(init_status) << ". Aborting."); + return Status::kInvalid; + } + + mcLaunchConfig_t launch_config; + launch_config.gridDim = {grid_dims.x, grid_dims.y, grid_dims.z}; + launch_config.blockDim = {block_dims.x, block_dims.y, block_dims.z}; + launch_config.dynamicSmemBytes = smem_size; + launch_config.stream = mc_stream; + + mcLaunchAttribute launch_attribute[1]; + launch_attribute[0].id = mcLaunchAttributeClusterDimension; + launch_attribute[0].val.clusterDim.x = cluster_dims.x; + launch_attribute[0].val.clusterDim.y = cluster_dims.y; + launch_attribute[0].val.clusterDim.z = cluster_dims.z; + + launch_config.attrs = launch_attribute; + launch_config.numAttrs = 1; + + MCTLASS_TRACE_HOST("ClusterLauncher: Launching GPC_CLUSTER_GRID GridDims = " + "(" << grid_dims.x << ", " << grid_dims.y << ", " << grid_dims.z << "), " + "And ClusterDims = " + "(" << cluster_dims.x << ", " << cluster_dims.y << ", " << cluster_dims.z << ")\n"); + + mcError_t status = mcLaunchKernelExC(&launch_config, kernel, kernel_params); + Return_Status(status); +#else + MCTLASS_TRACE_HOST("ClusterLauncher: MCTLASS_SM90_CLUSTER_LAUNCH_ENABLED not defined! Aborting cluster launch."); + return Status::kInvalid; +#endif + } +}; + +namespace detail { + +template +void* checked_addressof(Arg&& arg) { + static_assert(! std::is_rvalue_reference_v || ! std::is_const_v, "You cannot take the address of a const rvalue reference (const T&&)."); + // We use std::addressof to ensure we get the address, + // in case the type has an overloaded operator&. + // Note that this precludes `const T&&` references. + return const_cast(reinterpret_cast(std::addressof(arg))); +} + +} // namespace detail + +//! Parameters for launch_on_cluster (see below). +struct ClusterLaunchParams { + //! Grid dimensions + dim3 grid_dims{1, 1, 1}; + + //! Block dimensions + dim3 block_dims{1, 1, 1}; + + //! Cluster dimensions + dim3 cluster_dims{1, 1, 1}; + + //! Number of bytes required for the kernel's shared memory. + int smem_size_in_bytes = 0; + + //! CUDA stream on which to launch the kernel. + mcStream_t mc_stream = nullptr; +}; + +/// @brief Launch the kernel on the stream using cluster launch. +/// +/// @param params Cluster launch parameters (see above). +/// @param kernel_ptr Pointer to the kernel function (see example). +/// @param args Zero or more arguments to pass to the kernel. +/// +/// @tparam Args Types of the arguments passed to the kernel. +/// Don't specify this/these template argument(s) explicitly. +/// +/// @return Status::Success on success, else an error code. +/// +/// @code +/// template +/// __global__ void kernel(A a, B b, C c); +/// +/// X x = get_x(); +/// Y y = get_y(); +/// Z z = get_z(); +/// +/// void const* kernel_ptr = +/// const_cast(reinterpret_cast( +/// &kernel)); +/// auto status = launch_on_cluster( +/// {grid_dims, block_dims, cluster_dims, sizeof(SharedMemory)}, +/// kernel_ptr, x, y, z); +/// @endcode +template +__host__ mctlass::Status +launch_kernel_on_cluster(const ClusterLaunchParams& params, + void const* kernel_ptr, + Args&& ... args) +{ + // Unfortunately, we find ourselves needing to pass in + // the parameters as an array of raw pointers. + void* kernel_params[] = { + detail::checked_addressof(std::forward(args))... + }; + return mctlass::ClusterLauncher::launch( + params.grid_dims, params.cluster_dims, params.block_dims, + params.smem_size_in_bytes, params.mc_stream, + kernel_ptr, kernel_params); +} + +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/complex.h b/csrc/mctlass/include/mctlass/complex.h new file mode 100644 index 0000000..1c4d18c --- /dev/null +++ b/csrc/mctlass/include/mctlass/complex.h @@ -0,0 +1,693 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include +#include + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/functional.h" +#include "mctlass/half.h" +#include "mctlass/real.h" + +#include "mctlass/bfloat16.h" +#include "mctlass/tfloat32.h" + +#include "mctlass/fast_math.h" + +#if !defined(__MACACC_RTC__) +#include +#endif + +namespace mctlass { + + + + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Enumeraed type describing a transformation on a complex value. +enum class ComplexTransform { + kNone, + kConjugate +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines ComplexTransform inversions +template +struct InvertComplexTransform; + +/// Invert ComplexTransform from kNone to kConjugate +template <> +struct InvertComplexTransform { + static ComplexTransform const transform = ComplexTransform::kConjugate; +}; + +/// Invert ComplexTransform from kConjugate to kNone +template <> +struct InvertComplexTransform { + static ComplexTransform const transform = ComplexTransform::kNone; +}; +///////////////////////////////////////////////////////////////////////////////////////////////// +////////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Accessors for CUDA complex types +// + +#if !defined(__MACACC_RTC__) +/// Returns the real part of the complex number +MCTLASS_HOST_DEVICE +float const &real(mcFloatComplex const &z) { return z.x; } + +/// Returns the real part of the complex number +MCTLASS_HOST_DEVICE +float &real(mcFloatComplex &z) { return z.x; } + +/// Returns the real part of the complex number +MCTLASS_HOST_DEVICE +double const &real(mcDoubleComplex const &z) { return z.x; } + +/// Returns the real part of the complex number +MCTLASS_HOST_DEVICE +double &real(mcDoubleComplex &z) { return z.x; } + +/// Returns the imaginary part of the complex number +MCTLASS_HOST_DEVICE +float const &imag(mcFloatComplex const &z) { return z.y; } + +/// Returns the imaginary part of the complex number +MCTLASS_HOST_DEVICE +float &imag(mcFloatComplex &z) { return z.y; } + +/// Returns the imaginary part of the complex number +MCTLASS_HOST_DEVICE +double const &imag(mcDoubleComplex const &z) { return z.y; } + +/// Returns the imaginary part of the complex number +MCTLASS_HOST_DEVICE +double &imag(mcDoubleComplex &z) { return z.y; } +#endif + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Class for representing and manipulating complex numbers with conversions from built-in CUDA +/// complex types. + +template +class complex +{ + public: + /// Type alias for scalar type + using value_type = T; + + private: + // + // Data members + // + + /// Real part + T _real; + + /// Imaginary part + T _imag; + + public: + +// +// Methods +// + + /// Default constructor + complex() = default; + + /// Constructor + MCTLASS_HOST_DEVICE + complex(T r) : _real(r), _imag(T(0)) {} + + /// Constructor + MCTLASS_HOST_DEVICE + complex(T r, T i) : _real(r), _imag(i) {} + + /// Constructor + template + MCTLASS_HOST_DEVICE + complex(complex const &z) : _real(static_cast(z.real())), _imag(static_cast(z.imag())) {} + + + #if !defined(__MACACC_RTC__) + /// Conversion from mcFloatComplex + MCTLASS_HOST_DEVICE + complex(mcFloatComplex const &z) : _real(static_cast(mcCrealf(z))), _imag(static_cast(mcCimagf(z))) {} + + /// Conversion from mcDoubleComplex + MCTLASS_HOST_DEVICE + complex(mcDoubleComplex const &z) : _real(static_cast(mcCreal(z))), _imag(static_cast(mcCimag(z))) {} + #endif + + /// Equality operator + MCTLASS_HOST_DEVICE bool operator==(complex const &rhs) const { + return this->real() == rhs.real() && this->imag() == rhs.imag(); + } + + /// Inequality operator + MCTLASS_HOST_DEVICE bool operator!=(complex const &rhs) const { + return !(*this == rhs); + } + + /// Addition + template + MCTLASS_HOST_DEVICE complex operator+(complex const &rhs) const { + return complex(this->real() + rhs.real(), this->imag() + rhs.imag()); + } + + /// Reduction into memory address. Components may update out of order. + template + MCTLASS_DEVICE void red(complex *ptr) const { + static_assert(platform::is_same::value, "Component type must match"); + mctlass::red reduce; + reduce(&ptr->_real, _real); + reduce(&ptr->_imag, _imag); + } + + /// Reduction into memory address. Components may update out of order. (Half specialization) + MCTLASS_DEVICE void red(complex *ptr) const { + static_assert(platform::is_same::value, "Component type must match"); + half2 *h2_ptr = reinterpret_cast(ptr); + half2 h2_data = reinterpret_cast(*this); + mctlass::red reduce; + reduce(h2_ptr, h2_data); + } + + /// Subtraction + template + MCTLASS_HOST_DEVICE complex operator-(complex const &rhs) const { + return complex(this->real() - rhs.real(), this->imag() - rhs.imag()); + } + + /// Multiplication + template + MCTLASS_HOST_DEVICE complex operator*(complex const &rhs) const { + return complex(this->real() * rhs.real() - this->imag() * rhs.imag(), + this->real() * rhs.imag() + this->imag() * rhs.real()); + } + + /// Scalar Multiplication + template + MCTLASS_HOST_DEVICE complex operator*(A const &s) const { + return complex(this->real() * s, this->imag() * s); + } + + /// Division + template + MCTLASS_HOST_DEVICE complex operator/(complex const &rhs) const { + T d = T(rhs.real() * rhs.real() + rhs.imag() * rhs.imag()); + + return complex( + (real() * rhs.real() + imag() * rhs.imag()) / d, + (imag() * rhs.real() - real() * rhs.imag()) / d + ); + } + + /// Scalar Division + template + MCTLASS_HOST_DEVICE complex operator/(A const &s) const { + return complex(this->real() / s, this->imag() / s); + } + + /// Addition + template + MCTLASS_HOST_DEVICE complex &operator+=(complex const &rhs) { + *this = *this + rhs; + return *this; + } + + /// Subtraction + template + MCTLASS_HOST_DEVICE complex &operator-=(complex const &rhs) { + *this = *this - rhs; + return *this; + } + + /// Multiplication + template + MCTLASS_HOST_DEVICE complex &operator*=(complex const &rhs) { + *this = *this * rhs; + return *this; + } + + /// Scalar multiplication + template + MCTLASS_HOST_DEVICE complex &operator*=(A s) { + *this = *this * s; + return *this; + } + + /// Division + template + MCTLASS_HOST_DEVICE complex &operator/=(complex const &rhs) { + *this = *this / rhs; + return *this; + } + + /// Accesses the real part of the complex number + MCTLASS_HOST_DEVICE + T const &real() const { return _real; } + + /// Accesses the real part of the complex number + MCTLASS_HOST_DEVICE + T &real() { return _real; } + + /// Accesses the imaginary part of the complex number + MCTLASS_HOST_DEVICE + T const &imag() const { return _imag; } + + /// Accesses the imaginary part of the complex number + MCTLASS_HOST_DEVICE + T &imag() { return _imag; } + + + #if !defined(__MACACC_RTC__) + /// Converts to mcFloatComplex + MCTLASS_HOST_DEVICE + explicit operator mcFloatComplex() const { return make_mcFloatComplex(float(real()), float(imag())); } + + /// Converts to mcDoubleComplex + MCTLASS_HOST_DEVICE + explicit operator mcDoubleComplex() const { return make_mcDoubleComplex(real(), imag()); } + #endif +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Accessors for complex template +// + +/// Returns the real part of the complex number +template +MCTLASS_HOST_DEVICE T const &real(complex const &z) { + return z.real(); +} + +/// Returns the real part of the complex number +template +MCTLASS_HOST_DEVICE T &real(complex &z) { + return z.real(); +} + +/// Returns the imaginary part of the complex number +template +MCTLASS_HOST_DEVICE T const &imag(complex const &z) { + return z.imag(); +} + +/// Returns the imaginary part of the complex number +template +MCTLASS_HOST_DEVICE T &imag(complex &z) { + return z.imag(); +} + +/// Returns the real part of the real number +template +MCTLASS_HOST_DEVICE T const &real(T const &r) { + return r; +} + +/// Returns the real part of the real number +template +MCTLASS_HOST_DEVICE T &real(T &r) { + return r; +} + +/// Returns the imaginary part of the real number +template +MCTLASS_HOST_DEVICE T const &imag(T const &r) { + return T(); +} + +/// Returns the imaginary part of the complex number +template +MCTLASS_HOST_DEVICE T &imag(T &r) { + return T(); +} + +// +// Output operators +// + +#if !defined(__MACACC_RTC__) +template +std::ostream &operator<<(std::ostream &out, complex const &z) { + T _r = real(z); + T _i = imag(z); + + if (bool(_i)) { + return out << _r << "+i" << _i; + } + return out << _r; +} +#endif + +// +// Non-member operators defined for complex types +// + + +// +// Non-member functions defined for complex numbers +// + +/// Returns the magnitude of the complex number +template +MCTLASS_HOST_DEVICE T abs(complex const &z) { + return sqrt(norm(z)); +} + +/// Returns the magnitude of the complex number +template +MCTLASS_HOST_DEVICE T arg(complex const &z) { + return atan2(imag(z), real(z)); +} + +/// Returns the squared magnitude of a real number +template +MCTLASS_HOST_DEVICE T norm(T const &z) { + return z * z; +} + +/// Returns the squared magnitude of a real number +template <> +MCTLASS_HOST_DEVICE int8_t norm(int8_t const &z) { + return static_cast(z * z); +} + +/// Returns the squared magnitude of a complex number +template +MCTLASS_HOST_DEVICE double norm(complex const &z) { + return real(z) * real(z) + imag(z) * imag(z); +} + +/// Norm-acmcmulate calmclation +template +MCTLASS_HOST_DEVICE R norm_acmcmulate(T const &x, R const & acmcmulator) { + return acmcmulator + static_cast(x) * static_cast(x); +} + +/// Norm acmcmulate specialized for complex types +template +MCTLASS_HOST_DEVICE R norm_acmcmulate(complex const &z, R const &acmcmulator) { + return acmcmulator + static_cast(real(z)) * static_cast(real(z)) + + static_cast(imag(z)) * static_cast(imag(z)); +} + +/// Returns the complex conjugate +MCTLASS_HOST_DEVICE float conj(float const &z) { + return z; +} + +/// Returns the complex conjugate +MCTLASS_HOST_DEVICE double conj(double const &z) { + return z; +} + +/// Returns the complex conjugate +template +MCTLASS_HOST_DEVICE complex conj(complex const &z) { + return complex(real(z), -imag(z)); +} +/// Indentity transform for non-complex types +template +MCTLASS_HOST_DEVICE T conj(T const &z) { + static_assert( !platform::is_same::value && + !platform::is_same::value && + !platform::is_same>::value && + !platform::is_same>::value, "May not be a complex data type"); + return z; +} + +/// Projects the complex number z onto the Riemann sphere +template +MCTLASS_HOST_DEVICE complex proj(complex const &z) { + T d = real(z) * real(z) + imag(z) * imag(z) + T(1); + return complex((T(2) * real(z)) / d, (T(2) * imag(z)) / d); +} + +/// Returns a complex number with magnitude r and phase theta +template +MCTLASS_HOST_DEVICE complex polar(T const &r, T const &theta = T()) { + return complex(r * cos(theta), r * sin(theta)); +} + +/// Computes the complex exponential of z. +template +MCTLASS_HOST_DEVICE complex exp(complex const &z) { + return complex(fast_exp(real(z)) * fast_cos(imag(z)), fast_exp(real(z)) * fast_sin(imag(z))); +} + +/// Computes the log of z +template +MCTLASS_HOST_DEVICE complex log(complex const &z) { + return complex(log(abs(z)), arg(z)); +} + +/// Computes the log base 10 of z +template +MCTLASS_HOST_DEVICE complex log10(complex const &z) { + return log(z) / T(log(T(10))); +} + +/// Computes the square root of complex number z +template +MCTLASS_HOST_DEVICE complex sqrt(complex const &z) { + return sqrt(T(2)) / T(2) * + complex(sqrt(sqrt(norm(z)) + real(z)), + (imag(z) < 0 ? T(-1) : T(1)) * sqrt(sqrt(norm(z)) - real(z))); +} + +/// Computes the cosine of complex z. +template +MCTLASS_HOST_DEVICE complex cos(complex const &z) { + return (exp(z) + exp(-z)) / T(2); +} + +/// Computes the sin of complex z. +template +MCTLASS_HOST_DEVICE complex sin(complex const &z) { + return (exp(-z) - exp(z)) * complex(T(0), T(1) / T(2)); +} + +/// Comparison +template +MCTLASS_HOST_DEVICE bool operator<(complex const &lhs, complex const &rhs) { + return true; +} + +////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex-valued type. +template +struct RealType< complex > +{ + using Type = T; + + /// Number of elements + static int const kExtent = 2; + + MCTLASS_HOST_DEVICE + static complex from_real(double x) { + return complex(static_cast(x)); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +MCTLASS_HOST_DEVICE +mctlass::complex from_real >(double r) { + return mctlass::complex(half_t(r)); +} + +template <> +MCTLASS_HOST_DEVICE +mctlass::complex from_real >(double r) { + return mctlass::complex(float(r)); +} + +template <> +MCTLASS_HOST_DEVICE +mctlass::complex from_real >(double r) { + return mctlass::complex(r); +} + +////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct is_complex { + static bool const value = false; +}; + +template +struct is_complex> { + static bool const value = true; +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// +// functional.h numeric specializations +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Squares with optional conversion +template +struct magnitude_squared, Output> { + MCTLASS_HOST_DEVICE + Output operator()(complex lhs) const { + multiplies mul_op; + + Output y_r = Output(lhs.real()); + Output y_i = Output(lhs.imag()); + + return mul_op(y_r, y_r) + mul_op(y_i, y_i); + } +}; + +/// Fused multiply-add +template +struct multiply_add, complex, complex> { + MCTLASS_HOST_DEVICE + complex operator()( + complex const &a, + complex const &b, + complex const &c) const { + + T real = c.real(); + T imag = c.imag(); + + real += a.real() * b.real(); + real += -a.imag() * b.imag(); + imag += a.real() * b.imag(); + imag += a.imag () * b.real(); + + return complex{ + real, + imag + }; + } +}; + +/// Fused multiply-add +template +struct multiply_add, T, complex> { + MCTLASS_HOST_DEVICE + complex operator()( + complex const &a, + T const &b, + complex const &c) const { + + T real = c.real(); + T imag = c.imag(); + + real += a.real() * b; + imag += a.imag () * b; + + return complex{ + real, + imag + }; + } +}; + +/// Fused multiply-add +template +struct multiply_add, complex> { + MCTLASS_HOST_DEVICE + complex operator()( + T const &a, + complex const &b, + complex const &c) const { + + T real = c.real(); + T imag = c.imag(); + + real += a * b.real(); + imag += a * b.imag(); + + return complex{ + real, + imag + }; + } +}; + +/// Conjugate +template +struct conjugate> { + MCTLASS_HOST_DEVICE + complex operator()(complex const &a) const { + return conj(a); + } +}; + +/// Computes the square of a difference with optional conversion +template +struct magnitude_squared_difference, Output> { + MCTLASS_HOST_DEVICE + Output operator()(complex lhs, complex rhs) const { + multiplies mul_op; + + Output y_r = Output(lhs.real()) - Output(rhs.real()); + Output y_i = Output(lhs.imag()) - Output(rhs.imag()); + + return mul_op(y_r, y_r) + mul_op(y_i, y_i); + } +}; + +/// Reduces value into the data pointed to by ptr (complex specialization) +template +struct red> { + MCTLASS_DEVICE + void operator()(complex *ptr, const complex &data) + { + data.red(ptr); + } +}; + + +////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/constants.h b/csrc/mctlass/include/mctlass/constants.h new file mode 100644 index 0000000..150213e --- /dev/null +++ b/csrc/mctlass/include/mctlass/constants.h @@ -0,0 +1,1239 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/* \file + \brief Boost-style constant definitions for floating-point types. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" + +#include "mctlass/complex.h" + +/////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace constants { + +/////////////////////////////////////////////////////////////////////////////////// + +// +// Primary templates +// + +/// Returns 1, the multiplicative identity element +template MCTLASS_HOST_DEVICE T one(); + +/// Returns 0, the additive identity element +template MCTLASS_HOST_DEVICE T zero(); + +/// Returns 2 +template MCTLASS_HOST_DEVICE T two(); + +/// Returns pi, approximately 3.141 +template MCTLASS_HOST_DEVICE T pi(); + +/// Returns 2 * pi +template MCTLASS_HOST_DEVICE T two_pi(); + +/// Returns pi / 2 +template MCTLASS_HOST_DEVICE T half_pi(); + +/// Returns sqrt(pi) +template MCTLASS_HOST_DEVICE T root_pi(); + +/// Returns sqrt(pi / 2) +template MCTLASS_HOST_DEVICE T root_half_pi(); + +/// Returns sqrt(2 * pi) +template MCTLASS_HOST_DEVICE T root_two_pi(); + +/// Returns sqrt(ln(4)) +template MCTLASS_HOST_DEVICE T root_ln_four(); + +/// Returns e, approximately 2.718... +template MCTLASS_HOST_DEVICE T e(); + +/// Returns (1/2) +template MCTLASS_HOST_DEVICE T half(); + +/// Returns sqrt(2), approximately 1.414... +template MCTLASS_HOST_DEVICE T root_two(); + +/// Returns sqrt(2)/2, approximately 0.707... +template MCTLASS_HOST_DEVICE T half_root_two(); + +/// Returns ln(2), approximately 0.693... +template MCTLASS_HOST_DEVICE T ln_two(); + +/// Returns ln(ln(2)), approximately -0.3665... +template MCTLASS_HOST_DEVICE T ln_ln_two(); + +/// Returns 1/3, approximately 0.333... +template MCTLASS_HOST_DEVICE T third(); + +/// Returns 2/3, approximately 0.666... +template MCTLASS_HOST_DEVICE T twothirds(); + +/// Returns pi - 3, approximately 0.1416... +template MCTLASS_HOST_DEVICE T pi_minus_three(); + +/// Returns 4 - pi, approximately 0.858... +template MCTLASS_HOST_DEVICE T four_minus_pi(); + + +///////////////////////////////////////////////////////////////////////////////////// + +// Specialization for double + +/// Returns 1, the multiplicative identity element (specialization for double) +template <> MCTLASS_HOST_DEVICE double one() { + uint64_t bits = 0x3ff0000000000000ull; + return reinterpret_cast(bits); +} + +/// Returns 1, the multiplicative identity element (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex one< complex >() { + return complex(one(), double()); +} + +/// Returns 0, the additive identity element (specialization for double) +template <> MCTLASS_HOST_DEVICE double zero() { + uint64_t bits = 0x0ull; + return reinterpret_cast(bits); +} + +/// Returns 0, the additive identity element (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex zero< complex >() { + return complex(zero(), double()); +} + +/// Returns 2 (specialization for double) +template <> MCTLASS_HOST_DEVICE double two() { + uint64_t bits = 0x4000000000000000ull; + return reinterpret_cast(bits); +} + +/// Returns 2 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex two< complex >() { + return complex(two(), double()); +} + +/// Returns pi, approximately 3.141 (specialization for double) +template <> MCTLASS_HOST_DEVICE double pi() { + uint64_t bits = 0x400921fb54442d18ull; + return reinterpret_cast(bits); +} + +/// Returns pi, approximately 3.141 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex pi< complex >() { + return complex(pi(), double()); +} + +/// Returns 2 * pi (specialization for double) +template <> MCTLASS_HOST_DEVICE double two_pi() { + uint64_t bits = 0x401921fb54442d18ull; + return reinterpret_cast(bits); +} + +/// Returns 2 * pi (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex two_pi< complex >() { + return complex(two_pi(), double()); +} + +/// Returns pi / 2 (specialization for double) +template <> MCTLASS_HOST_DEVICE double half_pi() { + uint64_t bits = 0x3ff921fb54442d18ull; + return reinterpret_cast(bits); +} + +/// Returns pi / 2 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half_pi< complex >() { + return complex(half_pi(), double()); +} + +/// Returns sqrt(pi) (specialization for double) +template <> MCTLASS_HOST_DEVICE double root_pi() { + uint64_t bits = 0x3ffc5bf891b4ef6aull; + return reinterpret_cast(bits); +} + +/// Returns sqrt(pi) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_pi< complex >() { + return complex(root_pi(), double()); +} + +/// Returns sqrt(pi / 2) (specialization for double) +template <> MCTLASS_HOST_DEVICE double root_half_pi() { + uint64_t bits = 0x3ff40d931ff62705ull; + return reinterpret_cast(bits); +} + +/// Returns sqrt(pi / 2) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_half_pi< complex >() { + return complex(root_half_pi(), double()); +} + +/// Returns sqrt(2 * pi) (specialization for double) +template <> MCTLASS_HOST_DEVICE double root_two_pi() { + uint64_t bits = 0x40040d931ff62705ull; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2 * pi) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_two_pi< complex >() { + return complex(root_two_pi(), double()); +} + +/// Returns sqrt(ln(4)) (specialization for double) +template <> MCTLASS_HOST_DEVICE double root_ln_four() { + uint64_t bits = 0x3ff2d6abe44afc43ull; + return reinterpret_cast(bits); +} + +/// Returns sqrt(ln(4)) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_ln_four< complex >() { + return complex(root_ln_four(), double()); +} + +/// Returns e, approximately 2.718... (specialization for double) +template <> MCTLASS_HOST_DEVICE double e() { + uint64_t bits = 0x4005bf0a8b145769ull; + return reinterpret_cast(bits); +} + +/// Returns e, approximately 2.718... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex e< complex >() { + return complex(e(), double()); +} + +/// Returns (1/2) (specialization for double) +template <> MCTLASS_HOST_DEVICE double half() { + uint64_t bits = 0x3fe0000000000000ull; + return reinterpret_cast(bits); +} + +/// Returns (1/2) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half< complex >() { + return complex(half(), double()); +} + +/// Returns sqrt(2), approximately 1.414... (specialization for double) +template <> MCTLASS_HOST_DEVICE double root_two() { + uint64_t bits = 0x3ff6a09e667f3bcdull; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2), approximately 1.414... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_two< complex >() { + return complex(root_two(), double()); +} + +/// Returns sqrt(2)/2, approximately 0.707... (specialization for double) +template <> MCTLASS_HOST_DEVICE double half_root_two() { + uint64_t bits = 0x3fe6a09e667f3bcdull; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2)/2, approximately 0.707... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half_root_two< complex >() { + return complex(half_root_two(), double()); +} + +/// Returns ln(2), approximately 0.693... (specialization for double) +template <> MCTLASS_HOST_DEVICE double ln_two() { + uint64_t bits = 0x3fe62e42fefa39efull; + return reinterpret_cast(bits); +} + +/// Returns ln(2), approximately 0.693... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex ln_two< complex >() { + return complex(ln_two(), double()); +} + +/// Returns ln(ln(2)), approximately -0.3665... (specialization for double) +template <> MCTLASS_HOST_DEVICE double ln_ln_two() { + uint64_t bits = 0xbfd774f29bdd6b9full; + return reinterpret_cast(bits); +} + +/// Returns ln(ln(2)), approximately -0.3665... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex ln_ln_two< complex >() { + return complex(ln_ln_two(), double()); +} + +/// Returns 1/3, approximately 0.333... (specialization for double) +template <> MCTLASS_HOST_DEVICE double third() { + uint64_t bits = 0x3fd5555555555555ull; + return reinterpret_cast(bits); +} + +/// Returns 1/3, approximately 0.333... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex third< complex >() { + return complex(third(), double()); +} + +/// Returns 2/3, approximately 0.666... (specialization for double) +template <> MCTLASS_HOST_DEVICE double twothirds() { + uint64_t bits = 0x3fe5555555555555ull; + return reinterpret_cast(bits); +} + +/// Returns 2/3, approximately 0.666... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex twothirds< complex >() { + return complex(twothirds(), double()); +} + +/// Returns pi - 3, approximately 0.1416... (specialization for double) +template <> MCTLASS_HOST_DEVICE double pi_minus_three() { + uint64_t bits = 0x3fc21fb54442d180ull; + return reinterpret_cast(bits); +} + +/// Returns pi - 3, approximately 0.1416... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex pi_minus_three< complex >() { + return complex(pi_minus_three(), double()); +} + +/// Returns 4 - pi, approximately 0.858... (specialization for double) +template <> MCTLASS_HOST_DEVICE double four_minus_pi() { + uint64_t bits = 0x3feb7812aeef4ba0ull; + return reinterpret_cast(bits); +} + +/// Returns 4 - pi, approximately 0.858... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex four_minus_pi< complex >() { + return complex(four_minus_pi(), double()); +} + +///////////////////////////////////////////////////////////////////////////////////// + +// Specialization for float + +/// Returns 1, the multiplicative identity element (specialization for float) +template <> MCTLASS_HOST_DEVICE float one() { + uint32_t bits = 0x3f800000u; + return reinterpret_cast(bits); +} + +/// Returns 1, the multiplicative identity element (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex one< complex >() { + return complex(one(), float()); +} + +/// Returns 0, the additive identity element (specialization for float) +template <> MCTLASS_HOST_DEVICE float zero() { + uint32_t bits = 0x0u; + return reinterpret_cast(bits); +} + +/// Returns 0, the additive identity element (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex zero< complex >() { + return complex(zero(), float()); +} + +/// Returns 2 (specialization for float) +template <> MCTLASS_HOST_DEVICE float two() { + uint32_t bits = 0x40000000u; + return reinterpret_cast(bits); +} + +/// Returns 2 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex two< complex >() { + return complex(two(), float()); +} + +/// Returns pi, approximately 3.141 (specialization for float) +template <> MCTLASS_HOST_DEVICE float pi() { + uint32_t bits = 0x40490fdbu; + return reinterpret_cast(bits); +} + +/// Returns pi, approximately 3.141 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex pi< complex >() { + return complex(pi(), float()); +} + +/// Returns 2 * pi (specialization for float) +template <> MCTLASS_HOST_DEVICE float two_pi() { + uint32_t bits = 0x40c90fdbu; + return reinterpret_cast(bits); +} + +/// Returns 2 * pi (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex two_pi< complex >() { + return complex(two_pi(), float()); +} + +/// Returns pi / 2 (specialization for float) +template <> MCTLASS_HOST_DEVICE float half_pi() { + uint32_t bits = 0x3fc90fdbu; + return reinterpret_cast(bits); +} + +/// Returns pi / 2 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half_pi< complex >() { + return complex(half_pi(), float()); +} + +/// Returns sqrt(pi) (specialization for float) +template <> MCTLASS_HOST_DEVICE float root_pi() { + uint32_t bits = 0x3fe2dfc5u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(pi) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_pi< complex >() { + return complex(root_pi(), float()); +} + +/// Returns sqrt(pi / 2) (specialization for float) +template <> MCTLASS_HOST_DEVICE float root_half_pi() { + uint32_t bits = 0x3fa06c99u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(pi / 2) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_half_pi< complex >() { + return complex(root_half_pi(), float()); +} + +/// Returns sqrt(2 * pi) (specialization for float) +template <> MCTLASS_HOST_DEVICE float root_two_pi() { + uint32_t bits = 0x40206c99u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2 * pi) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_two_pi< complex >() { + return complex(root_two_pi(), float()); +} + +/// Returns sqrt(ln(4)) (specialization for float) +template <> MCTLASS_HOST_DEVICE float root_ln_four() { + uint32_t bits = 0x3f96b55fu; + return reinterpret_cast(bits); +} + +/// Returns sqrt(ln(4)) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_ln_four< complex >() { + return complex(root_ln_four(), float()); +} + +/// Returns e, approximately 2.718... (specialization for float) +template <> MCTLASS_HOST_DEVICE float e() { + uint32_t bits = 0x402df854u; + return reinterpret_cast(bits); +} + +/// Returns e, approximately 2.718... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex e< complex >() { + return complex(e(), float()); +} + +/// Returns (1/2) (specialization for float) +template <> MCTLASS_HOST_DEVICE float half() { + uint32_t bits = 0x3f000000u; + return reinterpret_cast(bits); +} + +/// Returns (1/2) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half< complex >() { + return complex(half(), float()); +} + +/// Returns sqrt(2), approximately 1.414... (specialization for float) +template <> MCTLASS_HOST_DEVICE float root_two() { + uint32_t bits = 0x3fb504f3u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2), approximately 1.414... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_two< complex >() { + return complex(root_two(), float()); +} + +/// Returns sqrt(2)/2, approximately 0.707... (specialization for float) +template <> MCTLASS_HOST_DEVICE float half_root_two() { + uint32_t bits = 0x3f3504f3u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2)/2, approximately 0.707... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half_root_two< complex >() { + return complex(half_root_two(), float()); +} + +/// Returns ln(2), approximately 0.693... (specialization for float) +template <> MCTLASS_HOST_DEVICE float ln_two() { + uint32_t bits = 0x3f317218u; + return reinterpret_cast(bits); +} + +/// Returns ln(2), approximately 0.693... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex ln_two< complex >() { + return complex(ln_two(), float()); +} + +/// Returns ln(ln(2)), approximately -0.3665... (specialization for float) +template <> MCTLASS_HOST_DEVICE float ln_ln_two() { + uint32_t bits = 0xbebba795u; + return reinterpret_cast(bits); +} + +/// Returns ln(ln(2)), approximately -0.3665... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex ln_ln_two< complex >() { + return complex(ln_ln_two(), float()); +} + +/// Returns 1/3, approximately 0.333... (specialization for float) +template <> MCTLASS_HOST_DEVICE float third() { + uint32_t bits = 0x3eaaaaabu; + return reinterpret_cast(bits); +} + +/// Returns 1/3, approximately 0.333... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex third< complex >() { + return complex(third(), float()); +} + +/// Returns 2/3, approximately 0.666... (specialization for float) +template <> MCTLASS_HOST_DEVICE float twothirds() { + uint32_t bits = 0x3f2aaaabu; + return reinterpret_cast(bits); +} + +/// Returns 2/3, approximately 0.666... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex twothirds< complex >() { + return complex(twothirds(), float()); +} + +/// Returns pi - 3, approximately 0.1416... (specialization for float) +template <> MCTLASS_HOST_DEVICE float pi_minus_three() { + uint32_t bits = 0x3e10fdaau; + return reinterpret_cast(bits); +} + +/// Returns pi - 3, approximately 0.1416... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex pi_minus_three< complex >() { + return complex(pi_minus_three(), float()); +} + +/// Returns 4 - pi, approximately 0.858... (specialization for float) +template <> MCTLASS_HOST_DEVICE float four_minus_pi() { + uint32_t bits = 0x3f5bc095u; + return reinterpret_cast(bits); +} + +/// Returns 4 - pi, approximately 0.858... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex four_minus_pi< complex >() { + return complex(four_minus_pi(), float()); +} + +///////////////////////////////////////////////////////////////////////////////////// + +// Specialization for tfloat32_t + +/// Returns 1, the multiplicative identity element (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t one() { + uint32_t bits = 0x3f801000u; + return reinterpret_cast(bits); +} + +/// Returns 1, the multiplicative identity element (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex one< complex >() { + return complex(one(), tfloat32_t()); +} + +/// Returns 0, the additive identity element (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t zero() { + uint32_t bits = 0x1000u; + return reinterpret_cast(bits); +} + +/// Returns 0, the additive identity element (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex zero< complex >() { + return complex(zero(), tfloat32_t()); +} + +/// Returns 2 (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t two() { + uint32_t bits = 0x40001000u; + return reinterpret_cast(bits); +} + +/// Returns 2 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex two< complex >() { + return complex(two(), tfloat32_t()); +} + +/// Returns pi, approximately 3.141 (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t pi() { + uint32_t bits = 0x40491fdbu; + return reinterpret_cast(bits); +} + +/// Returns pi, approximately 3.141 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex pi< complex >() { + return complex(pi(), tfloat32_t()); +} + +/// Returns 2 * pi (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t two_pi() { + uint32_t bits = 0x40c91fdbu; + return reinterpret_cast(bits); +} + +/// Returns 2 * pi (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex two_pi< complex >() { + return complex(two_pi(), tfloat32_t()); +} + +/// Returns pi / 2 (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t half_pi() { + uint32_t bits = 0x3fc91fdbu; + return reinterpret_cast(bits); +} + +/// Returns pi / 2 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half_pi< complex >() { + return complex(half_pi(), tfloat32_t()); +} + +/// Returns sqrt(pi) (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t root_pi() { + uint32_t bits = 0x3fe2efc5u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(pi) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_pi< complex >() { + return complex(root_pi(), tfloat32_t()); +} + +/// Returns sqrt(pi / 2) (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t root_half_pi() { + uint32_t bits = 0x3fa07c99u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(pi / 2) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_half_pi< complex >() { + return complex(root_half_pi(), tfloat32_t()); +} + +/// Returns sqrt(2 * pi) (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t root_two_pi() { + uint32_t bits = 0x40207c99u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2 * pi) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_two_pi< complex >() { + return complex(root_two_pi(), tfloat32_t()); +} + +/// Returns sqrt(ln(4)) (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t root_ln_four() { + uint32_t bits = 0x3f96c55fu; + return reinterpret_cast(bits); +} + +/// Returns sqrt(ln(4)) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_ln_four< complex >() { + return complex(root_ln_four(), tfloat32_t()); +} + +/// Returns e, approximately 2.718... (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t e() { + uint32_t bits = 0x402e0854u; + return reinterpret_cast(bits); +} + +/// Returns e, approximately 2.718... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex e< complex >() { + return complex(e(), tfloat32_t()); +} + +/// Returns (1/2) (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t half() { + uint32_t bits = 0x3f001000u; + return reinterpret_cast(bits); +} + +/// Returns (1/2) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half< complex >() { + return complex(half(), tfloat32_t()); +} + +/// Returns sqrt(2), approximately 1.414... (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t root_two() { + uint32_t bits = 0x3fb514f3u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2), approximately 1.414... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_two< complex >() { + return complex(root_two(), tfloat32_t()); +} + +/// Returns sqrt(2)/2, approximately 0.707... (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t half_root_two() { + uint32_t bits = 0x3f3514f3u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2)/2, approximately 0.707... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half_root_two< complex >() { + return complex(half_root_two(), tfloat32_t()); +} + +/// Returns ln(2), approximately 0.693... (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t ln_two() { + uint32_t bits = 0x3f318218u; + return reinterpret_cast(bits); +} + +/// Returns ln(2), approximately 0.693... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex ln_two< complex >() { + return complex(ln_two(), tfloat32_t()); +} + +/// Returns ln(ln(2)), approximately -0.3665... (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t ln_ln_two() { + uint32_t bits = 0xbebbb795u; + return reinterpret_cast(bits); +} + +/// Returns ln(ln(2)), approximately -0.3665... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex ln_ln_two< complex >() { + return complex(ln_ln_two(), tfloat32_t()); +} + +/// Returns 1/3, approximately 0.333... (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t third() { + uint32_t bits = 0x3eaabaabu; + return reinterpret_cast(bits); +} + +/// Returns 1/3, approximately 0.333... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex third< complex >() { + return complex(third(), tfloat32_t()); +} + +/// Returns 2/3, approximately 0.666... (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t twothirds() { + uint32_t bits = 0x3f2abaabu; + return reinterpret_cast(bits); +} + +/// Returns 2/3, approximately 0.666... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex twothirds< complex >() { + return complex(twothirds(), tfloat32_t()); +} + +/// Returns pi - 3, approximately 0.1416... (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t pi_minus_three() { + uint32_t bits = 0x3e110daau; + return reinterpret_cast(bits); +} + +/// Returns pi - 3, approximately 0.1416... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex pi_minus_three< complex >() { + return complex(pi_minus_three(), tfloat32_t()); +} + +/// Returns 4 - pi, approximately 0.858... (specialization for tfloat32_t) +template <> MCTLASS_HOST_DEVICE tfloat32_t four_minus_pi() { + uint32_t bits = 0x3f5bd095u; + return reinterpret_cast(bits); +} + +/// Returns 4 - pi, approximately 0.858... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex four_minus_pi< complex >() { + return complex(four_minus_pi(), tfloat32_t()); +} + +///////////////////////////////////////////////////////////////////////////////////// + +// Specialization for half_t + +/// Returns 1, the multiplicative identity element (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t one() { + uint16_t bits = 0x3c00u; + return reinterpret_cast(bits); +} + +/// Returns 1, the multiplicative identity element (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex one< complex >() { + return complex(one(), half_t()); +} + +/// Returns 0, the additive identity element (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t zero() { + uint16_t bits = 0x0u; + return reinterpret_cast(bits); +} + +/// Returns 0, the additive identity element (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex zero< complex >() { + return complex(zero(), half_t()); +} + +/// Returns 2 (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t two() { + uint16_t bits = 0x4000u; + return reinterpret_cast(bits); +} + +/// Returns 2 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex two< complex >() { + return complex(two(), half_t()); +} + +/// Returns pi, approximately 3.141 (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t pi() { + uint16_t bits = 0x4248u; + return reinterpret_cast(bits); +} + +/// Returns pi, approximately 3.141 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex pi< complex >() { + return complex(pi(), half_t()); +} + +/// Returns 2 * pi (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t two_pi() { + uint16_t bits = 0x4648u; + return reinterpret_cast(bits); +} + +/// Returns 2 * pi (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex two_pi< complex >() { + return complex(two_pi(), half_t()); +} + +/// Returns pi / 2 (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t half_pi() { + uint16_t bits = 0x3e48u; + return reinterpret_cast(bits); +} + +/// Returns pi / 2 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half_pi< complex >() { + return complex(half_pi(), half_t()); +} + +/// Returns sqrt(pi) (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t root_pi() { + uint16_t bits = 0x3f17u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(pi) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_pi< complex >() { + return complex(root_pi(), half_t()); +} + +/// Returns sqrt(pi / 2) (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t root_half_pi() { + uint16_t bits = 0x3d03u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(pi / 2) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_half_pi< complex >() { + return complex(root_half_pi(), half_t()); +} + +/// Returns sqrt(2 * pi) (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t root_two_pi() { + uint16_t bits = 0x4103u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2 * pi) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_two_pi< complex >() { + return complex(root_two_pi(), half_t()); +} + +/// Returns sqrt(ln(4)) (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t root_ln_four() { + uint16_t bits = 0x3cb6u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(ln(4)) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_ln_four< complex >() { + return complex(root_ln_four(), half_t()); +} + +/// Returns e, approximately 2.718... (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t e() { + uint16_t bits = 0x4170u; + return reinterpret_cast(bits); +} + +/// Returns e, approximately 2.718... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex e< complex >() { + return complex(e(), half_t()); +} + +/// Returns (1/2) (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t half() { + uint16_t bits = 0x3800u; + return reinterpret_cast(bits); +} + +/// Returns (1/2) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half< complex >() { + return complex(half(), half_t()); +} + +/// Returns sqrt(2), approximately 1.414... (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t root_two() { + uint16_t bits = 0x3da8u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2), approximately 1.414... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_two< complex >() { + return complex(root_two(), half_t()); +} + +/// Returns sqrt(2)/2, approximately 0.707... (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t half_root_two() { + uint16_t bits = 0x39a8u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2)/2, approximately 0.707... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half_root_two< complex >() { + return complex(half_root_two(), half_t()); +} + +/// Returns ln(2), approximately 0.693... (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t ln_two() { + uint16_t bits = 0x398cu; + return reinterpret_cast(bits); +} + +/// Returns ln(2), approximately 0.693... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex ln_two< complex >() { + return complex(ln_two(), half_t()); +} + +/// Returns ln(ln(2)), approximately -0.3665... (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t ln_ln_two() { + uint16_t bits = 0xb5ddu; + return reinterpret_cast(bits); +} + +/// Returns ln(ln(2)), approximately -0.3665... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex ln_ln_two< complex >() { + return complex(ln_ln_two(), half_t()); +} + +/// Returns 1/3, approximately 0.333... (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t third() { + uint16_t bits = 0x3555u; + return reinterpret_cast(bits); +} + +/// Returns 1/3, approximately 0.333... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex third< complex >() { + return complex(third(), half_t()); +} + +/// Returns 2/3, approximately 0.666... (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t twothirds() { + uint16_t bits = 0x3955u; + return reinterpret_cast(bits); +} + +/// Returns 2/3, approximately 0.666... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex twothirds< complex >() { + return complex(twothirds(), half_t()); +} + +/// Returns pi - 3, approximately 0.1416... (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t pi_minus_three() { + uint16_t bits = 0x3088u; + return reinterpret_cast(bits); +} + +/// Returns pi - 3, approximately 0.1416... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex pi_minus_three< complex >() { + return complex(pi_minus_three(), half_t()); +} + +/// Returns 4 - pi, approximately 0.858... (specialization for half_t) +template <> MCTLASS_HOST_DEVICE half_t four_minus_pi() { + uint16_t bits = 0x3adeu; + return reinterpret_cast(bits); +} + +/// Returns 4 - pi, approximately 0.858... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex four_minus_pi< complex >() { + return complex(four_minus_pi(), half_t()); +} + +///////////////////////////////////////////////////////////////////////////////////// + +// Specialization for bfloat16_t + +/// Returns 1, the multiplicative identity element (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t one() { + uint16_t bits = 0x3f80u; + return reinterpret_cast(bits); +} + +/// Returns 1, the multiplicative identity element (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex one< complex >() { + return complex(one(), bfloat16_t()); +} + +/// Returns 0, the additive identity element (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t zero() { + uint16_t bits = 0x0u; + return reinterpret_cast(bits); +} + +/// Returns 0, the additive identity element (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex zero< complex >() { + return complex(zero(), bfloat16_t()); +} + +/// Returns 2 (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t two() { + uint16_t bits = 0x4000u; + return reinterpret_cast(bits); +} + +/// Returns 2 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex two< complex >() { + return complex(two(), bfloat16_t()); +} + +/// Returns pi, approximately 3.141 (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t pi() { + uint16_t bits = 0x4049u; + return reinterpret_cast(bits); +} + +/// Returns pi, approximately 3.141 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex pi< complex >() { + return complex(pi(), bfloat16_t()); +} + +/// Returns 2 * pi (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t two_pi() { + uint16_t bits = 0x40c9u; + return reinterpret_cast(bits); +} + +/// Returns 2 * pi (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex two_pi< complex >() { + return complex(two_pi(), bfloat16_t()); +} + +/// Returns pi / 2 (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t half_pi() { + uint16_t bits = 0x3fc9u; + return reinterpret_cast(bits); +} + +/// Returns pi / 2 (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half_pi< complex >() { + return complex(half_pi(), bfloat16_t()); +} + +/// Returns sqrt(pi) (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t root_pi() { + uint16_t bits = 0x3fe3u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(pi) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_pi< complex >() { + return complex(root_pi(), bfloat16_t()); +} + +/// Returns sqrt(pi / 2) (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t root_half_pi() { + uint16_t bits = 0x3fa0u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(pi / 2) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_half_pi< complex >() { + return complex(root_half_pi(), bfloat16_t()); +} + +/// Returns sqrt(2 * pi) (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t root_two_pi() { + uint16_t bits = 0x4020u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2 * pi) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_two_pi< complex >() { + return complex(root_two_pi(), bfloat16_t()); +} + +/// Returns sqrt(ln(4)) (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t root_ln_four() { + uint16_t bits = 0x3f97u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(ln(4)) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_ln_four< complex >() { + return complex(root_ln_four(), bfloat16_t()); +} + +/// Returns e, approximately 2.718... (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t e() { + uint16_t bits = 0x402eu; + return reinterpret_cast(bits); +} + +/// Returns e, approximately 2.718... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex e< complex >() { + return complex(e(), bfloat16_t()); +} + +/// Returns (1/2) (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t half() { + uint16_t bits = 0x3f00u; + return reinterpret_cast(bits); +} + +/// Returns (1/2) (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half< complex >() { + return complex(half(), bfloat16_t()); +} + +/// Returns sqrt(2), approximately 1.414... (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t root_two() { + uint16_t bits = 0x3fb5u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2), approximately 1.414... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex root_two< complex >() { + return complex(root_two(), bfloat16_t()); +} + +/// Returns sqrt(2)/2, approximately 0.707... (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t half_root_two() { + uint16_t bits = 0x3f35u; + return reinterpret_cast(bits); +} + +/// Returns sqrt(2)/2, approximately 0.707... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex half_root_two< complex >() { + return complex(half_root_two(), bfloat16_t()); +} + +/// Returns ln(2), approximately 0.693... (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t ln_two() { + uint16_t bits = 0x3f31u; + return reinterpret_cast(bits); +} + +/// Returns ln(2), approximately 0.693... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex ln_two< complex >() { + return complex(ln_two(), bfloat16_t()); +} + +/// Returns ln(ln(2)), approximately -0.3665... (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t ln_ln_two() { + uint16_t bits = 0xbebcu; + return reinterpret_cast(bits); +} + +/// Returns ln(ln(2)), approximately -0.3665... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex ln_ln_two< complex >() { + return complex(ln_ln_two(), bfloat16_t()); +} + +/// Returns 1/3, approximately 0.333... (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t third() { + uint16_t bits = 0x3eabu; + return reinterpret_cast(bits); +} + +/// Returns 1/3, approximately 0.333... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex third< complex >() { + return complex(third(), bfloat16_t()); +} + +/// Returns 2/3, approximately 0.666... (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t twothirds() { + uint16_t bits = 0x3f2bu; + return reinterpret_cast(bits); +} + +/// Returns 2/3, approximately 0.666... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex twothirds< complex >() { + return complex(twothirds(), bfloat16_t()); +} + +/// Returns pi - 3, approximately 0.1416... (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t pi_minus_three() { + uint16_t bits = 0x3e11u; + return reinterpret_cast(bits); +} + +/// Returns pi - 3, approximately 0.1416... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex pi_minus_three< complex >() { + return complex(pi_minus_three(), bfloat16_t()); +} + +/// Returns 4 - pi, approximately 0.858... (specialization for bfloat16_t) +template <> MCTLASS_HOST_DEVICE bfloat16_t four_minus_pi() { + uint16_t bits = 0x3f5cu; + return reinterpret_cast(bits); +} + +/// Returns 4 - pi, approximately 0.858... (specialization for complex) +template <> MCTLASS_HOST_DEVICE complex four_minus_pi< complex >() { + return complex(four_minus_pi(), bfloat16_t()); +} +/////////////////////////////////////////////////////////////////////////////////// + +} // namespace constants +} // namespace mctlass + +/////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/conv2d_problem_size.h b/csrc/mctlass/include/mctlass/conv/conv2d_problem_size.h new file mode 100644 index 0000000..1902fb9 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/conv2d_problem_size.h @@ -0,0 +1,645 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief This file contains definitions and utility functions for describing convolution problem sizes. + + Conv2dProblem desciption: + activation (NHWC), + filter (KRSC), + output (NPQK), + pading (pad_h, pad_w), + stride (stride_h, stride_w), + dilation (dilation_h, dilation_w). + + Free functions to map: + Map tensor extents (Conv2d -> ImplicitGemm) : implicit_gemm_tensor_[a|b|c]_extent(ConvolutionOperator) + Map tensor sizes (Conv2d -> ImplicitGemm) : implicit_gemm_tensor_[a|b|c]_size(ConvolutionOperator) + Map tensor problem sizes (Conv2d -> ImplicitGemm): implicit_gemm_problem_size(ConvolutionOperator) +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/tensor_coord.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/functional.h" + +namespace mctlass { +namespace conv { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Problem size structure +struct Conv2dProblemSize { + + // Conv2d strictly problem size parameters + int N, H, W, C, P, Q, K, R, S; + int pad_h, pad_w; + int stride_h, stride_w; + int dilation_h, dilation_w; + Mode mode; + + // Conv2d implementation-related parameters + int split_k_slices; + int groups; + + // + // Methods + // + +public: + MCTLASS_HOST_DEVICE + Conv2dProblemSize(): + N(0), H(0), W(0), C(0), P(0), Q(0), K(0), R(0), S(0), + pad_h(0), pad_w(0), stride_h(1), stride_w(1), dilation_h(1), dilation_w(1), + mode(Mode::kConvolution), split_k_slices(1), groups(1) { } + + /// Constructor for default padding, stride, dilation, and split-K + MCTLASS_HOST_DEVICE + Conv2dProblemSize( + int N, + int H, + int W, + int C, + int P, + int Q, + int K, + int R, + int S, + Mode mode + ): + N(N), H(H), W(W), C(C), P(P), Q(Q), K(K), R(R), S(S), + pad_h(R / 2), pad_w(S / 2), stride_h(1), stride_w(1), dilation_h(1), dilation_w(1), + mode(mode), split_k_slices(1), groups (1) { } + + /// Constructor + MCTLASS_HOST_DEVICE + Conv2dProblemSize( + int N, + int H, + int W, + int C, + int K, + int R, + int S, + int P, + int Q, + int pad_h, + int pad_w, + int stride_h, + int stride_w, + int dilation_h, + int dilation_w, + Mode mode, + int split_k_slices = 1, + int groups = 1 + ): + N(N), H(H), W(W), C(C), K(K), R(R), S(S), P(P), Q(Q), + pad_h(pad_h), pad_w(pad_w), stride_h(stride_h), stride_w(stride_w), + dilation_h(dilation_h), dilation_w(dilation_w), + mode(mode), split_k_slices(split_k_slices), groups (groups) { } + + /// Constructs convolution problem size from mctlass Tensor4DCoord and MatrixCoord + // set user-defined output size and sets P and Q (include all data members in ctor) + MCTLASS_HOST_DEVICE + Conv2dProblemSize( + mctlass::Tensor4DCoord input_size, // NHWC + mctlass::Tensor4DCoord filter_size, // KRSC + mctlass::Tensor4DCoord padding, // pad_h, _, pad_w, _ + mctlass::MatrixCoord stride, // stride_h, stride_w + mctlass::MatrixCoord dilation, // dilation_h, dilation_w + mctlass::Tensor4DCoord output_size, // NPQK + mctlass::conv::Mode mode = mctlass::conv::Mode::kCrossCorrelation, + int split_k_slices = 1, + int groups = 1 + ): + N(input_size.n()), H(input_size.h()), W(input_size.w()), C(input_size.c()), + K(filter_size.n()), R(filter_size.h()), S(filter_size.w()), + pad_h(padding[0]), pad_w(padding[2]), + stride_h(stride.row()), stride_w(stride.column()), + dilation_h(dilation.row()), dilation_w(dilation.column()), + P(output_size.h()), Q(output_size.w()), + mode(mode), split_k_slices(split_k_slices), groups(groups) {} + + /// Constructs convolution problem size from mctlass Tensor4DCoord and MatrixCoord + // computes output size and sets P and Q (skip output from ctor arguments) + MCTLASS_HOST_DEVICE + Conv2dProblemSize( + mctlass::Tensor4DCoord input_size, // NHWC + mctlass::Tensor4DCoord filter_size, // KRSC + mctlass::Tensor4DCoord padding, // pad_h, _, pad_w, _ + mctlass::MatrixCoord stride, // stride_h, stride_w + mctlass::MatrixCoord dilation, // dilation_h, dilation_w + mctlass::conv::Mode mode = mctlass::conv::Mode::kCrossCorrelation, + int split_k_slices = 1, + int groups = 1 + ): + N(input_size.n()), H(input_size.h()), W(input_size.w()), C(input_size.c()), + K(filter_size.n()), R(filter_size.h()), S(filter_size.w()), + pad_h(padding[0]), pad_w(padding[2]), + stride_h(stride.row()), stride_w(stride.column()), + dilation_h(dilation.row()), dilation_w(dilation.column()), + mode(mode), split_k_slices(split_k_slices), groups(groups) { + // set output P and Q + P = ((H + pad_h * 2 - R * dilation_h) / stride_h) + 1; + Q = ((W + pad_w * 2 - S * dilation_w) / stride_w) + 1; + } + + /// Constructs convolution problem size from mctlass Tensor4DCoord and MatrixCoord + // set user-defined output size and sets P and Q (skip padding, striding, and dilation) + MCTLASS_HOST_DEVICE + Conv2dProblemSize( + mctlass::Tensor4DCoord input_size, // NHWC + mctlass::Tensor4DCoord filter_size, // KRSC + mctlass::Tensor4DCoord output_size, // NPQK + mctlass::conv::Mode mode = mctlass::conv::Mode::kCrossCorrelation, + int split_k_slices = 1, + int groups = 1 + ): + N(input_size.n()), H(input_size.h()), W(input_size.w()), C(input_size.c()), + K(filter_size.n()), R(filter_size.h()), S(filter_size.w()), + P(output_size.h()), Q(output_size.w()), + pad_h(R / 2), pad_w(S / 2), stride_h(1), stride_w(1), + dilation_h(1), dilation_w(1), + mode(mode), split_k_slices(split_k_slices), groups(groups) {} + + // Reset covolution mode in the problem + MCTLASS_HOST_DEVICE + Conv2dProblemSize reset_mode(mctlass::conv::Mode mode_) { + Conv2dProblemSize tmp(*this); + tmp.mode = mode_; + return tmp; + } + + // Reset covolution mode in the problem + MCTLASS_HOST_DEVICE + Conv2dProblemSize reset_split_k_slices(int split_k_slices_) { + Conv2dProblemSize tmp(*this); + tmp.split_k_slices = split_k_slices_; + return tmp; + } + + /// Equality operator (ignores mode and split_k_slice) + MCTLASS_HOST_DEVICE + bool operator==(Conv2dProblemSize const &conv) const { + return ( + (N == conv.N) && (H == conv.H) && (W == conv.W) && (C == conv.C) && + (K == conv.K) && (R == conv.R) && (S == conv.S) && + (P == conv.P) && (Q == conv.Q) && + (pad_h == conv.pad_h) && (pad_w == conv.pad_w) && + (stride_h == conv.stride_h) && (stride_w == conv.stride_w) && + (dilation_h == conv.dilation_h) && (dilation_w == conv.dilation_w) + ); + } + + /// Inequality operator + MCTLASS_HOST_DEVICE + bool operator!=(Conv2dProblemSize const &rhs) const { + return !(*this == rhs); + } + + /// Returns activation extent as Tensor4DCoord + MCTLASS_HOST_DEVICE + mctlass::Tensor4DCoord activation_extent() const { + + return mctlass::Tensor4DCoord ({N, H, W, C}); + } + + /// Returns filter extent as Tensor4DCoord + MCTLASS_HOST_DEVICE + mctlass::Tensor4DCoord filter_extent() const { + + return mctlass::Tensor4DCoord ({K, R, S, C / groups}); + } + + /// Returns output extent as Tensor4DCoord + MCTLASS_HOST_DEVICE + mctlass::Tensor4DCoord output_extent() const { + + return mctlass::Tensor4DCoord ({N, P, Q, K}); + } + + /// Returns activation size in number of elements + MCTLASS_HOST_DEVICE + int64_t activation_size() const { + + return (N * H * W * C); + } + + /// Returns filter size in number of elements + MCTLASS_HOST_DEVICE + int64_t filter_size() const { + + return (K * R * S * C / groups); + } + + /// Returns output size in number of elements + MCTLASS_HOST_DEVICE + int64_t output_size() const { + + return (N * P * Q * K); + } + + /// Returns padding as Tensor4DCoord + MCTLASS_HOST_DEVICE + mctlass::Tensor4DCoord padding() const { + + return mctlass::Tensor4DCoord ({pad_h, pad_h, pad_w, pad_w}); + } + + /// Returns stride as MatrixCoord + MCTLASS_HOST_DEVICE + mctlass::MatrixCoord stride() const { + + return mctlass::MatrixCoord ({stride_h, stride_w}); + } + + /// Returns dilation as MatrixCoord + MCTLASS_HOST_DEVICE + mctlass::MatrixCoord dilation() const { + + return mctlass::MatrixCoord ({dilation_h, dilation_w}); + } + + ///////////////////////////////////////////////////////////////// + // Methods used for strided dgrad implementation + ///////////////////////////////////////////////////////////////// + /// Number of filter r positions to accumulate in gemm-k dim + MCTLASS_HOST_DEVICE + int num_gemm_k_filter_r(int r) const { + return ((R - r + stride_h - 1) / stride_h); + } + + /// Number of filter s positions to accumulate in gemm-k dim + MCTLASS_HOST_DEVICE + int num_gemm_k_filter_s(int s) const { + return ((S - s + stride_w - 1) / stride_w); + } + + /// Number of filter positions to accumulate in gemm-k dim + MCTLASS_HOST_DEVICE + int num_gemm_k_filter_positions(int r, int s) const { + return num_gemm_k_filter_r(r) * num_gemm_k_filter_s(s); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// ImplicitGemm helper functions // +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Determine the problem size of the implicit GEMM operation +MCTLASS_HOST_DEVICE +mctlass::gemm::GemmCoord implicit_gemm_problem_size( + Operator conv_operator, + Conv2dProblemSize const &problem_size) { + // Compute problem size + switch (conv_operator) { + case Operator::kFprop: + return gemm::GemmCoord( + problem_size.N * problem_size.P * problem_size.Q, + problem_size.K, + problem_size.R * problem_size.S * problem_size.C / problem_size.groups + ); + case Operator::kDgrad: + return gemm::GemmCoord( + problem_size.N * problem_size.H * problem_size.W, + problem_size.C, + problem_size.R * problem_size.S * problem_size.K + ); + case Operator::kWgrad: + return gemm::GemmCoord( + problem_size.K, + problem_size.R * problem_size.S * problem_size.C, + problem_size.N * problem_size.P * problem_size.Q + ); + default: + break; + } + return gemm::GemmCoord(); +} + +// Determine the number of gemm_k iterations for conv2d problem using implicit gemm algorithm +MCTLASS_HOST_DEVICE +int implicit_gemm_k_iterations( + Operator conv_operator, + int threadblock_K, + Conv2dProblemSize const &problem_size, + IteratorAlgorithm algorithm = IteratorAlgorithm::kAnalytic, + GroupMode group_mode = GroupMode::kNone, + int threadblock_N = 0) { + + int iterations = 0; + + if (group_mode == GroupMode::kNone) { + + if (algorithm == IteratorAlgorithm::kFixedChannels) { + + int positions_per_iteration = threadblock_K / problem_size.C; + switch (conv_operator) { + case Operator::kFprop: + iterations = (problem_size.R * problem_size.S + positions_per_iteration - 1 ) / positions_per_iteration; + break; + + default: + break; + } + } + else if (algorithm == IteratorAlgorithm::kFewChannels) { + + switch (conv_operator) { + case Operator::kFprop: + iterations = (problem_size.R * problem_size.S * problem_size.C + threadblock_K - 1 ) / threadblock_K; + break; + + default: + break; + } + } + else { + int elements_per_split_k_slice = 0; + + switch (conv_operator) { + case Operator::kFprop: + elements_per_split_k_slice = (problem_size.C + problem_size.split_k_slices - 1) / problem_size.split_k_slices; + iterations = problem_size.R * problem_size.S * ((elements_per_split_k_slice + threadblock_K - 1) / threadblock_K); + break; + + case Operator::kDgrad: + elements_per_split_k_slice = (problem_size.K + problem_size.split_k_slices - 1) / problem_size.split_k_slices; + iterations = problem_size.R * problem_size.S * ((elements_per_split_k_slice + threadblock_K - 1) / threadblock_K); + break; + + case Operator::kWgrad: + elements_per_split_k_slice = (problem_size.N * problem_size.P * problem_size.Q + problem_size.split_k_slices - 1) / problem_size.split_k_slices; + iterations = (elements_per_split_k_slice + threadblock_K - 1) / threadblock_K; + break; + + default: + break; + } + } + + } else if (group_mode == GroupMode::kDepthwise) { + int channels_per_cta = threadblock_N; + + if (algorithm == IteratorAlgorithm::kAnalytic) { + switch (conv_operator) { + case Operator::kFprop: + iterations = problem_size.R * problem_size.S * + ((channels_per_cta + threadblock_K - 1) / threadblock_K); + break; + + default: + break; + } + } + } else { // Group conv + + int channels_per_group = problem_size.C / problem_size.groups; + int k_per_group = problem_size.K / problem_size.groups; + + if (algorithm == IteratorAlgorithm::kAnalytic) { + switch (conv_operator) { + case Operator::kFprop: + iterations = problem_size.R * problem_size.S * ((channels_per_group + threadblock_K - 1) / threadblock_K); + // In group conv, if k_per_group < threadblock_N, one Threadblock will calculate multiple groups + if (problem_size.groups != 1) { + if (k_per_group < threadblock_N) { + iterations *= threadblock_N / k_per_group; + } + } + break; + + default: + break; + } + } else if (algorithm == IteratorAlgorithm::kOptimized) { + // Current optimized iterator only support GroupMode::kSingleGroup + if (group_mode == GroupMode::kSingleGroup) { + switch (conv_operator) { + case Operator::kFprop: + iterations = problem_size.R * problem_size.S * ((channels_per_group + threadblock_K - 1) / threadblock_K); + break; + + default: + break; + } + } + } + + } + + return iterations; +} + + +template +MCTLASS_HOST_DEVICE +int depthwise_gemm_k_iterations( + Operator conv_operator, + int threadblock_K, + Conv2dProblemSize const &problem_size, + IteratorAlgorithm algorithm = IteratorAlgorithm::kAnalytic, + GroupMode group_mode = GroupMode::kNone, + int threadblock_N = 0) { + + int n = problem_size.N; + int p = (problem_size.P + Output_P - 1) / Output_P; + int q = (problem_size.Q + Output_Q - 1) / Output_Q; + + int iterations = (n * p * q + problem_size.split_k_slices - 1) / problem_size.split_k_slices; + return iterations; +} + + +MCTLASS_HOST_DEVICE +int implicit_gemm_k_iterations_per_channel( + Operator conv_operator, + int threadblock_K, + Conv2dProblemSize const &problem_size, + IteratorAlgorithm algorithm = IteratorAlgorithm::kAnalytic) { + + int iterations = 0; //0 means not applicable + if (algorithm == IteratorAlgorithm::kAnalytic || algorithm == IteratorAlgorithm::kOptimized) { + switch (conv_operator) { + case Operator::kFprop: + iterations = problem_size.R * problem_size.S; + break; + + case Operator::kDgrad: + iterations = problem_size.R * problem_size.S; + break; + + default: + break; + } + } + return iterations; +} + +//////////////////////////////////////////////////////////////////////////////// +// Mapping function (ImplicitGemm A, B, C -> Conv Activation, Filter, Output) +//////////////////////////////////////////////////////////////////////////////// +/// Returns ImplicitGemm tensor A extent as Tensor4DCoord +MCTLASS_HOST_DEVICE +mctlass::Tensor4DCoord implicit_gemm_tensor_a_extent( + Operator conv_operator, + Conv2dProblemSize const &problem_size) { + switch (conv_operator) { + case mctlass::conv::Operator::kFprop: return problem_size.activation_extent(); + case mctlass::conv::Operator::kDgrad: return problem_size.output_extent(); + case mctlass::conv::Operator::kWgrad: return problem_size.output_extent(); + default : break; + } + return mctlass::Tensor4DCoord(); +} + +/// Returns ImplicitGemm tensor B extent as Tensor4DCoord +MCTLASS_HOST_DEVICE +mctlass::Tensor4DCoord implicit_gemm_tensor_b_extent( + Operator conv_operator, + Conv2dProblemSize const &problem_size) { + switch (conv_operator) { + case mctlass::conv::Operator::kFprop: return problem_size.filter_extent(); + case mctlass::conv::Operator::kDgrad: return problem_size.filter_extent(); + case mctlass::conv::Operator::kWgrad: return problem_size.activation_extent(); + default : break; + } + return mctlass::Tensor4DCoord(); +} + +/// Returns ImplicitGemm tensor C extent as Tensor4DCoord +MCTLASS_HOST_DEVICE +mctlass::Tensor4DCoord implicit_gemm_tensor_c_extent( + Operator conv_operator, + Conv2dProblemSize const &problem_size) { + switch (conv_operator) { + case mctlass::conv::Operator::kFprop: return problem_size.output_extent(); + case mctlass::conv::Operator::kDgrad: return problem_size.activation_extent(); + case mctlass::conv::Operator::kWgrad: return problem_size.filter_extent(); + default : break; + } + return mctlass::Tensor4DCoord(); +} + +/// Returns ImplicitGemm tensor A size in number of elements +MCTLASS_HOST_DEVICE +int64_t implicit_gemm_tensor_a_size( + Operator conv_operator, + Conv2dProblemSize const &problem_size) { + switch (conv_operator) { + case mctlass::conv::Operator::kFprop: return problem_size.activation_size(); + case mctlass::conv::Operator::kDgrad: return problem_size.output_size(); + case mctlass::conv::Operator::kWgrad: return problem_size.output_size(); + default : break; + } + return 0; +} + +/// Returns ImplicitGemm tensor B size in number of elements +MCTLASS_HOST_DEVICE +int64_t implicit_gemm_tensor_b_size( + Operator conv_operator, + Conv2dProblemSize const &problem_size) { + switch (conv_operator) { + case mctlass::conv::Operator::kFprop: return problem_size.filter_size(); + case mctlass::conv::Operator::kDgrad: return problem_size.filter_size(); + case mctlass::conv::Operator::kWgrad: return problem_size.activation_size(); + default : break; + } + return 0; +} + +/// Returns ImplicitGemm tensor C size in number of elements +MCTLASS_HOST_DEVICE +int64_t implicit_gemm_tensor_c_size( + Operator conv_operator, + Conv2dProblemSize const &problem_size) { + switch (conv_operator) { + case mctlass::conv::Operator::kFprop: return problem_size.output_size(); + case mctlass::conv::Operator::kDgrad: return problem_size.activation_size(); + case mctlass::conv::Operator::kWgrad: return problem_size.filter_size(); + default : break; + } + return 0; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Strided dgrad helper functions // +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Returns number of CTAs tile M to cover valid MMAs per starting filter postion +MCTLASS_HOST_DEVICE +int strided_dgrad_tile_m_per_filter( + Conv2dProblemSize const &problem_size, + int tile_size_m) { + + // Compute NHW rows in Dx output that needs MMA per starting filter position + int rows_h_per_filter = (problem_size.H + problem_size.stride_h - 1) / problem_size.stride_h; + int rows_w_per_filter = (problem_size.W + problem_size.stride_w - 1) / problem_size.stride_w; + int rows_nhw_per_filter = problem_size.N * rows_h_per_filter * rows_w_per_filter; + + // Number of CTAs tile M to cover valid MMAs per starting filter postion + int tile_m_per_filter = (rows_nhw_per_filter + tile_size_m - 1) / tile_size_m; + + return tile_m_per_filter; +} + +// Computes starting Dx coord (h, w) for given starting filter postion +MCTLASS_HOST_DEVICE +void strided_dgrad_starting_coords( + Conv2dProblemSize const &problem_size, + FastDivmod const &stride_h_divmod, FastDivmod const &stride_w_divmod, + int r, int s, + int &start_h, int &start_w) { + + // function locals for remainder by fast divmod + int pad_h_rem_, pad_w_rem_; + + // start_h = std::abs(problem_size.stride_h - ((problem_size.pad_h % problem_size.stride_h) - r)) % problem_size.stride_h; + stride_h_divmod.divmod(pad_h_rem_, problem_size.pad_h); + int r_ = absolute_value(problem_size.stride_h - (pad_h_rem_ - r)); + stride_h_divmod.divmod(start_h, r_); + + //start_w = std::abs(problem_size.stride_w - ((problem_size.pad_w % problem_size.stride_w) - s)) % problem_size.stride_w; + stride_w_divmod.divmod(pad_w_rem_, problem_size.pad_w); + int s_ = absolute_value(problem_size.stride_w - (pad_w_rem_ - s)); + stride_w_divmod.divmod(start_w, s_); +} + +} // namespace conv +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/conv3d_problem_size.h b/csrc/mctlass/include/mctlass/conv/conv3d_problem_size.h new file mode 100644 index 0000000..880b74d --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/conv3d_problem_size.h @@ -0,0 +1,477 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief This file contains definitions and utility functions for describing convolution problem sizes. + + Conv3dProblem desciption: + activation (NDHWC), + filter (KTRSC), + output (NZPQK), + pading (pad_d, pad_h, pad_w), + stride (stride_d, stride_h, stride_w), + dilation (dilation_d, dilation_h, dilation_w). + + Free functions to map: + Map tensor extents (Conv3d -> ImplicitGemm) : implicit_gemm_tensor_[a|b|c]_extent(ConvolutionOperator) + Map tensor sizes (Conv3d -> ImplicitGemm) : implicit_gemm_tensor_[a|b|c]_size(ConvolutionOperator) + Map tensor problem sizes (Conv3d -> ImplicitGemm): implicit_gemm_problem_size(ConvolutionOperator) +*/ + +#pragma once + +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" + +namespace mctlass { +namespace conv { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Problem size structure +struct Conv3dProblemSize : public Conv2dProblemSize { + // + // Type definitions + // + + // 3D coordinate for padding, stride, and dilation in (d, h, w) dimensions + using Coord3D = Coord<3>; + + // + // Data members + // + + // Conv3d strictly problem size parameters + int D, T, Z; // input depth, filter depth, output depth + int pad_d; // padding in depth dimension + int stride_d; // stride in depth dimension + int dilation_d; // dilation in depth dimension + + // + // Methods + // +public: + MCTLASS_HOST_DEVICE + Conv3dProblemSize(): + D(0), T(0), Z(0), + pad_d(0), + stride_d(1), + dilation_d(1), + Conv2dProblemSize() { } + + /// Constructor for default padding, stride, dilation, and split-K + MCTLASS_HOST_DEVICE + Conv3dProblemSize( + int N, + int D, + int H, + int W, + int C, + int Z, + int P, + int Q, + int K, + int T, + int R, + int S, + Mode mode + ): + D(D), T(T), Z(Z), + pad_d(T / 2), stride_d(1), dilation_d(1), + Conv2dProblemSize(N, H, W, C, P, Q, K, R, S, mode) { } + + /// Constructor + MCTLASS_HOST_DEVICE + Conv3dProblemSize( + int N, + int D, + int H, + int W, + int C, + int K, + int T, + int R, + int S, + int Z, + int P, + int Q, + int pad_d, + int pad_h, + int pad_w, + int stride_d, + int stride_h, + int stride_w, + int dilation_d, + int dilation_h, + int dilation_w, + Mode mode, + int split_k_slices = 1, + int groups = 1 + ): + D(D), T(T), Z(Z), + pad_d(pad_d), stride_d(stride_d), dilation_d(dilation_d), + Conv2dProblemSize( + N, H, W, C, K, R, S, P, Q, + pad_h, pad_w, + stride_h, stride_w, + dilation_h, dilation_w, + mode, split_k_slices, groups) { } + + /// Constructs convolution problem size from mctlass Tensor5DCoord and Coord3D + // set *user-defined* output size and sets Z, P, and Q (include all data members in ctor) + MCTLASS_HOST_DEVICE + Conv3dProblemSize( + mctlass::Tensor5DCoord input_size, // NDHWC + mctlass::Tensor5DCoord filter_size, // KTRSC + Coord3D padding, // pad_d, pad_h, pad_w + Coord3D stride, // stride_d, stride_h, stride_w + Coord3D dilation, // dilation_d, dilation_h, dilation_w + mctlass::Tensor5DCoord output_size, // NZPQK + mctlass::conv::Mode mode = mctlass::conv::Mode::kCrossCorrelation, + int split_k_slices = 1, + int groups = 1 + ): + D(input_size.d()), T(filter_size.d()), Z(output_size.d()), + pad_d(padding[0]), stride_d(stride[0]), dilation_d(dilation[0]), + Conv2dProblemSize( + {input_size.n(), input_size.h(), input_size.w(), input_size.c()}, + {filter_size.n(), filter_size.h(), filter_size.w(), filter_size.c()}, + {padding[1], padding[1], padding[2], padding[2]}, + {stride[1], stride[2]}, + {dilation[1], dilation[2]}, + {output_size.n(), output_size.h(), output_size.w(), output_size.c()}, + mode, split_k_slices, groups + ) { } + + /// Constructs convolution problem size from mctlass Tensor5DCoord and Coord3D + // *computes* output size and sets Z, P and Q (include all data members in ctor) + MCTLASS_HOST_DEVICE + Conv3dProblemSize( + mctlass::Tensor5DCoord input_size, // NDHWC + mctlass::Tensor5DCoord filter_size, // KTRSC + Coord3D padding, // pad_d, pad_h, pad_w + Coord3D stride, // stride_d, stride_h, stride_w + Coord3D dilation, // dilation_d, dilation_h, dilation_w + mctlass::conv::Mode mode = mctlass::conv::Mode::kCrossCorrelation, + int split_k_slices = 1, + int groups = 1 + ): + D(input_size.d()), T(filter_size.d()), + pad_d(padding[0]), stride_d(stride[0]), dilation_d(dilation[0]), + Conv2dProblemSize( + {input_size.n(), input_size.h(), input_size.w(), input_size.c()}, + {filter_size.n(), filter_size.h(), filter_size.w(), filter_size.c()}, + {padding[1], padding[1], padding[2], padding[2]}, + {stride[1], stride[2]}, + {dilation[1], dilation[2]}, + mode, split_k_slices, groups + ) { + // set output Z + Z = ((D + pad_d * 2 - T * dilation_d) / stride_d) + 1; + } + + /// Equality operator (ignores mode and split_k_slice) + MCTLASS_HOST_DEVICE + bool operator==(Conv3dProblemSize const &conv) const { + return ( + (N == conv.N) && (D == conv.D) && (H == conv.H) && (W == conv.W) && (C == conv.C) && + (K == conv.K) && (T == conv.T) && (R == conv.R) && (S == conv.S) && + (Z == conv.Z) &&(P == conv.P) && (Q == conv.Q) && + (pad_d == conv.pad_d) && (pad_h == conv.pad_h) && (pad_w == conv.pad_w) && + (stride_d == conv.stride_d) && (stride_h == conv.stride_h) && (stride_w == conv.stride_w) && + (dilation_d == conv.dilation_d) && (dilation_h == conv.dilation_h) && (dilation_w == conv.dilation_w) + ); + } + + /// Inequality operator + MCTLASS_HOST_DEVICE + bool operator!=(Conv3dProblemSize const &rhs) const { + return !(*this == rhs); + } + + // Reset covolution mode in the problem + MCTLASS_HOST_DEVICE + Conv3dProblemSize reset_mode(mctlass::conv::Mode mode_) { + Conv3dProblemSize tmp(*this); + tmp.mode = mode_; + return tmp; + } + + // Reset covolution mode in the problem + MCTLASS_HOST_DEVICE + Conv3dProblemSize reset_split_k_slices(int split_k_slices_) { + Conv3dProblemSize tmp(*this); + tmp.split_k_slices = split_k_slices_; + return tmp; + } + + /// Returns activation extent as Tensor5DCoord + MCTLASS_HOST_DEVICE + mctlass::Tensor5DCoord activation_extent() const { + + return mctlass::Tensor5DCoord ({N, D, H, W, C}); + } + + /// Returns filter extent as Tensor5DCoord + MCTLASS_HOST_DEVICE + mctlass::Tensor5DCoord filter_extent() const { + + return mctlass::Tensor5DCoord ({K, T, R, S, C}); + } + + /// Returns output extent as Tensor5DCoord + MCTLASS_HOST_DEVICE + mctlass::Tensor5DCoord output_extent() const { + + return mctlass::Tensor5DCoord ({N, Z, P, Q, K}); + } + + /// Returns activation size in number of elements + MCTLASS_HOST_DEVICE + int64_t activation_size() const { + + return (N * D * H * W * C); + } + + /// Returns filter size in number of elements + MCTLASS_HOST_DEVICE + int64_t filter_size() const { + + return (K * T * R * S * C); + } + + /// Returns output size in number of elements + MCTLASS_HOST_DEVICE + int64_t output_size() const { + + return (N * Z * P * Q * K); + } + + /// Returns output extent as Tensor5DCoord + MCTLASS_HOST_DEVICE + Coord3D padding() const { + + return Coord3D ({pad_d, pad_h, pad_w}); + } + + /// Returns stride as MatrixCoord + MCTLASS_HOST_DEVICE + Coord3D stride() const { + + return Coord3D ({stride_d, stride_h, stride_w}); + } + + /// Returns dilation as MatrixCoord + MCTLASS_HOST_DEVICE + Coord3D dilation() const { + + return Coord3D ({dilation_d, dilation_h, dilation_w}); + } + +}; + + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// ImplicitGemm helper functions // +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Determine the problem size of the implicit GEMM operation +MCTLASS_HOST_DEVICE +mctlass::gemm::GemmCoord implicit_gemm_problem_size( + Operator conv_operator, + Conv3dProblemSize const &problem_size) { + // Compute problem size + switch (conv_operator) { + case Operator::kFprop: + return gemm::GemmCoord( + problem_size.N * problem_size.Z * problem_size.P * problem_size.Q, + problem_size.K, + problem_size.T * problem_size.R * problem_size.S * problem_size.C + ); + case Operator::kDgrad: + return gemm::GemmCoord( + problem_size.N * problem_size.D * problem_size.H * problem_size.W, + problem_size.C, + problem_size.T * problem_size.R * problem_size.S * problem_size.K + ); + case Operator::kWgrad: + return gemm::GemmCoord( + problem_size.K, + problem_size.T * problem_size.R * problem_size.S * problem_size.C, + problem_size.N * problem_size.Z * problem_size.P * problem_size.Q + ); + default: + break; + } + return gemm::GemmCoord(); +} + +// Determine the number of gemm_k iterations for conv2d problem using implicit gemm algorithm +MCTLASS_HOST_DEVICE +int implicit_gemm_k_iterations( + Operator conv_operator, + int threadblock_K, + Conv3dProblemSize const &problem_size, + IteratorAlgorithm algorithm = IteratorAlgorithm::kAnalytic, + GroupMode group_mode = GroupMode::kNone, + int threadblock_N = 0) { + + int iterations = 0; + int elements_per_split_k_slice = 0; + if (group_mode == GroupMode::kNone) { + switch (conv_operator) { + case Operator::kFprop: + elements_per_split_k_slice = (problem_size.C + problem_size.split_k_slices - 1) / problem_size.split_k_slices; + iterations = problem_size.T * problem_size.R * problem_size.S * ((elements_per_split_k_slice + threadblock_K - 1) / threadblock_K); + break; + + case Operator::kDgrad: + elements_per_split_k_slice = (problem_size.K + problem_size.split_k_slices - 1) / problem_size.split_k_slices; + iterations = problem_size.T * problem_size.R * problem_size.S * ((elements_per_split_k_slice + threadblock_K - 1) / threadblock_K); + break; + + case Operator::kWgrad: + elements_per_split_k_slice = (problem_size.N * problem_size.Z * problem_size.P * problem_size.Q + problem_size.split_k_slices - 1) / problem_size.split_k_slices; + iterations = (elements_per_split_k_slice + threadblock_K - 1) / threadblock_K; + break; + + default: + break; + } + } else if (group_mode == GroupMode::kDepthwise) { + int channels_per_cta = threadblock_N; + + if (algorithm == IteratorAlgorithm::kAnalytic) { + switch (conv_operator) { + case Operator::kFprop: + iterations = problem_size.T * problem_size.R * problem_size.S * + ((channels_per_cta + threadblock_K - 1) / threadblock_K); + break; + + default: + break; + } + } + } + + return iterations; +} + +//////////////////////////////////////////////////////////////////////////////// +// Mapping function (ImplicitGemm A, B, C -> Conv Activation, Filter, Output) +//////////////////////////////////////////////////////////////////////////////// +/// Returns ImplicitGemm tensor A extent as Tensor5DCoord +MCTLASS_HOST_DEVICE +mctlass::Tensor5DCoord implicit_gemm_tensor_a_extent( + Operator conv_operator, + Conv3dProblemSize const &problem_size) { + switch (conv_operator) { + case mctlass::conv::Operator::kFprop: return problem_size.activation_extent(); + case mctlass::conv::Operator::kDgrad: return problem_size.output_extent(); + case mctlass::conv::Operator::kWgrad: return problem_size.output_extent(); + default : break; + } + return mctlass::Tensor5DCoord(); +} + +/// Returns ImplicitGemm tensor B extent as Tensor5DCoord +MCTLASS_HOST_DEVICE +mctlass::Tensor5DCoord implicit_gemm_tensor_b_extent( + Operator conv_operator, + Conv3dProblemSize const &problem_size) { + switch (conv_operator) { + case mctlass::conv::Operator::kFprop: return problem_size.filter_extent(); + case mctlass::conv::Operator::kDgrad: return problem_size.filter_extent(); + case mctlass::conv::Operator::kWgrad: return problem_size.activation_extent(); + default : break; + } + return mctlass::Tensor5DCoord(); +} + +/// Returns ImplicitGemm tensor C extent as Tensor5DCoord +MCTLASS_HOST_DEVICE +mctlass::Tensor5DCoord implicit_gemm_tensor_c_extent( + Operator conv_operator, + Conv3dProblemSize const &problem_size) { + switch (conv_operator) { + case mctlass::conv::Operator::kFprop: return problem_size.output_extent(); + case mctlass::conv::Operator::kDgrad: return problem_size.activation_extent(); + case mctlass::conv::Operator::kWgrad: return problem_size.filter_extent(); + default : break; + } + return mctlass::Tensor5DCoord(); +} + +/// Returns ImplicitGemm tensor A size in number of elements +MCTLASS_HOST_DEVICE +int64_t implicit_gemm_tensor_a_size( + Operator conv_operator, + Conv3dProblemSize const &problem_size) { + switch (conv_operator) { + case mctlass::conv::Operator::kFprop: return problem_size.activation_size(); + case mctlass::conv::Operator::kDgrad: return problem_size.output_size(); + case mctlass::conv::Operator::kWgrad: return problem_size.output_size(); + default : break; + } + return 0; +} + +/// Returns ImplicitGemm tensor B size in number of elements +MCTLASS_HOST_DEVICE +int64_t implicit_gemm_tensor_b_size( + Operator conv_operator, + Conv3dProblemSize const &problem_size) { + switch (conv_operator) { + case mctlass::conv::Operator::kFprop: return problem_size.filter_size(); + case mctlass::conv::Operator::kDgrad: return problem_size.filter_size(); + case mctlass::conv::Operator::kWgrad: return problem_size.activation_size(); + default : break; + } + return 0; +} + +/// Returns ImplicitGemm tensor C size in number of elements +MCTLASS_HOST_DEVICE +int64_t implicit_gemm_tensor_c_size( + Operator conv_operator, + Conv3dProblemSize const &problem_size) { + switch (conv_operator) { + case mctlass::conv::Operator::kFprop: return problem_size.output_size(); + case mctlass::conv::Operator::kDgrad: return problem_size.activation_size(); + case mctlass::conv::Operator::kWgrad: return problem_size.filter_size(); + default : break; + } + return 0; +} + +} // namespace conv +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/convolution.h b/csrc/mctlass/include/mctlass/conv/convolution.h new file mode 100644 index 0000000..7f666f6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/convolution.h @@ -0,0 +1,168 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief + +This file contains definitions and utility functions for describing convolution problem sizes in terms of +activation (NHWC), filter (KRSC), output (NPQK), padding (pad_h, pad_w), stride (stride_h, stride_w), and +dilation (dilation_h, dilation_w). Furthermore, it defines helper functions to map CUTLASS's implicit gemm +tensor extents, sizes, and data types to that of the convolution's extents, sizes, and data types. + + * Mapping convolutions to Gemm computation * + +Cutlass implements convolutions with the Implicit Gemm algorithm. This algorithm performs a gemm +(general matrix-matrix multiply) on the convolution tensors Activation, Filter, and Output. +The underlying gemm operation follows the standard gemm definition: + + C = A * B + C + + A and B are input matrices + C is source and output matrix + + +For the three convolutional operators (Fprop, Dgrad, Wgrad), ImplicitGemm matrices A, B, and C are mapped +to convolution tensors Activation, Filter and Output as described in the table below. + + ___________________________________________________________________________ + ConvolutionalOperator | A | B | C + ___________________________________________________________________________ + | | | | | + | Fprop | Activation | Filter | Output | + | Dgrad | Output | Filter | Activation | + | Wgrad | Output | Activation | Filter | + ___________________________________________________________________________ + +In convolution codebase, DO NOT mix using (A, B, C) with (Activation, Filter, Output). + +For example, it's confusing and error prone to document a convolution class or function +as operating on "A, B, Output." Instead, use the mapping functions below, +and adhere to using either A, B, C or Activation, Filter, Output. + +Map elements' data types (ImplicitGemm -> Conv): GemmToConvElementMap +Map elements' data types (Conv -> ImplicitGemm): ConvToGemmElementMap +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/tensor_coord.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" + +namespace mctlass { +namespace conv { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Convolutional operator +enum class Operator { + kFprop, + kDgrad, + kWgrad +}; + +/// Distinguishes convolution from cross correlation +enum class Mode { + kCrossCorrelation, + kConvolution +}; + +/// Selects among several implementation variants trading off performance with simplicity +enum class IteratorAlgorithm { + kAnalytic, ///< functionally correct in all cases but lower performance + kOptimized, ///< optimized for R <= 32, S <= 32 and unity-stride dgrad + kFixedChannels, ///< Analytic algorithm optimized for fixed channel count (C == AccessSize) + kFewChannels, ///< Analytic algorithm optimized for few channels (C divisible by AccessSize) + kFixedStrideDilation ///< Optimized for fixed stride and dilation +}; + +/// Distinguishes among partial specializations that accelerate certain problems where convolution +/// stride is unit. +enum class StrideSupport { + kStrided, ///< arbitrary convolution stride + kUnity, ///< unit convolution stride + kFixed ///< fixed convolution stride +}; + +/// Identifies split-K mode +enum class SplitKMode { + kNone, + kSerial, + kParallel +}; + +/// Identifies group mode +enum class GroupMode { + kNone, + kSingleGroup, ///< One CTA calculates one group or less + kMultipleGroup, ///< One CTA calculates multiple groups + kDepthwise ///< One CTA calculates cta_n groups (problem_size.C == problem_size.K == problem_size.groups) +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Shape of a tensor +template < + int N = 1, + int H = 1, + int W = 1, + int C = 1 +> +struct TensorNHWCShape { + static int const kN = N; + static int const kH = H; + static int const kW = W; + static int const kC = C; + + static int const kHW = H * W; + static int const kNHW = N * kHW; + static int const kNHWC = N * H * W * C; + + static int const kCount = kNHWC; + + // + // Static member functions + // + + /// Returns a Coord object + MCTLASS_HOST_DEVICE + static Coord<4> toCoord() { + return make_Coord(kN, kH, kW, kC); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace conv +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/device/direct_convolution.h b/csrc/mctlass/include/mctlass/conv/device/direct_convolution.h new file mode 100644 index 0000000..64c332a --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/device/direct_convolution.h @@ -0,0 +1,269 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/* \file + \brief Template for device-level Depthwise Convolution +*/ + +#pragma once + +#include + +#include "mctlass/mctlass.h" +#include "mctlass/device_kernel.h" +#include "mctlass/conv/convolution.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +class DirectConvolution { +public: + + using UnderlyingKernel = DirectConvolutionKernel_; + + using ElementA = typename UnderlyingKernel::ElementA; + using LayoutA = typename UnderlyingKernel::LayoutA; + using ElementB = typename UnderlyingKernel::ElementB; + using LayoutB = typename UnderlyingKernel::LayoutB; + using ElementC = typename UnderlyingKernel::ElementC; + using LayoutC = typename UnderlyingKernel::LayoutC; + using ElementAccumulator = typename UnderlyingKernel::ElementAccumulator; + using ElementCompute = typename UnderlyingKernel::ElementCompute; + using OperatorClass = typename UnderlyingKernel::OperatorClass; + using ArchTag = typename UnderlyingKernel::ArchTag; + using ThreadblockShape = typename UnderlyingKernel::ThreadblockShape; + using WarpShape = typename UnderlyingKernel::WarpShape; + using InstructionShape = typename UnderlyingKernel::InstructionShape; + using ThreadblockSwizzle = typename UnderlyingKernel::ThreadblockSwizzle; + using EpilogueOutputOp = typename UnderlyingKernel::EpilogueOutputOp; + static int const kStages = UnderlyingKernel::kStages; + static int const kConvDim = UnderlyingKernel::kConvDim; + using WarpMmaOperator = typename UnderlyingKernel::WarpMmaOperator; + using ArchMmaOperator = typename UnderlyingKernel::ArchMmaOperator; + using MathOperator = typename UnderlyingKernel::MathOperator; + + static mctlass::conv::Operator const kConvolutionalOperator = UnderlyingKernel::kConvolutionalOperator; + static mctlass::conv::IteratorAlgorithm const kIteratorAlgorithm = UnderlyingKernel::kIteratorAlgorithm; + static mctlass::conv::StrideSupport const kStrideSupport = UnderlyingKernel::kStrideSupport; + static mctlass::conv::GroupMode const kGroupMode = UnderlyingKernel::kGroupMode; + + static int const kWarpCount = + (ThreadblockShape::kM / WarpShape::kM) * + (ThreadblockShape::kN / WarpShape::kN) * + (ThreadblockShape::kK / WarpShape::kK); + + /// Argument structure + using Arguments = typename UnderlyingKernel::Arguments; + + using ReorderKernel = typename UnderlyingKernel::ReorderKernel; + + private: + + /// Kernel parameters object + typename UnderlyingKernel::Params params_; + +public: + + /// Constructs Implicit GEMM + DirectConvolution() { } + + /// Determines whether the Implicit GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + // dispatch to iterators + Status status = UnderlyingKernel::Mma::IteratorA::can_implement(args.problem_size); + if (Status::kSuccess != status) { + return status; + } + + status = UnderlyingKernel::Mma::IteratorB::can_implement(args.problem_size); + if (Status::kSuccess != status) { + return status; + } + + if (kGroupMode != conv::GroupMode::kDepthwise) { + return Status::kErrorInvalidProblem; + } + + // C and K should be multiple of groups + if (args.problem_size.K != args.problem_size.groups && + args.problem_size.C != args.problem_size.groups) { + return Status::kErrorInvalidProblem; + } + + + static int const kAlignmentC = UnderlyingKernel::Epilogue::OutputTileIterator::kElementsPerAccess; + if (kConvolutionalOperator == conv::Operator::kFprop) { + if (args.problem_size.K % kAlignmentC) + return Status::kErrorMisalignedOperand; + } else if (kConvolutionalOperator == conv::Operator::kDgrad) { + if (args.problem_size.C % kAlignmentC) + return Status::kErrorMisalignedOperand; + } else if (kConvolutionalOperator == conv::Operator::kWgrad) { + if (args.problem_size.C % kAlignmentC) + return Status::kErrorMisalignedOperand; + } + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape( + threadblock_swizzle.get_tiled_shape( + kConvolutionalOperator, + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.problem_size.split_k_slices)); + + if (!(grid.y <= std::numeric_limits::max() && + grid.z <= std::numeric_limits::max())) { + + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + return 0; + } + + /// Initializes GEMM state from arguments. + Status initialize( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + // initialize the params structure from the arguments + params_ = typename UnderlyingKernel::Params( + args, + static_cast(workspace) + ); + + int smem_size = int(sizeof(typename UnderlyingKernel::SharedStorage)); + + if (smem_size >= (48 << 10)) { + mcError_t result = mcFuncSetAttribute(mctlass::Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + /// Initializes GEMM state from arguments. + Status update(Arguments const &args, void *workspace = nullptr) { + + // update the params structure from the arguments + params_.ptr_A = args.ref_A.data(); + params_.ptr_B = args.ref_B.data(); + params_.ptr_C = args.ref_C.data(); + params_.ptr_D = args.ref_D.data(); + params_.output_op = args.output_op; + params_.ptr_reordered_B = args.ref_reordered_B.data();; + params_.semaphore = static_cast(workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + // Launch reorder kernel + if (params_.ptr_reordered_B != nullptr) { + dim3 grid = ReorderKernel::get_grid_shape(params_); + dim3 block = ReorderKernel::get_block_shape(); + + mctlass::Kernel<<>>(params_); + } + + // Launch main kernel + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(32 * kWarpCount, 1, 1); + + // Dynamic SMEM size based on input params. + int smem_size = int(params_.get_smem_size()); + + // Make sure we can use that much shared memory. + mcError_t status = + mcFuncSetAttribute(mctlass::Kernel, mcFuncAttributeMaxDynamicSharedMemorySize, smem_size); + if (status != mcSuccess) + return Status::kErrorInternal; + + + mctlass::Kernel<<>>(params_); + + mcError_t result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } + + int get_smem_size() { return int(params_.get_smem_size()); } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} +} +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/device/implicit_gemm_convolution.h b/csrc/mctlass/include/mctlass/conv/device/implicit_gemm_convolution.h new file mode 100644 index 0000000..6fe1fa2 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/device/implicit_gemm_convolution.h @@ -0,0 +1,328 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/* \file + \brief Template for device-level Implicit GEMM Convolution +*/ + +#pragma once + +#include + +#include "mctlass/mctlass.h" +#include "mctlass/device_kernel.h" +#include "mctlass/conv/convolution.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +class ImplicitGemmConvolution { +public: + + using UnderlyingKernel = ImplicitGemmKernel_; + + using ElementA = typename UnderlyingKernel::ElementA; + using LayoutA = typename UnderlyingKernel::LayoutA; + using ElementB = typename UnderlyingKernel::ElementB; + using LayoutB = typename UnderlyingKernel::LayoutB; + using ElementC = typename UnderlyingKernel::ElementC; + using LayoutC = typename UnderlyingKernel::LayoutC; + using ElementAccumulator = typename UnderlyingKernel::ElementAccumulator; + using ElementCompute = typename UnderlyingKernel::ElementCompute; + using OperatorClass = typename UnderlyingKernel::OperatorClass; + using ArchTag = typename UnderlyingKernel::ArchTag; + using ThreadblockShape = typename UnderlyingKernel::ThreadblockShape; + using WarpShape = typename UnderlyingKernel::WarpShape; + using InstructionShape = typename UnderlyingKernel::InstructionShape; + using ThreadblockSwizzle = typename UnderlyingKernel::ThreadblockSwizzle; + using EpilogueOutputOp = typename UnderlyingKernel::EpilogueOutputOp; + static int const kStages = UnderlyingKernel::kStages; + static int const kConvDim = UnderlyingKernel::kConvDim; + using WarpMmaOperator = typename UnderlyingKernel::WarpMmaOperator; + using ArchMmaOperator = typename UnderlyingKernel::ArchMmaOperator; + using MathOperator = typename UnderlyingKernel::MathOperator; + + static mctlass::conv::Operator const kConvolutionalOperator = UnderlyingKernel::kConvolutionalOperator; + static mctlass::conv::IteratorAlgorithm const kIteratorAlgorithm = UnderlyingKernel::kIteratorAlgorithm; + static mctlass::conv::StrideSupport const kStrideSupport = UnderlyingKernel::kStrideSupport; + static mctlass::conv::GroupMode const kGroupMode = UnderlyingKernel::kGroupMode; + + static int const kWarpCount = + (ThreadblockShape::kM / WarpShape::kM) * + (ThreadblockShape::kN / WarpShape::kN) * + (ThreadblockShape::kK / WarpShape::kK); + + /// Argument structure + using Arguments = typename UnderlyingKernel::Arguments; + +private: + + /// Kernel parameters object + typename UnderlyingKernel::Params params_; + +public: + + /// Constructs Implicit GEMM + ImplicitGemmConvolution() { } + + /// Determines whether the Implicit GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + // dispatch to iterators + Status status = UnderlyingKernel::Mma::IteratorA::can_implement(args.problem_size); + if (Status::kSuccess != status) { + return status; + } + + status = UnderlyingKernel::Mma::IteratorB::can_implement(args.problem_size); + if (Status::kSuccess != status) { + return status; + } + + // check group conv constraint + if (args.problem_size.groups != 1) { + if (kGroupMode == conv::GroupMode::kNone) { + return Status::kErrorInvalidProblem; + } + + // C and K should be multiple of groups + if (args.problem_size.K % args.problem_size.groups || + args.problem_size.C % args.problem_size.groups) { + return Status::kErrorInvalidProblem; + } + + // split-k is not supported + if (args.problem_size.split_k_slices != 1) { + return Status::kErrorInvalidProblem; + } + + int k_per_group = args.problem_size.K / args.problem_size.groups; + // k_per_group should be multiple of ThreadblockShape N, one CTA calculate one group + if (kGroupMode == conv::GroupMode::kSingleGroup && k_per_group % ThreadblockShape::kN) { + return Status::kErrorInvalidProblem; + } + // ThreadblockShape::kN should be divisible by k_per_group, one CTA calculate multiple groups + if (kGroupMode == conv::GroupMode::kMultipleGroup && ThreadblockShape::kN % k_per_group) { + return Status::kErrorInvalidProblem; + } + + // current optimized iterator algo only supports SingleGroup mode + if (kIteratorAlgorithm == IteratorAlgorithm::kOptimized && + kGroupMode != conv::GroupMode::kSingleGroup) { + return Status::kErrorInvalidProblem; + } + } + + static int const kAlignmentC = UnderlyingKernel::Epilogue::OutputTileIterator::kElementsPerAccess; + if (kConvolutionalOperator == conv::Operator::kFprop) { + if (args.problem_size.K % kAlignmentC) + return Status::kErrorMisalignedOperand; + } else if (kConvolutionalOperator == conv::Operator::kDgrad) { + if (args.problem_size.C % kAlignmentC) + return Status::kErrorMisalignedOperand; + } else if (kConvolutionalOperator == conv::Operator::kWgrad) { + if (args.problem_size.C % kAlignmentC) + return Status::kErrorMisalignedOperand; + } + + // check for unsupported problem sizes for strided dgrad implementation + if (kConvolutionalOperator == conv::Operator::kDgrad && + kStrideSupport == conv::StrideSupport::kStrided) { + + // split-k (serial or parallel) is not supported for strided dgrad + if(args.problem_size.split_k_slices > 1) { + return Status::kErrorNotSupported; + } + + // dilation > {1x1} is not supported for strided dgrad + if(args.problem_size.dilation_h > 1 || args.problem_size.dilation_w > 1) { + return Status::kErrorNotSupported; + } + } + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape( + threadblock_swizzle.get_tiled_shape( + kConvolutionalOperator, + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.problem_size.split_k_slices)); + + if (!(grid.y <= std::numeric_limits::max() && + grid.z <= std::numeric_limits::max())) { + + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + size_t workspace_bytes = 0; + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + kConvolutionalOperator, + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.problem_size.split_k_slices); + + if(args.split_k_mode == SplitKMode::kParallel) { + + // Split-K parallel: CTAs in k-dimension write the partial results in a temporary workspace. + // The user needs to call a reduction operator to optain the final output tensor + workspace_bytes = + sizeof(ElementAccumulator) * + size_t(mctlass::conv::implicit_gemm_tensor_c_size(kConvolutionalOperator, args.problem_size)) * + size_t(grid_tiled_shape.k()); + } + + else if(args.split_k_mode == SplitKMode::kSerial && args.problem_size.split_k_slices > 1) { + + // Split-K serial: The user workspace is used to store semaphore and serialize writing the + // final reduced output to user's output tensor + workspace_bytes = sizeof(int) * size_t(grid_tiled_shape.m()) * size_t(grid_tiled_shape.n()); + } + + return workspace_bytes; + } + + /// Initializes GEMM state from arguments. + Status initialize( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + if (args.problem_size.split_k_slices > 1) { + + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + mcError_t status = mcMemsetAsync(workspace, 0, get_workspace_size(args), stream); + + if (status != mcSuccess) { + return Status::kErrorInternal; + } + } + + // initialize the params structure from the arguments + params_ = typename UnderlyingKernel::Params( + args, + static_cast(workspace) + ); + + int smem_size = int(sizeof(typename UnderlyingKernel::SharedStorage)); + + if (smem_size >= (64 << 10)) { + mcError_t result = mcFuncSetAttribute(mctlass::Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + /// Initializes GEMM state from arguments. + Status update(Arguments const &args, void *workspace = nullptr) { + + // update the params structure from the arguments + params_.ptr_A = args.ref_A.data(); + params_.ptr_B = args.ref_B.data(); + params_.ptr_C = args.ref_C.data(); + params_.ptr_D = args.ref_D.data(); + params_.output_op = args.output_op; + params_.semaphore = static_cast(workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(64 * kWarpCount, 1, 1); + + int smem_size = int(sizeof(typename UnderlyingKernel::SharedStorage)); + + mctlass::Kernel<<>>(params_); + + mcError_t result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} +} +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/device/implicit_gemm_convolution_fusion.h b/csrc/mctlass/include/mctlass/conv/device/implicit_gemm_convolution_fusion.h new file mode 100644 index 0000000..cf01fff --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/device/implicit_gemm_convolution_fusion.h @@ -0,0 +1,268 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/* \file + \brief Template for device-level fused activation's scale+bias+relu and Implicit GEMM Convolution +*/ + +#pragma once + +#include + +#include "mctlass/mctlass.h" +#include "mctlass/device_kernel.h" +#include "mctlass/conv/convolution.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +class ImplicitGemmConvolutionFusion { +public: + + using ImplicitGemmFusionKernel = ImplicitGemmFusionKernel_; + + using ElementA = typename ImplicitGemmFusionKernel::ElementA; + using LayoutA = typename ImplicitGemmFusionKernel::LayoutA; + using ElementB = typename ImplicitGemmFusionKernel::ElementB; + using LayoutB = typename ImplicitGemmFusionKernel::LayoutB; + +// using ElementScaleBias = typename ImplicitGemmFusionKernel::ElementScaleBias; +// using LayoutScaleBias = typename ImplicitGemmFusionKernel::LayoutScaleBias; + + using ElementC = typename ImplicitGemmFusionKernel::ElementC; + using LayoutC = typename ImplicitGemmFusionKernel::LayoutC; + using ElementAccumulator = typename ImplicitGemmFusionKernel::ElementAccumulator; + using ElementCompute = typename ImplicitGemmFusionKernel::ElementCompute; + using OperatorClass = typename ImplicitGemmFusionKernel::OperatorClass; + using ArchTag = typename ImplicitGemmFusionKernel::ArchTag; + using ThreadblockShape = typename ImplicitGemmFusionKernel::ThreadblockShape; + using WarpShape = typename ImplicitGemmFusionKernel::WarpShape; + using InstructionShape = typename ImplicitGemmFusionKernel::InstructionShape; + using ThreadblockSwizzle = typename ImplicitGemmFusionKernel::ThreadblockSwizzle; + using EpilogueOutputOp = typename ImplicitGemmFusionKernel::EpilogueOutputOp; + static int const kStages = ImplicitGemmFusionKernel::kStages; + static int const kConvDim = ImplicitGemmFusionKernel::kConvDim; + using WarpMmaOperator = typename ImplicitGemmFusionKernel::WarpMmaOperator; + using ArchMmaOperator = typename ImplicitGemmFusionKernel::ArchMmaOperator; + using MathOperator = typename ImplicitGemmFusionKernel::MathOperator; + + static mctlass::conv::Operator const kConvolutionalOperator = ImplicitGemmFusionKernel::kConvolutionalOperator; + static mctlass::conv::IteratorAlgorithm const kIteratorAlgorithm = ImplicitGemmFusionKernel::kIteratorAlgorithm; + + static int const kWarpCount = + (ThreadblockShape::kM / WarpShape::kM) * + (ThreadblockShape::kN / WarpShape::kN) * + (ThreadblockShape::kK / WarpShape::kK); + + /// Argument structure + using Arguments = typename ImplicitGemmFusionKernel::Arguments; + +private: + + /// Kernel parameters object + typename ImplicitGemmFusionKernel::Params params_; + +public: + + /// Constructs Implicit GEMM + ImplicitGemmConvolutionFusion() { } + + /// Determines whether the Implicit GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + // dispatch to iterators + Status status = ImplicitGemmFusionKernel::Mma::IteratorA::can_implement(args.problem_size); + if (Status::kSuccess != status) { + return status; + } + + status = ImplicitGemmFusionKernel::Mma::IteratorB::can_implement(args.problem_size); + if (Status::kSuccess != status) { + return status; + } + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape( + threadblock_swizzle.get_tiled_shape( + mctlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, args.problem_size), + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.problem_size.split_k_slices)); + + if (!(grid.y <= std::numeric_limits::max() && + grid.z <= std::numeric_limits::max())) { + + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + size_t workspace_bytes = 0; + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + mctlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, args.problem_size), + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.problem_size.split_k_slices); + + if(args.split_k_mode == SplitKMode::kParallel) { + + // Split-K parallel: CTAs in k-dimension write the partial results in a temporary workspace. + // The user needs to call a reduction operator to optain the final output tensor + workspace_bytes = + sizeof(ElementAccumulator) * + size_t(mctlass::conv::implicit_gemm_tensor_c_size(kConvolutionalOperator, args.problem_size)) * + size_t(grid_tiled_shape.k()); + } + + else if(args.split_k_mode == SplitKMode::kSerial && args.problem_size.split_k_slices > 1) { + + // Split-K serial: The user workspace is used to store semaphore and serialize writing the + // final reduced output to user's output tensor + workspace_bytes = sizeof(int) * size_t(grid_tiled_shape.m()) * size_t(grid_tiled_shape.n()); + } + + return workspace_bytes; + } + + /// Initializes GEMM state from arguments. + Status initialize( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + if (args.problem_size.split_k_slices > 1) { + + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + mcError_t status = mcMemsetAsync(workspace, 0, get_workspace_size(args), stream); + + if (status != mcSuccess) { + return Status::kErrorInternal; + } + } + + // initialize the params structure from the arguments + params_ = typename ImplicitGemmFusionKernel::Params( + args, + static_cast(workspace) + ); + + int smem_size = int(sizeof(typename ImplicitGemmFusionKernel::SharedStorage)); + + if (smem_size >= (48 << 10)) { + mcError_t result = mcFuncSetAttribute(mctlass::Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + /// Initializes Impicit GEMM state from arguments. + Status update(Arguments const &args, void *workspace = nullptr) { + + // update the params structure from the arguments + params_.ptr_A = args.ref_A.data(); + params_.ptr_B = args.ref_B.data(); + params_.ptr_scale = args.ref_A_scale.data(); + params_.ptr_bias = args.ref_A_bias.data(); + params_.ptr_C = args.ref_C.data(); + params_.ptr_D = args.ref_D.data(); + params_.output_op = args.output_op; + params_.semaphore = static_cast(workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(32 * kWarpCount, 1, 1); + + int smem_size = int(sizeof(typename ImplicitGemmFusionKernel::SharedStorage)); + + mctlass::Kernel<<>>(params_); + + mcError_t result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} +} +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d.h b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d.h new file mode 100644 index 0000000..999bcdf --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d.h @@ -0,0 +1,272 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level implicit GEMM convolution definitions for threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/conv/threadblock/threadblock_swizzle.h" +#include "mctlass/epilogue/threadblock/default_epilogue_simt.h" +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_with_broadcast.h" +#include "mctlass/epilogue/threadblock/default_epilogue_with_reduction.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/threadblock/conv2d_tile_iterator.h" +#include "mctlass/conv/threadblock/implicit_gemm_pipelined.h" +#include "mctlass/conv/threadblock/implicit_gemm_multistage.h" +#include "mctlass/conv/threadblock/implicit_gemm_fprop_fusion_multistage.h" +#include "mctlass/conv/threadblock/implicit_gemm_wgrad_fusion_multistage.h" +#include "mctlass/conv/kernel/implicit_gemm_convolution.h" +#include "mctlass/conv/kernel/implicit_gemm_convolution_fusion.h" +#include "mctlass/conv/kernel/implicit_gemm_convolution_strided_dgrad.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +template < + typename ArchTag, + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename OutputOp +> +struct DefaultConvEpilogue { + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputOp, + OutputOp::kCount + >::Epilogue; +}; + +template < + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename OutputOp +> +struct DefaultConvEpilogue< + arch::Sm70, + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputOp +> { + + using Epilogue = typename epilogue::threadblock::DefaultEpilogueVoltaTensorOp< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputOp, + OutputOp::kCount + >::Epilogue; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ArchTag, + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename ElementOutput, + typename ElementTensor, + typename ElementVector, + typename OutputOp, + int ElementsPerAccess +> +struct DefaultConvEpilogueWithBroadcastTensorOp { + using Epilogue = typename epilogue::threadblock::DefaultEpilogueWithBroadcastTensorOp< + Shape, + WarpMmaTensorOp, + PartitionsK, + ElementOutput, + ElementTensor, + ElementVector, + OutputOp, + ElementsPerAccess + >::Epilogue; +}; + +template < + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename ElementOutput, + typename ElementTensor, + typename ElementVector, + typename OutputOp, + int ElementsPerAccess +> +struct DefaultConvEpilogueWithBroadcastTensorOp< + arch::Sm70, + Shape, + WarpMmaTensorOp, + PartitionsK, + ElementOutput, + ElementTensor, + ElementVector, + OutputOp, + ElementsPerAccess + > { + using Epilogue = typename epilogue::threadblock::DefaultEpilogueWithBroadcastVoltaTensorOp< + Shape, + WarpMmaTensorOp, + PartitionsK, + ElementOutput, + ElementTensor, + ElementVector, + OutputOp, + ElementsPerAccess + >::Epilogue; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ArchTag, + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename ElementOutput, + typename OutputOp, + typename ReductionOp, + int ElementsPerAccess +> +struct DefaultConvEpilogueWithReductionTensorOp { + using Epilogue = typename epilogue::threadblock::DefaultEpilogueWithReductionTensorOp< + Shape, + WarpMmaTensorOp, + PartitionsK, + ElementOutput, + OutputOp, + ReductionOp, + ElementsPerAccess + >::Epilogue; +}; + +template < + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename ElementOutput, + typename OutputOp, + typename ReductionOp, + int ElementsPerAccess +> +struct DefaultConvEpilogueWithReductionTensorOp< + arch::Sm70, + Shape, + WarpMmaTensorOp, + PartitionsK, + ElementOutput, + OutputOp, + ReductionOp, + ElementsPerAccess + > { + using Epilogue = typename epilogue::threadblock::DefaultEpilogueWithReductionVoltaTensorOp< + Shape, + WarpMmaTensorOp, + PartitionsK, + ElementOutput, + OutputOp, + ReductionOp, + ElementsPerAccess + >::Epilogue; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Defaults for strided Dgrad +template < + typename ArchTag, + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename OutputOp +> +struct DefaultConvEpilogueStridedDgrad { + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOpStridedDgrad< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputOp, + OutputOp::kCount + >::Epilogue; +}; + +template < + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename OutputOp +> +struct DefaultConvEpilogueStridedDgrad< + arch::Sm70, + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputOp +> { + + using Epilogue = typename epilogue::threadblock::DefaultEpilogueVoltaTensorOpStridedDgrad< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputOp, + OutputOp::kCount + >::Epilogue; +}; + +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_dgrad.h b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_dgrad.h new file mode 100644 index 0000000..d4b41c1 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_dgrad.h @@ -0,0 +1,1927 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level implicit GEMM convolution definitions combine threadblock-scoped + matrix multiply-add with the appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/conv/kernel/default_conv2d.h" + +#include "mctlass/conv/threadblock/conv2d_dgrad_output_gradient_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_dgrad_output_gradient_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv2d_dgrad_filter_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_dgrad_filter_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv2d_tile_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv2dDgrad +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided, + /// Access granularity of A matrix in units of elements + int AlignmentA = 128 / mctlass::sizeof_bits::value, + /// Access granularity of B matrix in units of elements + int AlignmentB = 128 / mctlass::sizeof_bits::value +> struct DefaultConv2dDgrad; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassTensorOp convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dDgrad specialization for Analytic IteratorAlgorithm Dgrad Strided and +// multistage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport::kStrided, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + StrideSupport::kStrided, + AccessTypeA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + StrideSupport::kStrided, + AccessTypeB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * AlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + CacheOpB, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOpStridedDgrad< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionStridedDgrad< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; +}; + +/// Defines a kernel for Conv2dDgrad specialization for Analytic IteratorAlgorithm Dgrad Strided +// and 2 stage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport::kStrided, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::TileIteratorStridedDgrad< + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + StrideSupport::kStrided, + AccessTypeA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::TileIteratorStridedDgrad< + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + StrideSupport::kStrided, + AccessTypeB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogueStridedDgrad< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionStridedDgrad< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dDgrad specialization for Analytic IteratorAlgorithm Dgrad Unity Strided +// and multistage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport::kUnity, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + StrideSupport::kUnity, + AccessTypeA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + StrideSupport::kUnity, + AccessTypeB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * AlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + CacheOpB, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; +}; + +/// Defines a kernel for Conv2dDgrad specialization for Analytic IteratorAlgorithm Dgrad Unity +// 2 stage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport::kUnity, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + StrideSupport::kUnity, + AccessTypeA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + StrideSupport::kUnity, + AccessTypeB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dDgrad specialization for optimized IteratorAlgorithm Dgrad Unity Strided +// and multistage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport::kUnity, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + StrideSupport::kUnity, + AccessTypeA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + StrideSupport::kUnity, + AccessTypeB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * AlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + CacheOpB, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; +}; + +/// Defines a kernel for Conv2dDgrad specialization for Optimized IteratorAlgorithm Dgrad Strided and +// multistage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport::kStrided, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + StrideSupport::kStrided, + AccessTypeA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + StrideSupport::kStrided, + AccessTypeB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * AlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + CacheOpB, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOpStridedDgrad< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionStridedDgrad< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; +}; + +/// Defines a kernel for Conv2dDgrad specialization for Optimized IteratorAlgorithm Dgrad Strided +// and 2 stage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport::kStrided, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::TileIteratorStridedDgrad< + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + StrideSupport::kStrided, + AccessTypeA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::TileIteratorStridedDgrad< + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + StrideSupport::kStrided, + AccessTypeB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogueStridedDgrad< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionStridedDgrad< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; +}; + +/// Defines a kernel for Conv2dDgrad specialization for Optimized IteratorAlgorithm Dgrad Unity +// 2 stage pipeline +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport::kUnity, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + StrideSupport::kUnity, + AccessTypeA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + StrideSupport::kUnity, + AccessTypeB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassSimt convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv2dDgrad specialization for Analytic IteratorAlgorithm, +/// multi-stage pipeline, and FFMA-based mainloop for SM80 + +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + conv::StrideSupport::kUnity, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + conv::StrideSupport::kUnity + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + conv::StrideSupport::kUnity + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + conv::StrideSupport::kStrided, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + conv::StrideSupport::kStrided + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + conv::StrideSupport::kStrided + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimtStridedDgrad< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionStridedDgrad< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dDgrad specialization for Optimized IteratorAlgorithm, +/// multi-stage pipeline, and FFMA-based mainloop for SM80 + +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport::kUnity, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + StrideSupport::kUnity + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + StrideSupport::kUnity + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + conv::StrideSupport::kStrided, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + conv::StrideSupport::kStrided + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + conv::StrideSupport::kStrided + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimtStridedDgrad< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionStridedDgrad< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; + +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dDgrad specialization for Analytic IteratorAlgorithm, +/// 2 stage pipeline, and FFMA-based mainloop for SM50 +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + conv::StrideSupport::kUnity, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + conv::StrideSupport::kUnity + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + conv::StrideSupport::kUnity + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; + +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + conv::StrideSupport::kStrided, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIteratorStridedDgrad< + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + conv::StrideSupport::kStrided + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIteratorStridedDgrad< + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + conv::StrideSupport::kStrided + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimtStridedDgrad< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionStridedDgrad< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dDgrad specialization for Optimized IteratorAlgorithm, +/// 2 stage pipeline, and FFMA-based mainloop for SM50 +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport::kUnity, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + StrideSupport::kUnity + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + StrideSupport::kUnity + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + conv::StrideSupport::kStrided, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIteratorStridedDgrad< + mctlass::conv::threadblock::Conv2dDgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + conv::StrideSupport::kStrided + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIteratorStridedDgrad< + mctlass::conv::threadblock::Conv2dDgradFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + conv::StrideSupport::kStrided + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimtStridedDgrad< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionStridedDgrad< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad + >; + +}; + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop.h b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop.h new file mode 100644 index 0000000..3597343 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop.h @@ -0,0 +1,1989 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level implicit GEMM convolution definitions combine threadblock-scoped + matrix multiply-add with the appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/conv/kernel/default_conv2d.h" + +#include "mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_fixed_channels.h" +#include "mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_few_channels.h" + +#include "mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_fixed_channels.h" +#include "mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_few_channels.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv2dFprop +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided, + /// Access granularity of A matrix in units of elements + int AlignmentA = 128 / mctlass::sizeof_bits::value, + /// Access granularity of B matrix in units of elements + int AlignmentB = 128 / mctlass::sizeof_bits::value +> struct DefaultConv2dFprop; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassTensorOp convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Analytic IteratorAlgorithm and multistage +/// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA, + AccessTypeA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB, + AccessTypeB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * AlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + CacheOpB, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Analytic IteratorAlgorithm and multistage +/// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kFixedChannels, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorFixedChannels< + mctlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA, + AccessTypeA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorFixedChannels< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB, + AccessTypeB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * AlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + CacheOpB, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Analytic IteratorAlgorithm and two stage +/// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kFixedChannels, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorFixedChannels< + mctlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA, + AccessTypeA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorFixedChannels< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB, + AccessTypeB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Analytic IteratorAlgorithm and multistage +/// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kFewChannels, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorFewChannels< + mctlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA, + AccessTypeA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorFewChannels< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB, + AccessTypeB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * AlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + CacheOpB, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +/// Defines a kernel for Conv2dFprop specialization for Analytic IteratorAlgorithm and multistage +/// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kFewChannels, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorFewChannels< + mctlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA, + AccessTypeA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorFewChannels< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB, + AccessTypeB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * AlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Analytic IteratorAlgorithm and multistage +/// pipeline with interleaved layout. +template < + typename ElementA, + typename ElementB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB, + int InterleavedK +> +struct DefaultConv2dFprop < + ElementA, + layout::TensorNCxHWx, + ElementB, + layout::TensorCxRSKx, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajorInterleaved, + ElementB, layout::RowMajorInterleaved, + ElementAccumulator, LayoutC, arch::OpClassTensorOp, + Stages, MathOperatorTag, true>; + + // Define iterators over tiles from the A operand + // Note GEMM shared memory threadmap is used here because conv global memory + // layout needs to be mapped to fprop which is similar to the crosswise + // layout which is used by the interleaved GEMM shared memory threadmap. + // The Interleaved GEMM global memory layout is similar to the congruous + // layout. + using ThreadMapA = typename MmaCore::SmemThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, layout::TensorNCxHWx, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + // Note GEMM shared memory threadmap is used here because conv global memory + // layout needs to be mapped to fprop which is similar to the crosswise + // layout which is used by the interleaved GEMM shared memory threadmap. + // The Interleaved GEMM global memory layout is similar to the congruous + // layout. + using ThreadMapB = typename MmaCore::SmemThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, layout::TensorCxRSKx, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Global, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultInterleavedConvEpilogue< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount, + InterleavedK + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Analytic IteratorAlgorithm +/// and 2 stage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA, + AccessTypeA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB, + AccessTypeB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Analytic IteratorAlgorithm and 2 stage +/// pipeline with interleaved layout. +template < + typename ElementA, + typename ElementB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB, + int InterleavedK +> +struct DefaultConv2dFprop < + ElementA, + layout::TensorNCxHWx, + ElementB, + layout::TensorCxRSKx, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajorInterleaved, + ElementB, layout::RowMajorInterleaved, + ElementAccumulator, LayoutC, arch::OpClassTensorOp, + 2, MathOperatorTag, true>; + + // Define iterators over tiles from the A operand + // Note GEMM shared memory threadmap is used here because conv global memory + // layout needs to be mapped to fprop which is similar to the crosswise + // layout which is used by the interleaved GEMM shared memory threadmap. + // The Interleaved GEMM global memory layout is similar to the congruous + // layout. + using ThreadMapA = typename MmaCore::SmemThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, layout::TensorNCxHWx, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + // Note GEMM shared memory threadmap is used here because conv global memory + // layout needs to be mapped to fprop which is similar to the crosswise + // layout which is used by the interleaved GEMM shared memory threadmap. + // The Interleaved GEMM global memory layout is similar to the congruous + // layout. + using ThreadMapB = typename MmaCore::SmemThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, layout::TensorCxRSKx, + ThreadMapB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultInterleavedConvEpilogue< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount, + InterleavedK + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Optimzed IteratorAlgorithm and +/// multistage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag + >; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + LayoutA, + ThreadMapA, + AccessTypeA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + LayoutB, + ThreadMapB, + AccessTypeB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * AlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + CacheOpB, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Optimzed IteratorAlgorithm and +// multistage pipeline with interleaved layout. +template < + typename ElementA, + typename ElementB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB, + int InterleavedK +> +struct DefaultConv2dFprop < + ElementA, + layout::TensorNCxHWx, + ElementB, + layout::TensorCxRSKx, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajorInterleaved, + ElementB, layout::RowMajorInterleaved, ElementAccumulator, LayoutC, arch::OpClassTensorOp, + Stages, MathOperatorTag, true + >; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::SmemThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + layout::TensorNCxHWx, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::SmemThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + layout::TensorCxRSKx, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Global, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultInterleavedConvEpilogue< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount, + InterleavedK + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Optimized IteratorAlgorithm +/// and 2 stage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + LayoutA, + ThreadMapA, + AccessTypeA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + LayoutB, + ThreadMapB, + AccessTypeB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Optimized IteratorAlgorithm and 2 stage +/// pipeline with interleaved layout. +template < + typename ElementA, + typename ElementB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB, + int InterleavedK +> +struct DefaultConv2dFprop < + ElementA, + layout::TensorNCxHWx, + ElementB, + layout::TensorCxRSKx, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajorInterleaved, + ElementB, layout::RowMajorInterleaved, + ElementAccumulator, LayoutC, arch::OpClassTensorOp, + 2, MathOperatorTag, true>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::SmemThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, layout::TensorNCxHWx, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::SmemThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, layout::TensorCxRSKx, + ThreadMapB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultInterleavedConvEpilogue< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount, + InterleavedK + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassSimt convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv2dFprop specialization for Analytic IteratorAlgorithm, +/// multi-stage pipeline, and FFMA-based mainloop for SM80 + +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Optimized IteratorAlgorithm, +/// multi-stage pipeline, and FFMA-based mainloop for SM80 + +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + LayoutA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + LayoutB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Analytic IteratorAlgorithm, +/// 2 stage pipeline, and FFMA-based mainloop for SM50 +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Optimized IteratorAlgorithm, +/// 2 stage pipeline, and FFMA-based mainloop for SM50 +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + LayoutA, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + LayoutB, + ThreadMapB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop_fusion.h b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop_fusion.h new file mode 100644 index 0000000..900b409 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop_fusion.h @@ -0,0 +1,357 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief + Default kernel-level fused activation's scale+bias+relu and implicit GEMM convolution + definitions that combine threadblock-scoped matrix multiply-add with the + appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/conv/kernel/default_conv2d.h" + +#include "mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/predicated_scale_bias_vector_access_iterator.h" +#include "mctlass/transform/threadblock/regular_scale_bias_vector_access_iterator.h" +#include "mctlass/gemm/warp/scale_bias_tile_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for fused batch norm and Conv2dFprop +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided +> struct DefaultConv2dFpropFusion; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassTensorOp convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Analytic IteratorAlgorithm and multistage +/// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv2dFpropFusion < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementScaleBias, + LayoutScaleBias, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using IteratorScaleBias = + mctlass::conv::threadblock::PredicatedScaleBiasVectorAccessIterator< + mctlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + using SmemIteratorScaleBias = + mctlass::transform::threadblock::RegularScaleBiasVectorAccessIterator< + mctlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static int const kThreadCount = 32; + + // Warp-level iterators to load scale and bias vectors + using WarpIteratorScaleBias = mctlass::gemm::warp::ScaleBiasTileIterator< + MatrixShape, ElementScaleBias, + LayoutScaleBias, MatrixShape, + typename WarpMmaTensorOp::IteratorA::Base::Policy, kThreadCount, + MmaCore::WarpCount::kK>; + + // Define the Mma + using Mma = threadblock::ImplicitGemmFpropFusionMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Global, + IteratorScaleBias, + SmemIteratorScaleBias, + arch::CacheOperation::Always, + MmaPolicy, + WarpIteratorScaleBias, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionFusion< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Optimzed IteratorAlgorithm and +/// multistage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv2dFpropFusion < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementScaleBias, + LayoutScaleBias, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag + >; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + LayoutA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + LayoutB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using IteratorScaleBias = + mctlass::conv::threadblock::PredicatedScaleBiasVectorAccessIterator< + mctlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + using SmemIteratorScaleBias = + mctlass::transform::threadblock::RegularScaleBiasVectorAccessIterator< + mctlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static int const kThreadCount = 32; + + // Warp-level iterators to load scale and bias vectors + using WarpIteratorScaleBias = mctlass::gemm::warp::ScaleBiasTileIterator< + MatrixShape, ElementScaleBias, + LayoutScaleBias, MatrixShape, + typename WarpMmaTensorOp::IteratorA::Base::Policy, kThreadCount, + MmaCore::WarpCount::kK>; + + // Define the Mma + using Mma = threadblock::ImplicitGemmFpropFusionMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Global, + IteratorScaleBias, + SmemIteratorScaleBias, + arch::CacheOperation::Always, + MmaPolicy, + WarpIteratorScaleBias, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionFusion< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop_with_broadcast.h b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop_with_broadcast.h new file mode 100644 index 0000000..02d5372 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop_with_broadcast.h @@ -0,0 +1,130 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Defines a GEMM with Reduction based on an existing UniversalGemm kernel. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/conv/kernel/default_conv2d_fprop.h" +#include "mctlass/conv/kernel/implicit_gemm_convolution_with_fused_epilogue.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_with_broadcast.h" +#include "mctlass/epilogue/threadblock/epilogue_with_broadcast.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided, + /// Access granularity of A matrix in units of elements + int AlignmentA = 128 / mctlass::sizeof_bits::value, + /// Access granularity of B matrix in units of elements + int AlignmentB = 128 / mctlass::sizeof_bits::value +> +struct DefaultConv2dFpropWithBroadcast { + + using ImplicitGemmBase = typename DefaultConv2dFprop< + ElementA, LayoutA, + ElementB, LayoutB, + ElementC, LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm, + StrideSupport, + AlignmentA, + AlignmentB + >::Kernel; + + // Define epilogue + using Epilogue = typename mctlass::conv::kernel::detail::DefaultConvEpilogueWithBroadcastTensorOp< + ArchTag, + typename ImplicitGemmBase::Epilogue::Shape, + typename ImplicitGemmBase::Epilogue::WarpMmaOperator, + ImplicitGemmBase::Epilogue::kPartitionsK, + ElementC, + typename EpilogueOutputOp::ElementT, + typename EpilogueOutputOp::ElementVector, + EpilogueOutputOp, + ImplicitGemmBase::Epilogue::kElementsPerAccess + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionWithFusedEpilogue< + typename ImplicitGemmBase::Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop_with_reduction.h b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop_with_reduction.h new file mode 100644 index 0000000..76cef11 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_fprop_with_reduction.h @@ -0,0 +1,130 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Defines a GEMM with Reduction based on an existing UniversalGemm kernel. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/conv/kernel/default_conv2d_fprop.h" +#include "mctlass/conv/kernel/implicit_gemm_convolution_with_fused_epilogue.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_with_reduction.h" +#include "mctlass/epilogue/threadblock/epilogue_with_reduction.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename EpilogueReductionOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided, + /// Access granularity of A matrix in units of elements + int AlignmentA = 128 / mctlass::sizeof_bits::value, + /// Access granularity of B matrix in units of elements + int AlignmentB = 128 / mctlass::sizeof_bits::value +> +struct DefaultConv2dFpropWithReduction { + + using ImplicitGemmBase = typename DefaultConv2dFprop< + ElementA, LayoutA, + ElementB, LayoutB, + ElementC, LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm, + StrideSupport, + AlignmentA, + AlignmentB + >::Kernel; + + // Define epilogue + using Epilogue = typename mctlass::conv::kernel::detail::DefaultConvEpilogueWithReductionTensorOp< + ArchTag, + typename ImplicitGemmBase::Epilogue::Shape, + typename ImplicitGemmBase::Epilogue::WarpMmaOperator, + ImplicitGemmBase::Epilogue::kPartitionsK, + ElementC, + EpilogueOutputOp, + EpilogueReductionOp, + ImplicitGemmBase::Epilogue::kElementsPerAccess + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionWithFusedEpilogue< + typename ImplicitGemmBase::Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_group_fprop.h b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_group_fprop.h new file mode 100644 index 0000000..554a331 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_group_fprop.h @@ -0,0 +1,622 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level implicit GEMM convolution definitions combine threadblock-scoped + matrix multiply-add with the appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/conv/kernel/default_conv2d.h" + +#include "mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_fixed_channels.h" +#include "mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_few_channels.h" + +#include "mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_fixed_channels.h" +#include "mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_few_channels.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv2dGroupFpro +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::GroupMode GroupMode, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided, + /// Access granularity of A matrix in units of elements + int AlignmentA = 128 / mctlass::sizeof_bits::value, + /// Access granularity of B matrix in units of elements + int AlignmentB = 128 / mctlass::sizeof_bits::value +> struct DefaultConv2dGroupFprop; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassTensorOp convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dGroupFprop specialization for Analytic IteratorAlgorithm and multistage +/// pipeline that supports all GroupMode. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::GroupMode GroupMode, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dGroupFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + GroupMode, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB +> { + + static_assert(std::is_same::value, + "Current group conv only support NHWC layout"); + static_assert(std::is_same::value, + "Current group conv only support NHWC layout"); + static_assert(std::is_same::value, + "Current group conv only support NHWC layout"); + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA, + AccessTypeA, + GroupMode + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB, + AccessTypeB, + GroupMode + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * AlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + CacheOpB, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv2dProblemSize, + GroupMode + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dGroupFprop specialization for Analytic IteratorAlgorithm and +/// 2 stage pipeline that supports all GroupMode. + +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::GroupMode GroupMode, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dGroupFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + GroupMode, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB +> { + + static_assert(std::is_same::value, + "Current group conv only support NHWC layout"); + static_assert(std::is_same::value, + "Current group conv only support NHWC layout"); + static_assert(std::is_same::value, + "Current group conv only support NHWC layout"); + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA, + AccessTypeA, + GroupMode + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB, + AccessTypeB, + GroupMode + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv2dProblemSize, + GroupMode + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dGroupFprop specialization for Optimized IteratorAlgorithm and multistage +/// pipeline that supports GroupMode::kSingleGroup. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dGroupFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + GroupMode::kSingleGroup, + IteratorAlgorithm::kOptimized, + StrideSupport, + AlignmentA, + AlignmentB +> { + + static_assert(std::is_same::value, + "Current group conv only support NHWC layout"); + static_assert(std::is_same::value, + "Current group conv only support NHWC layout"); + static_assert(std::is_same::value, + "Current group conv only support NHWC layout"); + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA, + AccessTypeA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB, + AccessTypeB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * AlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + CacheOpB, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv2dProblemSize, + GroupMode::kSingleGroup + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dGroupFprop specialization for Optimized IteratorAlgorithm and +/// 2 stage pipeline that supports GroupMode::kSingleGroup. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dGroupFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + GroupMode::kSingleGroup, + IteratorAlgorithm::kOptimized, + StrideSupport, + AlignmentA, + AlignmentB +> { + + static_assert(std::is_same::value, + "Current group conv only support NHWC layout"); + static_assert(std::is_same::value, + "Current group conv only support NHWC layout"); + static_assert(std::is_same::value, + "Current group conv only support NHWC layout"); + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + LayoutA, + ThreadMapA, + AccessTypeA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + LayoutB, + ThreadMapB, + AccessTypeB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv2dProblemSize, + GroupMode::kSingleGroup + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_wgrad.h b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_wgrad.h new file mode 100644 index 0000000..c2f1ef0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_wgrad.h @@ -0,0 +1,1011 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level implicit GEMM convolution definitions combine threadblock-scoped + matrix multiply-add with the appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/conv/kernel/default_conv2d.h" + +#include "mctlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv2d_tile_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided, + /// Access granularity of A matrix in units of elements + int AlignmentA = 128 / mctlass::sizeof_bits::value, + /// Access granularity of B matrix in units of elements + int AlignmentB = 128 / mctlass::sizeof_bits::value +> struct DefaultConv2dWgrad; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassTensorOp convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad specialization for Analytic IteratorAlgorithm and multistage +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::Conv2dWgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + AccessTypeA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::Conv2dWgradActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + AccessTypeB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad specialization for Analytic IteratorAlgorithm and two +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dWgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + AccessTypeA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dWgradActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + AccessTypeB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad specialization for Optimized IteratorAlgorithm and multistage +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::Conv2dWgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + AccessTypeA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::Conv2dWgradActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + AccessTypeB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + MmaPolicy, + Stages + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad specialization for Optimized IteratorAlgorithm and two +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultConv2dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::AlignedArray; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dWgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + AccessTypeA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dWgradActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB, + AccessTypeB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + kPartitionsK, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassSimt convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv2dWgrad specialization for Analytic IteratorAlgorithm, +/// multi-stage pipeline, and FFMA-based mainloop for SM80 + +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AccessTypeA, + int AccessTypeB +> +struct DefaultConv2dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AccessTypeA, + AccessTypeB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dWgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dWgradActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad specialization for Optimized IteratorAlgorithm, +/// multi-stage pipeline, and FFMA-based mainloop for SM80 + +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AccessTypeA, + int AccessTypeB +> +struct DefaultConv2dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport, + AccessTypeA, + AccessTypeB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dWgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dWgradActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad specialization for Analytic IteratorAlgorithm, +/// 2 stage pipeline, and FFMA-based mainloop for SM50 +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AccessTypeA, + int AccessTypeB +> +struct DefaultConv2dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport, + AccessTypeA, + AccessTypeB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dWgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dWgradActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad specialization for Optimized IteratorAlgorithm, +/// 2 stage pipeline, and FFMA-based mainloop for SM50 +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AccessTypeA, + int AccessTypeB +> +struct DefaultConv2dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport, + AccessTypeA, + AccessTypeB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dWgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dWgradActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad + >; + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_wgrad_fusion.h b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_wgrad_fusion.h new file mode 100644 index 0000000..02bb480 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_conv2d_wgrad_fusion.h @@ -0,0 +1,325 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief + Default kernel-level implicit GEMM convolution definitions combine threadblock-scoped + matrix multiply-add with the appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/conv/kernel/default_conv2d.h" + +#include "mctlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv2d_tile_iterator.h" +#include "mctlass/conv/threadblock/predicated_scale_bias_vector_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided +> struct DefaultConv2dWgradFusion; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassTensorOp convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad specialization for Analytic IteratorAlgorithm and multistage +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv2dWgradFusion < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementScaleBias, + LayoutScaleBias, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dWgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dWgradActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using IteratorScaleBias = + mctlass::conv::threadblock::PredicatedScaleBiasVectorIterator< + mctlass::MatrixShape<1, WarpShape::kN>, + ElementScaleBias, + LayoutScaleBias>; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmWgradFusionMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + IteratorScaleBias, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionFusion< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad specialization for Optimized IteratorAlgorithm and multistage +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv2dWgradFusion < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementScaleBias, + LayoutScaleBias, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv2dWgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv2dWgradActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using IteratorScaleBias = + mctlass::conv::threadblock::PredicatedScaleBiasVectorIterator< + mctlass::MatrixShape<1, WarpShape::kN>, + ElementScaleBias, + LayoutScaleBias>; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmWgradFusionMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + IteratorScaleBias, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionFusion< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_dgrad.h b/csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_dgrad.h new file mode 100644 index 0000000..bf8387a --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_dgrad.h @@ -0,0 +1,303 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level implicit GEMM convolution definitions combine threadblock-scoped + matrix multiply-add with the appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/conv/kernel/default_conv2d.h" + +#include "mctlass/conv/threadblock/conv3d_dgrad_output_gradient_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv3d_dgrad_filter_tile_access_iterator_optimized.h" + +#include "mctlass/conv/threadblock/conv3d_dgrad_output_gradient_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv3d_dgrad_filter_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_tile_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv3dDgrad +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided +> struct DefaultConv3dDgrad; + +/// Defines a kernel for Conv3dDgrad specialization for Analytic IteratorAlgorithm Dgrad Strided +// and multistage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv3dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic, + StrideSupport::kStrided +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv3dDgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + StrideSupport::kStrided + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv3dDgradFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Global, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad, + Conv3dProblemSize + >; +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dDgrad specialization for Optimized IteratorAlgorithm Dgrad Strided +// and multistage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv3dDgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport::kUnity +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv3dDgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA, + StrideSupport::kUnity + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + + using IteratorB = + mctlass::conv::threadblock::Conv3dDgradFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Global, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kDgrad, + Conv3dProblemSize + >; +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_fprop.h b/csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_fprop.h new file mode 100644 index 0000000..98a8fb9 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_fprop.h @@ -0,0 +1,515 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level implicit GEMM convolution definitions combine threadblock-scoped + matrix multiply-add with the appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/conv/kernel/default_conv2d.h" + +#include "mctlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_optimized.h" + + +#include "mctlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_analytic.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv2dFprop +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided +> struct DefaultConv3dFprop; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dFprop specialization for Analytic Iterator Algorithm +/// and 2 stage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag +> +struct DefaultConv3dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv3dFpropActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv3dFpropFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dFprop specialization for Analytic IteratorAlgorithm and multistage +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv3dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv3dFpropActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv3dFpropFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Global, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dFprop specialization for Optimized Iterator Algorithm +/// and 2 stage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag +> +struct DefaultConv3dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv3dFpropActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + LayoutA, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv3dFpropFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + LayoutB, + ThreadMapB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dFprop specialization for Optimized IteratorAlgorithm and multistage +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv3dFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv3dFpropActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + LayoutA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + + using IteratorB = + mctlass::conv::threadblock::Conv3dFpropFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + LayoutB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Global, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_fprop_fusion.h b/csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_fprop_fusion.h new file mode 100644 index 0000000..70c8c89 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_fprop_fusion.h @@ -0,0 +1,360 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level fused activation's scale+bias+relu and implicit GEMM convolution + definitions that combine threadblock-scoped matrix multiply-add with the + appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/conv/kernel/default_conv2d.h" + +#include "mctlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/predicated_scale_bias_vector_access_iterator.h" +#include "mctlass/transform/threadblock/regular_scale_bias_vector_access_iterator.h" +#include "mctlass/gemm/warp/scale_bias_tile_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for fused batch norm and Conv3dFprop +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided +> struct DefaultConv3dFpropFusion; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassTensorOp convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dFprop specialzation for Analytic IteratorAlgorithm and multistage +/// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv3dFpropFusion < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementScaleBias, + LayoutScaleBias, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv3dFpropActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv3dFpropFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using IteratorScaleBias = + mctlass::conv::threadblock::PredicatedScaleBiasVectorAccessIterator< + mctlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + using SmemIteratorScaleBias = + mctlass::transform::threadblock::RegularScaleBiasVectorAccessIterator< + mctlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static int const kThreadCount = 32; + + // Warp-level iterators to load scale and bias vectors + using WarpIteratorScaleBias = mctlass::gemm::warp::ScaleBiasTileIterator< + MatrixShape, ElementScaleBias, + LayoutScaleBias, MatrixShape, + typename WarpMmaTensorOp::IteratorA::Base::Policy, kThreadCount, + MmaCore::WarpCount::kK>; + + // Define the Mma + using Mma = threadblock::ImplicitGemmFpropFusionMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Global, + IteratorScaleBias, + SmemIteratorScaleBias, + arch::CacheOperation::Always, + MmaPolicy, + WarpIteratorScaleBias, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionFusion< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dFprop specialzation for Optimzed IteratorAlgorithm and +/// multistage pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementScaleBias, + typename LayoutScaleBias, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv3dFpropFusion < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementScaleBias, + LayoutScaleBias, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::RowMajor, + ElementB, layout::ColumnMajor, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, MathOperatorTag + >; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv3dFpropActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + LayoutA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv3dFpropFilterTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + LayoutB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + /// Define iterators over tiles from scale/bias vectors + using IteratorScaleBias = + mctlass::conv::threadblock::PredicatedScaleBiasVectorAccessIterator< + mctlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + using SmemIteratorScaleBias = + mctlass::transform::threadblock::RegularScaleBiasVectorAccessIterator< + mctlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + static int const kThreadCount = 32; + + // Warp-level iterators to load scale and bias vectors + using WarpIteratorScaleBias = mctlass::gemm::warp::ScaleBiasTileIterator< + MatrixShape, ElementScaleBias, + LayoutScaleBias, MatrixShape, + typename WarpMmaTensorOp::IteratorA::Base::Policy, kThreadCount, + MmaCore::WarpCount::kK>; + + // Define the Mma + using Mma = threadblock::ImplicitGemmFpropFusionMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Global, + IteratorScaleBias, + SmemIteratorScaleBias, + arch::CacheOperation::Always, + MmaPolicy, + WarpIteratorScaleBias, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolutionFusion< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_wgrad.h b/csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_wgrad.h new file mode 100644 index 0000000..c65ddea --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_conv3d_wgrad.h @@ -0,0 +1,509 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level implicit GEMM convolution definitions combine threadblock-scoped + matrix multiply-add with the appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/conv/kernel/default_conv2d.h" + +#include "mctlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_optimized.h" +#include "mctlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_optimized.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv2dWgrad +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kOptimized, + conv::StrideSupport StrideSupport = StrideSupport::kStrided +> struct DefaultConv3dWgrad; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dWgrad specialization for Analytic IteratorAlgorithm and multistage +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv3dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kAnalytic +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv3dWgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv3dWgradActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv3dWgrad specialization for Analytic IteratorAlgorithm and two +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag +> +struct DefaultConv3dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kAnalytic +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv3dWgradOutputGradientTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv3dWgradActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, + ThreadMapB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a kernel for Conv3dWgrad specialization for Optimized IteratorAlgorithm and multistage +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag +> +struct DefaultConv3dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::Conv3dWgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::Conv3dWgradActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmMultistage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + arch::CacheOperation::Always, + IteratorB, + SmemIteratorB, + arch::CacheOperation::Always, + MmaPolicy, + Stages + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad, + Conv3dProblemSize + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv3dWgrad specialization for Optimized IteratorAlgorithm and two +// pipeline. +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag +> +struct DefaultConv3dWgrad < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, + IteratorAlgorithm::kOptimized +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, layout::ColumnMajor, + ElementB, layout::RowMajor, ElementAccumulator, layout::RowMajor, OperatorClass, + 2, MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv3dWgradOutputGradientTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementA, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv3dWgradActivationTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, + ThreadMapB + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaTensorOp = typename MmaCore::MmaTensorOp; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::ImplicitGemmPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + // Define the epilogue + using Epilogue = typename detail::DefaultConvEpilogue< + ArchTag, + ThreadblockShape, + WarpMmaTensorOp, + 1, + EpilogueOutputOp + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kWgrad, + Conv3dProblemSize + >; +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/default_depthwise_fprop.h b/csrc/mctlass/include/mctlass/conv/kernel/default_depthwise_fprop.h new file mode 100644 index 0000000..13dc4de --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/default_depthwise_fprop.h @@ -0,0 +1,588 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level Depthwise implicit GEMM convolution definitions combine threadblock-scoped + matrix multiply-add with the appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/conv/kernel/default_conv2d.h" +#include "mctlass/conv/kernel/direct_convolution.h" + +#include "mctlass/conv/threadblock/depthwise_mma_core_with_lane_access_size.h" + +#include "mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h" +#include "mctlass/conv/threadblock/depthwise_fprop_pipelined.h" + +// Direct Conv Related Header files +#include "mctlass/conv/threadblock/depthwise_fprop_activation_tile_access_iterator_direct_conv_optimized.h" +#include "mctlass/conv/threadblock/depthwise_fprop_activation_tile_access_iterator_direct_conv_fixed_stride_dilation.h" + +#include "mctlass/conv/threadblock/depthwise_fprop_filter_tile_access_iterator_direct_conv_optimized.h" +#include "mctlass/conv/threadblock/depthwise_fprop_direct_conv_multistage.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Conv2dFprop +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kAnalytic, + conv::StrideSupport StrideSupport = StrideSupport::kStrided, + /// Access granularity of A matrix in units of elements + int AlignmentA = 128 / mctlass::sizeof_bits::value, + /// Access granularity of B matrix in units of elements + int AlignmentB = mctlass::sizeof_bits::value / mctlass::sizeof_bits::value +> struct DefaultDepthwiseFprop; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for DepthwiseFprop with direct convolution algorithm +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename OperatorClass, + typename ArchTag, + typename ThreadblockShape, + typename ThreadBlockOutputShape, + typename FilterShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kAnalytic, + conv::StrideSupport StrideSupport = StrideSupport::kStrided, + // MatrixShape + typename StrideShape = mctlass::MatrixShape<-1, -1>, + // MatrixShape< Height, Width> + typename DilationShape = mctlass::MatrixShape<-1, -1>, + /// Access granularity of A matrix in units of elements + int AlignmentA = 128 / mctlass::sizeof_bits::value, + /// Access granularity of B matrix in units of elements + int AlignmentB = 128 / mctlass::sizeof_bits::value +> struct DefaultDepthwiseDirect2dConvFprop; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// OpClassSimt convolutions +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Depthwise specialization for Analytic IteratorAlgorithm +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + int AlignmentA, + int AlignmentB +> +struct DefaultDepthwiseFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + MathOperatorTag, // mctlass::arch::OpMultiplyAdd + IteratorAlgorithm::kAnalytic, + StrideSupport, + AlignmentA, + AlignmentB +> { + + // Define the core components from GEMM + using MmaCore = typename mctlass::conv::threadblock::DepthwiseMmaCoreWithLaneAccessSize< + ThreadblockShape, + WarpShape, + InstructionShape, + ElementA, + layout::RowMajor, + ElementB, + layout::ColumnMajor, + ElementAccumulator, + layout::RowMajor, + arch::OpClassSimt, + 128, + sizeof_bits::value, + 2, + MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropActivationTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementA, LayoutA, + ThreadMapA + > + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::TileIterator< + mctlass::conv::threadblock::Conv2dFpropFilterTileAccessIteratorAnalytic< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB, + AccessTypeB, + mctlass::conv::GroupMode::kDepthwise + > + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + + // Define the Mma + using Mma = threadblock::DepthwiseFpropPipelined< + ThreadblockShape, + IteratorA, + SmemIteratorA, + IteratorB, + SmemIteratorB, + ElementC, + LayoutC, + MmaPolicy + >; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + // Define the kernel + using Kernel = mctlass::conv::kernel::ImplicitGemmConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv2dProblemSize, + mctlass::conv::GroupMode::kDepthwise + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Depthwise specialization for direct 2d conv implementation, +/// multiple stage pipeline, and SIMT-based mainloop +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename ThreadBlockOutputShape, + typename FilterShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + typename StrideShape, + typename DilationShape, + int AlignmentA, + int AlignmentB +> +struct DefaultDepthwiseDirect2dConvFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + ThreadBlockOutputShape, + FilterShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kOptimized, + StrideSupport, + StrideShape, + DilationShape, + AlignmentA, + AlignmentB +> { + // One warp handles the entrie groups per cta. + static_assert(ThreadblockShape::kN == WarpShape::kN, + "ThreadblockShape::kN should be same as WarpShape::kN "); + static_assert(ThreadblockShape::kK == FilterShape::kCount && WarpShape::kK == FilterShape::kCount, + "ThreadblockShape::kK and WarpShape::kK should be same as filter size"); + static_assert(ThreadblockShape::kM % WarpShape::kM == 0, + "ThreadblockShape::kM must be divisible by WarpShape shape::kM"); + static_assert(ThreadBlockOutputShape::kN, "ThreadBlockOutputShape::kN should be 1"); + + // Define the core components from GEMM + using MmaCore = typename mctlass::conv::threadblock::DepthwiseDirectConvMmaCoreWithLaneAccessSize< + ThreadblockShape, + ThreadBlockOutputShape, + FilterShape, + WarpShape, + InstructionShape, + ElementA, + layout::RowMajor, + ElementB, + layout::ColumnMajor, + ElementAccumulator, + layout::RowMajor, + arch::OpClassSimt, + 128, + 128, + Stages, + MathOperatorTag>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::DepthwiseFpropActivationDirect2dConvTileAccessIteratorOptimized< + mctlass::MatrixShape, // < outputShape:KMNK, groups per cta> + ThreadBlockOutputShape, + ElementA, LayoutA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::DepthwiseFpropFilterDirectConvTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + using ThreadOutputShape = typename MmaCore::ThreadOutputShape; + static mctlass::arch::CacheOperation::Kind const CacheOpA = + ((sizeof_bits::value * AlignmentA) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * AlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultDirectConvEpilogueSimt< + ThreadblockShape, // < outputShape:KMNK, groups per cta> + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount, + ThreadOutputShape, + ThreadBlockOutputShape + >::Epilogue; + + // Define the Mma + using Mma = threadblock::DepthwiseFpropDirectConvMultipleStage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + CacheOpA, + IteratorB, + SmemIteratorB, + CacheOpB, + MmaPolicy, + Stages, + Epilogue + >; + + // Define the kernel + using Kernel = mctlass::conv::kernel::DirectConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv2dProblemSize, + mctlass::conv::GroupMode::kDepthwise, + ThreadBlockOutputShape + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Defines a kernel for Depthwise specialization for direct 2d conv implementation, +/// multiple stage pipeline, and SIMT-based mainloop +template < + typename ElementA, + typename LayoutA, + typename ElementB, + typename LayoutB, + typename ElementC, + typename LayoutC, + typename ElementAccumulator, + typename ArchTag, + typename ThreadblockShape, + typename ThreadBlockOutputShape, + typename FilterShape, + typename WarpShape, + typename InstructionShape, + typename EpilogueOutputOp, + typename ThreadblockSwizzle, + int Stages, + typename MathOperatorTag, + conv::StrideSupport StrideSupport, + typename StrideShape, + typename DilationShape, + int AlignmentA, + int AlignmentB +> +struct DefaultDepthwiseDirect2dConvFprop < + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + ThreadBlockOutputShape, + FilterShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MathOperatorTag, + IteratorAlgorithm::kFixedStrideDilation, + StrideSupport, + StrideShape, + DilationShape, + AlignmentA, + AlignmentB +> { + + + + // One warp handles the entrie groups per cta. + static_assert(ThreadblockShape::kN == WarpShape::kN, + "ThreadblockShape::kN should be same as WarpShape::kN "); + static_assert(ThreadblockShape::kK == FilterShape::kCount && WarpShape::kK == FilterShape::kCount, + "ThreadblockShape::kK and WarpShape::kK should be same as filter size"); + static_assert(ThreadblockShape::kM % WarpShape::kM == 0, + "ThreadblockShape::kM must be divisible by WarpShape shape::kM"); + static_assert(ThreadBlockOutputShape::kN, "ThreadBlockOutputShape::kN should be 1"); + + static_assert(StrideShape::kRow >= 0 && StrideShape::kColumn >= 0, "Stride should be fixed"); + static_assert(DilationShape::kRow >= 0 && DilationShape::kColumn >= 0, "Stride should be fixed"); + + // Activations loaded by threadblock + static int const ActivationShapeH = (ThreadBlockOutputShape::kH - 1) * StrideShape::kRow + + (FilterShape::kRow - 1) * DilationShape::kRow + 1; + + static int const ActivationShapeW = (ThreadBlockOutputShape::kW - 1) * StrideShape::kColumn + + (FilterShape::kColumn - 1) * DilationShape::kColumn + 1; + + using ActivationShape = + mctlass::conv::TensorNHWCShape<1, ActivationShapeH, ActivationShapeW, ThreadblockShape::kN >; + + // Define the core components from GEMM + using MmaCore = typename mctlass::conv::threadblock::DepthwiseDirectConvMmaCoreWithLaneAccessSize< + ThreadblockShape, + ThreadBlockOutputShape, + FilterShape, + WarpShape, + InstructionShape, + ElementA, + layout::RowMajor, + ElementB, + layout::ColumnMajor, + ElementAccumulator, + layout::RowMajor, + arch::OpClassSimt, + 128, + 128, + Stages, + MathOperatorTag, + IteratorAlgorithm::kFixedStrideDilation, + StrideShape, + DilationShape, + ActivationShape>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using IteratorA = + mctlass::conv::threadblock::DepthwiseFpropActivationDirect2dConvTileAccessIteratorFixedStrideDilation< + mctlass::MatrixShape, // < outputShape:KMNK, groups per cta> + ThreadBlockOutputShape, + StrideShape, + DilationShape, + ActivationShape, + ElementA, LayoutA, + ThreadMapA + >; + + using SmemIteratorA = typename MmaCore::SmemIteratorA; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::AlignedArray; + using IteratorB = + mctlass::conv::threadblock::DepthwiseFpropFilterDirectConvTileAccessIteratorOptimized< + mctlass::MatrixShape, + ElementB, LayoutB, + ThreadMapB + >; + + using SmemIteratorB = typename MmaCore::SmemIteratorB; + + // Warp-level GEMM components + using WarpMmaSimtOp = typename MmaCore::MmaWarpSimt; + using MmaPolicy = typename MmaCore::MmaPolicy; + using ThreadOutputShape = typename MmaCore::ThreadOutputShape; + static mctlass::arch::CacheOperation::Kind const CacheOpA = + ((sizeof_bits::value * AlignmentA) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * AlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the epilogue + using Epilogue = typename epilogue::threadblock::DefaultDirectConvEpilogueSimt< + ThreadblockShape, // < outputShape:KMNK, groups per cta> + WarpMmaSimtOp, + EpilogueOutputOp, + EpilogueOutputOp::kCount, + ThreadOutputShape, + ThreadBlockOutputShape + >::Epilogue; + + // Define the Mma + using Mma = threadblock::DepthwiseFpropDirectConvMultipleStage< + ThreadblockShape, + IteratorA, + SmemIteratorA, + CacheOpA, + IteratorB, + SmemIteratorB, + CacheOpB, + MmaPolicy, + Stages, + Epilogue, + IteratorAlgorithm::kFixedStrideDilation + >; + + // Define the kernel + using Kernel = mctlass::conv::kernel::DirectConvolution< + Mma, + Epilogue, + ThreadblockSwizzle, + conv::Operator::kFprop, + Conv2dProblemSize, + mctlass::conv::GroupMode::kDepthwise, + ThreadBlockOutputShape + >; +}; + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/direct_convolution.h b/csrc/mctlass/include/mctlass/conv/kernel/direct_convolution.h new file mode 100644 index 0000000..ce8a699 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/direct_convolution.h @@ -0,0 +1,505 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a multi-staged Depthwise Convolution kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/aligned_buffer.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/semaphore.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/epilogue/threadblock/output_iterator_parameter.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Parameters structure +template > ///! OutputShape per ThreadBlock +struct DirectConvolutionParams { + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using ThreadBlockOutputShape = ThreadBlockOutputShape_; + static Operator const kConvolutionalOperator = ConvOperator; + using ConvProblemSize = ConvProblemSize_; + using Arguments = Arguments_; + using ConvOutputIteratorParameter = ConvOutputIteratorParameter_; + + using ThreadblockShape = typename Mma::Shape; + static IteratorAlgorithm const kIteratorAlgorithm = Mma::IteratorA::kIteratorAlgorithm; + static conv::GroupMode const kGroupMode = GroupMode_; + static int const kStages = Mma::kStages; + + ConvProblemSize problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + gemm::GemmCoord implicit_gemm_problem_size; + int swizzle_log_tile; + int smem_size_; + + int gemm_k_iterations; + int gemm_k_iterations_per_channel; + typename Mma::IteratorA::Params iterator_A; + typename Mma::IteratorA::Element const *ptr_A; + typename Mma::IteratorB::Params iterator_B; + typename Mma::IteratorB::Element const *ptr_B; + typename Mma::IteratorB::Element *ptr_reordered_B; + typename Epilogue::OutputTileIterator::Params iterator_C; + typename Epilogue::OutputTileIterator::Element *ptr_C; + typename Epilogue::OutputTileIterator::Params iterator_D; + typename Epilogue::OutputTileIterator::Element *ptr_D; + typename EpilogueOutputOp::Params output_op; + int *semaphore; + SplitKMode split_k_mode; + int split_k_slices; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + DirectConvolutionParams() : swizzle_log_tile(0), gemm_k_iterations(0) {} + + /// + MCTLASS_HOST_DEVICE + DirectConvolutionParams(Arguments const &args, int *semaphore = nullptr) + : problem_size(args.problem_size), + implicit_gemm_problem_size( + mctlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, args.problem_size)), + iterator_A(Mma::IteratorA::getParams(args.problem_size, args.ref_A.layout())), + ptr_A(args.ref_A.data()), + iterator_B(Mma::IteratorB::getParams(args.problem_size, args.ref_B.layout())), + ptr_B(args.ref_B.data()), + ptr_reordered_B(args.ref_reordered_B.data()), + iterator_C(ConvOutputIteratorParameter::layout(args.ref_C), args.problem_size), + ptr_C(args.ref_C.data()), + iterator_D(ConvOutputIteratorParameter::layout(args.ref_D), args.problem_size), + ptr_D(args.ref_D.data()), + output_op(args.output_op), + semaphore(semaphore), + split_k_mode(args.split_k_mode), + split_k_slices(args.problem_size.split_k_slices) { + gemm_k_iterations = + depthwise_gemm_k_iterations(kConvolutionalOperator, + ThreadblockShape::kK, + args.problem_size, + kIteratorAlgorithm, + kGroupMode, + ThreadblockShape::kN); + + gemm_k_iterations_per_channel = implicit_gemm_k_iterations_per_channel( + kConvolutionalOperator, ThreadblockShape::kK, args.problem_size, kIteratorAlgorithm); + + ThreadblockSwizzle threadblock_swizzle; + + grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + kConvolutionalOperator, + problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.problem_size.split_k_slices); + + swizzle_log_tile = threadblock_swizzle.get_log_tile(grid_tiled_shape); + + // Dynamic SMEM usage because stride and dilation are runtime params. + smem_size_ = (iterator_A.activation_size * kStages + iterator_B.filter_size); + } + + MCTLASS_HOST_DEVICE + int get_smem_size() { + // Dynamic Smem Size + return smem_size_; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +template +struct ReorderKernel { + using Params = Params_; + using ElementB = ElementB_; + + union SharedStorage {}; + + static unsigned int const kReorderKernelThreadPerCTA = 128; + + MCTLASS_HOST_DEVICE + ReorderKernel() {} + + MCTLASS_HOST_DEVICE + static dim3 get_grid_shape(Params const ¶ms) { + return dim3{static_cast( + (params.problem_size.filter_size() + kReorderKernelThreadPerCTA - 1) / + kReorderKernelThreadPerCTA), + 1, + 1}; + } + + MCTLASS_HOST_DEVICE + static dim3 get_block_shape() { return dim3{kReorderKernelThreadPerCTA, 1, 1}; } + + MCTLASS_HOST_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + int64_t m = static_cast(params.problem_size.groups); + int64_t n = static_cast(params.problem_size.filter_size() / params.problem_size.K); + const ElementB *src_with_type = static_cast(params.ptr_B); + ElementB *dst_with_type = static_cast(params.ptr_reordered_B); + + int64_t linear_index = blockIdx.x * kReorderKernelThreadPerCTA + threadIdx.x; + int64_t index_m = linear_index / n; + int64_t index_n = linear_index % n; + int64_t new_linear_index = index_m + index_n * m; + + if (linear_index < m * n) { + dst_with_type[new_linear_index] = src_with_type[linear_index]; + } + return; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + conv::Operator ConvOperator, ///! Convolutional operator (Fprop, Dgrad, Wgrad) + typename ConvProblemSize_ = Conv2dProblemSize, ///! Convolutional operator on 2D or 3D problem + conv::GroupMode GroupMode_ = conv::GroupMode::kNone, ///! Group mode + typename ThreadBlockOutputShape_ = mctlass::conv::TensorNHWCShape<1, 1, 1, 1> +> +struct DirectConvolution { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using ThreadBlockOutputShape = ThreadBlockOutputShape_; + static Operator const kConvolutionalOperator = ConvOperator; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename EpilogueOutputOp::ElementOutput; + + /// Set output tensor C layout + using LayoutC = LayoutA; + + using ElementAccumulator = typename EpilogueOutputOp::ElementAccumulator; + using ElementCompute = typename EpilogueOutputOp::ElementCompute; + + using WarpMmaOperator = typename Mma::Policy::Operator; + + using ArchMmaOperator = typename WarpMmaOperator::ArchMmaOperator; + using MathOperator = typename ArchMmaOperator::Operator; + + using OperatorClass = typename WarpMmaOperator::OperatorClass; + using ArchTag = typename WarpMmaOperator::ArchTag; + + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename WarpMmaOperator::Shape; + using InstructionShape = typename mctlass::gemm::GemmShape<1, 1, 1>; + + static int const kStages = Mma::kStages; + static IteratorAlgorithm const kIteratorAlgorithm = Mma::IteratorA::kIteratorAlgorithm; + static StrideSupport const kStrideSupport = Mma::IteratorA::kStrideSupport; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + using TensorRefA = typename Mma::IteratorA::TensorRef; + using TensorRefB = typename Mma::IteratorB::TensorRef; + using TensorRefC = mctlass::TensorRef; + + /// Check iterator A and B convolution dimension are the same and + // set device::ImplicitGemmConvolution::kConvDim + static_assert(Mma::IteratorA::kConvDim == Mma::IteratorB::kConvDim, + "Convolution on different different dimensions is not supported"); + static int const kConvDim = Mma::IteratorA::kConvDim; + + /// Conv dimension and problem size structure (Conv2d or Conv3d) + using ConvProblemSize = ConvProblemSize_; + + static conv::GroupMode const kGroupMode = GroupMode_; + + + // + // + // + using ConvOutputIteratorParameter = epilogue::threadblock::ConvOutputIteratorParameter< + LayoutC, + typename Epilogue::OutputTileIterator::Layout, + TensorRefC, + ConvOperator, + ConvProblemSize + >; + + + /// Argument structure + struct Arguments { + + // + // Data members + // + + ConvProblemSize problem_size; + TensorRefA ref_A; + TensorRefB ref_B; + TensorRefB ref_reordered_B; + TensorRefC ref_C; + TensorRefC ref_D; + typename EpilogueOutputOp::Params output_op; + SplitKMode split_k_mode; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + MCTLASS_HOST_DEVICE + Arguments( + ConvProblemSize const & problem_size + ): + problem_size(problem_size) { } + + MCTLASS_HOST_DEVICE + Arguments( + ConvProblemSize const & problem_size, + TensorRefA const & ref_A, + TensorRefB const & ref_B, + TensorRefC const & ref_C, + TensorRefC const & ref_D, + typename EpilogueOutputOp::Params const & output_op, + TensorRefB const & ref_reordered_B = nullptr, + SplitKMode const & split_k_mode = SplitKMode::kSerial + ): + problem_size(problem_size), + ref_A(ref_A), + ref_B(ref_B), + ref_C(ref_C), + ref_D(ref_D), + output_op(output_op), + ref_reordered_B(ref_reordered_B), + split_k_mode(split_k_mode) + { + + } + + }; + + using Params = + typename mctlass::conv::kernel::DirectConvolutionParams; + + using ReorderKernel = typename mctlass::conv::kernel::ReorderKernel; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + DirectConvolution() { } + + /// Executes one ImplicitGEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_idx = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if threadblock is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_idx.m() || + params.grid_tiled_shape.n() <= threadblock_tile_idx.n()) { + + return; + } + + // Compute position within threadblock + int thread_idx = threadIdx.x; + int iterator_column_offset = 0; + int filter_row_offset = 0; + if (kGroupMode != GroupMode::kNone) { + if (kGroupMode == GroupMode::kDepthwise) { + iterator_column_offset += threadblock_tile_idx.n() * Mma::Shape::kN; + } + } + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.iterator_A, + params.problem_size, + params.ptr_A, + thread_idx, + MatrixCoord( + threadblock_tile_idx.m() + threadblock_tile_idx.k(), + iterator_column_offset + ) + ); + + typename Mma::IteratorB iterator_B( + params.iterator_B, + params.problem_size, + params.ptr_reordered_B, + thread_idx, + MatrixCoord( + filter_row_offset, + iterator_column_offset + ) + ); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 64, 0); + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // Compute logical position within grid + threadblock_tile_idx = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + + MatrixCoord threadblock_offset( + threadblock_tile_idx.m() + threadblock_tile_idx.k(), + threadblock_tile_idx.n() * Mma::Shape::kN + ); + + // Tile iterator writing to destination tensor + typename Epilogue::OutputTileIterator iterator_D( + params.iterator_D, + params.ptr_D, + ConvOutputIteratorParameter::extent(params.problem_size), + thread_idx, + threadblock_offset + ); + + // Tile iterator reading from source accumulator tensor + typename Epilogue::OutputTileIterator iterator_C( + params.iterator_C, + params.ptr_C, + ConvOutputIteratorParameter::extent(params.problem_size), + thread_idx, + threadblock_offset + ); + + + // Construct the epilogue + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + + // Compute threadblock-scoped matrix multiply-add + // Epilogue is fused in the mainloop + mma(params.gemm_k_iterations, + accumulators, + iterator_A, + params.iterator_A, + iterator_B, + params.iterator_B, + accumulators, + epilogue, + output_op, + iterator_D, + iterator_C, + params.split_k_slices); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution.h b/csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution.h new file mode 100644 index 0000000..5989e56 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution.h @@ -0,0 +1,456 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined Implicit GEMM kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/aligned_buffer.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/semaphore.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/epilogue/threadblock/output_iterator_parameter.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + conv::Operator ConvOperator, ///! Convolutional operator (Fprop, Dgrad, Wgrad) + typename ConvProblemSize_ = Conv2dProblemSize, ///! Convolutional operator on 2D or 3D problem + conv::GroupMode GroupMode_ = conv::GroupMode::kNone ///! Group mode +> +struct ImplicitGemmConvolution { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static Operator const kConvolutionalOperator = ConvOperator; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename EpilogueOutputOp::ElementOutput; + + /// Set output tensor C layout + using LayoutC = LayoutA; + + using ElementAccumulator = typename EpilogueOutputOp::ElementAccumulator; + using ElementCompute = typename EpilogueOutputOp::ElementCompute; + + using WarpMmaOperator = typename Mma::Policy::Operator; + + using ArchMmaOperator = typename WarpMmaOperator::ArchMmaOperator; + using MathOperator = typename ArchMmaOperator::Operator; + + using OperatorClass = typename WarpMmaOperator::OperatorClass; + using ArchTag = typename WarpMmaOperator::ArchTag; + + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename WarpMmaOperator::Shape; + using InstructionShape = typename ArchMmaOperator::Shape; + + static int const kStages = Mma::kStages; + static IteratorAlgorithm const kIteratorAlgorithm = Mma::IteratorA::kIteratorAlgorithm; + static StrideSupport const kStrideSupport = Mma::IteratorA::kStrideSupport; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 64 * WarpCount::kCount; + + using TensorRefA = typename Mma::IteratorA::TensorRef; + using TensorRefB = typename Mma::IteratorB::TensorRef; + using TensorRefC = mctlass::TensorRef; + + /// Check iterator A and B convolution dimension are the same and + // set device::ImplicitGemmConvolution::kConvDim + static_assert(Mma::IteratorA::kConvDim == Mma::IteratorB::kConvDim, + "Convolution on different different dimensions is not supported"); + static int const kConvDim = Mma::IteratorA::kConvDim; + + /// Conv dimension and problem size structure (Conv2d or Conv3d) + using ConvProblemSize = ConvProblemSize_; + + static conv::GroupMode const kGroupMode = GroupMode_; + + /// Wgrad C stride idx for implicit gemm algorithm + // Conv2d row-major matrix C (KxRSC) + // Conv3d row-major matrix C (KxTRSC) + static int const kWgradCStrideIdx = + platform::is_same::value ? 2 : 3; + + /// This chooses the appropriate stride element of the C tensor. + static int const kTensorCStrideIdx = + (kConvolutionalOperator == conv::Operator::kWgrad ? kWgradCStrideIdx : 0); + + // + // + // + using ConvOutputIteratorParameter = epilogue::threadblock::ConvOutputIteratorParameter< + LayoutC, + typename Epilogue::OutputTileIterator::Layout, + TensorRefC, + ConvOperator, + ConvProblemSize + >; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + ConvProblemSize problem_size; + TensorRefA ref_A; + TensorRefB ref_B; + TensorRefC ref_C; + TensorRefC ref_D; + typename EpilogueOutputOp::Params output_op; + SplitKMode split_k_mode; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + MCTLASS_HOST_DEVICE + Arguments( + ConvProblemSize const & problem_size + ): + problem_size(problem_size) { } + + MCTLASS_HOST_DEVICE + Arguments( + ConvProblemSize const & problem_size, + TensorRefA const & ref_A, + TensorRefB const & ref_B, + TensorRefC const & ref_C, + TensorRefC const & ref_D, + typename EpilogueOutputOp::Params const & output_op, + SplitKMode const & split_k_mode = SplitKMode::kSerial + ): + problem_size(problem_size), + ref_A(ref_A), + ref_B(ref_B), + ref_C(ref_C), + ref_D(ref_D), + output_op(output_op), + split_k_mode(split_k_mode) + { + + } + + }; + + /// Parameters structure + struct Params { + ConvProblemSize problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + gemm::GemmCoord implicit_gemm_problem_size; + int swizzle_log_tile; + + int gemm_k_iterations; + int gemm_k_iterations_per_channel; + typename Mma::IteratorA::Params iterator_A; + typename Mma::IteratorA::Element const *ptr_A; + typename Mma::IteratorB::Params iterator_B; + typename Mma::IteratorB::Element const *ptr_B; + typename Epilogue::OutputTileIterator::Params iterator_C; + typename Epilogue::OutputTileIterator::Element *ptr_C; + typename Epilogue::OutputTileIterator::Params iterator_D; + typename Epilogue::OutputTileIterator::Element *ptr_D; + typename EpilogueOutputOp::Params output_op; + int *semaphore; + SplitKMode split_k_mode; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): swizzle_log_tile(0), gemm_k_iterations(0) { } + + /// + MCTLASS_HOST_DEVICE + Params( + Arguments const &args, + int *semaphore = nullptr + ): + problem_size(args.problem_size), + implicit_gemm_problem_size(mctlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, args.problem_size)), + iterator_A(Mma::IteratorA::getParams(args.problem_size, args.ref_A.layout())), + ptr_A(args.ref_A.data()), + iterator_B(args.problem_size, args.ref_B.layout()), + ptr_B(args.ref_B.data()), + iterator_C(ConvOutputIteratorParameter::layout(args.ref_C)), + ptr_C(args.ref_C.data()), + iterator_D(ConvOutputIteratorParameter::layout(args.ref_D)), + ptr_D(args.ref_D.data()), + output_op(args.output_op), + semaphore(semaphore), + split_k_mode(args.split_k_mode) + { + gemm_k_iterations = implicit_gemm_k_iterations( + kConvolutionalOperator, + ThreadblockShape::kK, + args.problem_size, + kIteratorAlgorithm, + kGroupMode, + ThreadblockShape::kN); + + gemm_k_iterations_per_channel = implicit_gemm_k_iterations_per_channel( + kConvolutionalOperator, ThreadblockShape::kK, args.problem_size, kIteratorAlgorithm); + + ThreadblockSwizzle threadblock_swizzle; + + grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + implicit_gemm_problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.problem_size.split_k_slices); + + swizzle_log_tile = threadblock_swizzle.get_log_tile(grid_tiled_shape); + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + ImplicitGemmConvolution() { } + + /// Executes one ImplicitGEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_idx = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_idx.m() || + params.grid_tiled_shape.n() <= threadblock_tile_idx.n()) { + + return; + } + + // Compute position within threadblock + int thread_idx = threadIdx.x; + int iterator_A_column_offset = threadblock_tile_idx.k() * Mma::Shape::kK; + if (kGroupMode != GroupMode::kNone) { + if (kGroupMode != GroupMode::kDepthwise) { + int k_per_group = params.problem_size.K / params.problem_size.groups; + int group_idx = threadblock_tile_idx.n() * Mma::Shape::kN / k_per_group; + int channels_per_group = params.problem_size.C / params.problem_size.groups; + iterator_A_column_offset += group_idx * channels_per_group; + } else { + iterator_A_column_offset += threadblock_tile_idx.n() * Mma::Shape::kN; + } + } + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.iterator_A, + params.problem_size, + params.ptr_A, + thread_idx, + MatrixCoord( + threadblock_tile_idx.m() * Mma::Shape::kM, + iterator_A_column_offset + ) + ); + + typename Mma::IteratorB iterator_B( + params.iterator_B, + params.problem_size, + params.ptr_B, + thread_idx, + MatrixCoord( + threadblock_tile_idx.k() * Mma::Shape::kK, + threadblock_tile_idx.n() * Mma::Shape::kN + ) + ); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + mma(params.gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators, params.gemm_k_iterations_per_channel); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // Construct the semaphore. + int block_idx = threadblock_tile_idx.m() + threadblock_tile_idx.n() * params.grid_tiled_shape.m(); + + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + // Compute logical position within grid + threadblock_tile_idx = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_idx.k(), params.grid_tiled_shape.k()); + } + + MatrixCoord threadblock_offset( + threadblock_tile_idx.m() * Mma::Shape::kM, + threadblock_tile_idx.n() * Mma::Shape::kN + ); + + // Tile iterator writing to destination tensor + typename Epilogue::OutputTileIterator iterator_D( + params.iterator_D, + params.ptr_D, + ConvOutputIteratorParameter::extent(params.problem_size), + thread_idx, + threadblock_offset + ); + + // Tile iterator reading from source accumulator tensor + typename Epilogue::OutputTileIterator iterator_C( + params.iterator_C, + params.ptr_C, + ConvOutputIteratorParameter::extent(params.problem_size), + thread_idx, + threadblock_offset + ); + + + // Construct the epilogue + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_idx.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_idx.k()); + + } + // Each split-k-slice writes to a unique tensor location + else if (params.split_k_mode == SplitKMode::kParallel) { + iterator_D.add_pointer_offset(threadblock_tile_idx.k() * + mctlass::conv::implicit_gemm_tensor_c_size(ConvOperator, params.problem_size)); + } + + // Run efficient epilogue + epilogue(output_op, iterator_D, accumulators, iterator_C); + + // + // Release the semaphore + // + + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_idx.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_idx.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution_fusion.h b/csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution_fusion.h new file mode 100644 index 0000000..b00a70e --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution_fusion.h @@ -0,0 +1,463 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined fused activation's scale+bias+relu and Implicit GEMM kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/aligned_buffer.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/semaphore.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/epilogue/threadblock/output_iterator_parameter.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + conv::Operator ConvOperator, ///! Convolutional operator (Fprop, Dgrad, Wgrad) + typename ConvProblemSize_ = Conv2dProblemSize ///! Convolutional operator on 2D or 3D problem +> +struct ImplicitGemmConvolutionFusion { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static Operator const kConvolutionalOperator = ConvOperator; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + + using ElementScaleBias = typename Mma::IteratorScaleBias::Element; + using LayoutScaleBias = typename Mma::IteratorScaleBias::Layout; + + using ElementC = typename EpilogueOutputOp::ElementOutput; + using LayoutC = LayoutA; + + using ElementAccumulator = typename EpilogueOutputOp::ElementAccumulator; + using ElementCompute = typename EpilogueOutputOp::ElementCompute; + + using WarpMmaOperator = typename Mma::Policy::Operator; + + using ArchMmaOperator = typename WarpMmaOperator::ArchMmaOperator; + using MathOperator = typename ArchMmaOperator::Operator; + + using OperatorClass = typename WarpMmaOperator::OperatorClass; + using ArchTag = typename WarpMmaOperator::ArchTag; + + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename WarpMmaOperator::Shape; + using InstructionShape = typename ArchMmaOperator::Shape; + + static int const kStages = Mma::kStages; + static IteratorAlgorithm const kIteratorAlgorithm = Mma::IteratorA::kIteratorAlgorithm; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + using TensorRefA = typename Mma::IteratorA::TensorRef; + using TensorRefB = typename Mma::IteratorB::TensorRef; + using TensorRefScaleBias = typename Mma::IteratorScaleBias::TensorRef; + using TensorRefC = mctlass::TensorRef; + + /// Check iterator A and B convolution dimension are the same and + // set device::ImplicitGemmConvolution::kConvDim + static_assert(Mma::IteratorA::kConvDim == Mma::IteratorB::kConvDim, + "Convolution on different different dimensions is not supported"); + static int const kConvDim = Mma::IteratorA::kConvDim; + + /// Conv dimension and problem size structure (Conv2d or Conv3d) + using ConvProblemSize = ConvProblemSize_; + + static conv::GroupMode const kGroupMode = conv::GroupMode::kNone; + + /// Wgrad C stride idx for implicit gemm algorithm + // Conv2d row-major matrix C (KxRSC) + // Conv3d row-major matrix C (KxTRSC) + static int const kWgradCStrideIdx = + platform::is_same::value ? 2 : 3; + + /// This chooses the appropriate stride element of the C tensor. + static int const kTensorCStrideIdx = + (kConvolutionalOperator == conv::Operator::kWgrad ? kWgradCStrideIdx : 0); + + // + // + // + using ConvOutputIteratorParameter = epilogue::threadblock::ConvOutputIteratorParameter< + LayoutC, + typename Epilogue::OutputTileIterator::Layout, + TensorRefC, + ConvOperator, + ConvProblemSize + >; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + ConvProblemSize problem_size; + TensorRefA ref_A; + TensorRefB ref_B; + TensorRefScaleBias ref_scale; + TensorRefScaleBias ref_bias; + TensorRefC ref_C; + TensorRefC ref_D; + typename EpilogueOutputOp::Params output_op; + SplitKMode split_k_mode; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + MCTLASS_HOST_DEVICE + Arguments( + ConvProblemSize const & problem_size + ): + problem_size(problem_size) { } + + MCTLASS_HOST_DEVICE + Arguments( + ConvProblemSize const & problem_size, + TensorRefA const & ref_A, + TensorRefB const & ref_B, + TensorRefScaleBias const & ref_scale, + TensorRefScaleBias const & ref_bias, + TensorRefC const & ref_C, + TensorRefC const & ref_D, + typename EpilogueOutputOp::Params const & output_op, + SplitKMode const & split_k_mode = SplitKMode::kSerial + ): + problem_size(problem_size), + ref_A(ref_A), + ref_B(ref_B), + ref_scale(ref_scale), + ref_bias(ref_bias), + ref_C(ref_C), + ref_D(ref_D), + output_op(output_op), + split_k_mode(split_k_mode) + { + + } + + }; + + /// Parameters structure + struct Params { + ConvProblemSize problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + gemm::GemmCoord implicit_gemm_problem_size; + int swizzle_log_tile; + int gemm_k_iterations; + typename Mma::IteratorA::Params iterator_A; + typename Mma::IteratorA::Element const *ptr_A; + typename Mma::IteratorB::Params iterator_B; + typename Mma::IteratorB::Element const *ptr_B; + typename Mma::IteratorScaleBias::Params iterator_scale_bias; + typename Mma::IteratorScaleBias::Element const *ptr_scale; + typename Mma::IteratorScaleBias::Element const *ptr_bias; + typename Epilogue::OutputTileIterator::Params iterator_C; + typename Epilogue::OutputTileIterator::Element *ptr_C; + typename Epilogue::OutputTileIterator::Params iterator_D; + typename Epilogue::OutputTileIterator::Element *ptr_D; + typename EpilogueOutputOp::Params output_op; + int *semaphore; + SplitKMode split_k_mode; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): swizzle_log_tile(0), gemm_k_iterations(0) { } + + /// + MCTLASS_HOST_DEVICE + Params( + Arguments const &args, + int *semaphore = nullptr + ): + problem_size(args.problem_size), + implicit_gemm_problem_size(mctlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, args.problem_size)), + iterator_A(Mma::IteratorA::getParams(args.problem_size, args.ref_A.layout())), + ptr_A(args.ref_A.data()), + iterator_B(args.problem_size, args.ref_B.layout()), + ptr_B(args.ref_B.data()), + iterator_scale_bias(args.problem_size, args.ref_scale.layout()), + ptr_scale(args.ref_scale.data()), + ptr_bias(args.ref_bias.data()), + iterator_C(ConvOutputIteratorParameter::layout(args.ref_C)), + ptr_C(args.ref_C.data()), + iterator_D(ConvOutputIteratorParameter::layout(args.ref_D)), + ptr_D(args.ref_D.data()), + output_op(args.output_op), + semaphore(semaphore), + split_k_mode(args.split_k_mode) + { + gemm_k_iterations = implicit_gemm_k_iterations(kConvolutionalOperator, ThreadblockShape::kK, args.problem_size); + + ThreadblockSwizzle threadblock_swizzle; + + grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + implicit_gemm_problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.problem_size.split_k_slices); + + swizzle_log_tile = threadblock_swizzle.get_log_tile(grid_tiled_shape); + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + ImplicitGemmConvolutionFusion() { } + + /// Executes one ImplicitGEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_idx = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_idx.m() || + params.grid_tiled_shape.n() <= threadblock_tile_idx.n()) { + + return; + } + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A operand + typename Mma::IteratorA iterator_A( + params.iterator_A, + params.problem_size, + params.ptr_A, + thread_idx, + MatrixCoord( + threadblock_tile_idx.m() * Mma::Shape::kM, + threadblock_tile_idx.k() * Mma::Shape::kK + ) + ); + + // Construct iterators to B operand + typename Mma::IteratorB iterator_B( + params.iterator_B, + params.problem_size, + params.ptr_B, + thread_idx, + MatrixCoord( + threadblock_tile_idx.k() * Mma::Shape::kK, + threadblock_tile_idx.n() * Mma::Shape::kN + ) + ); + + // Construct iterators to A scale/bias vector + typename Mma::IteratorScaleBias iterator_scale_bias( + params.iterator_scale_bias, + params.problem_size, + params.ptr_scale, + params.ptr_bias, + thread_idx, + MatrixCoord( + 0, (kConvolutionalOperator == conv::Operator::kFprop) ? + (threadblock_tile_idx.k() * Mma::Shape::kK) : + // Wgrad + (threadblock_tile_idx.n() * Mma::Shape::kN) + ) + ); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + mma(params.gemm_k_iterations, accumulators, iterator_A, + iterator_B, iterator_scale_bias, accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // Construct the semaphore. + int block_idx = threadblock_tile_idx.m() + threadblock_tile_idx.n() * params.grid_tiled_shape.m(); + + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + // Compute logical position within grid + threadblock_tile_idx = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_idx.k(), params.grid_tiled_shape.k()); + } + + MatrixCoord threadblock_offset( + threadblock_tile_idx.m() * Mma::Shape::kM, + threadblock_tile_idx.n() * Mma::Shape::kN + ); + + // Tile iterator writing to destination tensor + typename Epilogue::OutputTileIterator iterator_D( + params.iterator_D, + params.ptr_D, + ConvOutputIteratorParameter::extent(params.problem_size), + thread_idx, + threadblock_offset + ); + + // Tile iterator reading from source accumulator tensor + typename Epilogue::OutputTileIterator iterator_C( + params.iterator_C, + params.ptr_C, + ConvOutputIteratorParameter::extent(params.problem_size), + thread_idx, + threadblock_offset + ); + + // Construct the epilogue + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_idx.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_idx.k()); + + } + // Each split-k-slice writes to a unique tensor location + else if (params.split_k_mode == SplitKMode::kParallel) { + iterator_D.add_pointer_offset(threadblock_tile_idx.k() * + mctlass::conv::implicit_gemm_tensor_c_size(ConvOperator, params.problem_size)); + } + + // Run efficient epilogue + epilogue(output_op, iterator_D, accumulators, iterator_C); + + // + // Release the semaphore + // + + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_idx.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_idx.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution_strided_dgrad.h b/csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution_strided_dgrad.h new file mode 100644 index 0000000..c5198ad --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution_strided_dgrad.h @@ -0,0 +1,492 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined Implicit GEMM kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/aligned_buffer.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/semaphore.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/epilogue/threadblock/output_iterator_parameter.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + conv::Operator ConvOperator, ///! Convolutional operator (Fprop, Dgrad, Wgrad) + typename ConvProblemSize_ = Conv2dProblemSize ///! Convolutional operator on 2D or 3D problem +> +struct ImplicitGemmConvolutionStridedDgrad { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static Operator const kConvolutionalOperator = ConvOperator; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename EpilogueOutputOp::ElementOutput; + + /// Set output tensor C layout + using LayoutC = LayoutA; + + using ElementAccumulator = typename EpilogueOutputOp::ElementAccumulator; + using ElementCompute = typename EpilogueOutputOp::ElementCompute; + + using WarpMmaOperator = typename Mma::Policy::Operator; + + using ArchMmaOperator = typename WarpMmaOperator::ArchMmaOperator; + using MathOperator = typename ArchMmaOperator::Operator; + + using OperatorClass = typename WarpMmaOperator::OperatorClass; + using ArchTag = typename WarpMmaOperator::ArchTag; + + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename WarpMmaOperator::Shape; + using InstructionShape = typename ArchMmaOperator::Shape; + + static int const kStages = Mma::kStages; + static IteratorAlgorithm const kIteratorAlgorithm = Mma::IteratorA::kIteratorAlgorithm; + static StrideSupport const kStrideSupport = Mma::IteratorA::kStrideSupport; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 64 * WarpCount::kCount; + + using TensorRefA = typename Mma::IteratorA::TensorRef; + using TensorRefB = typename Mma::IteratorB::TensorRef; + using TensorRefC = mctlass::TensorRef; + + /// Check iterator A and B convolution dimension are the same and + // set device::ImplicitGemmConvolution::kConvDim + static_assert(Mma::IteratorA::kConvDim == Mma::IteratorB::kConvDim, + "Convolution on different different dimensions is not supported"); + static int const kConvDim = Mma::IteratorA::kConvDim; + + /// Conv dimension and problem size structure (Conv2d or Conv3d) + using ConvProblemSize = ConvProblemSize_; + + static conv::GroupMode const kGroupMode = conv::GroupMode::kNone; + + /// Wgrad C stride idx for implicit gemm algorithm + // Conv2d row-major matrix C (KxRSC) + // Conv3d row-major matrix C (KxTRSC) + static int const kWgradCStrideIdx = + platform::is_same::value ? 2 : 3; + + /// This chooses the appropriate stride element of the C tensor. + static int const kTensorCStrideIdx = + (kConvolutionalOperator == conv::Operator::kWgrad ? kWgradCStrideIdx : 0); + + // Strided dgrad uses a specialized threadblock swizzle for functionality and performance + static_assert((platform::is_same::value) || + (platform::is_same>::value) || + (platform::is_same>::value) || + (platform::is_same>::value), + "Needs ThreadblockSwizzle type specialized for strided dgrad"); + + // + // + // + using ConvOutputIteratorParameter = epilogue::threadblock::ConvOutputIteratorParameter< + LayoutC, + typename Epilogue::OutputTileIterator::Layout, + TensorRefC, + ConvOperator, + ConvProblemSize + >; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + ConvProblemSize problem_size; + TensorRefA ref_A; + TensorRefB ref_B; + TensorRefC ref_C; + TensorRefC ref_D; + typename EpilogueOutputOp::Params output_op; + SplitKMode split_k_mode; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + MCTLASS_HOST_DEVICE + Arguments( + ConvProblemSize const & problem_size + ): + problem_size(problem_size) { } + + MCTLASS_HOST_DEVICE + Arguments( + ConvProblemSize const & problem_size, + TensorRefA const & ref_A, + TensorRefB const & ref_B, + TensorRefC const & ref_C, + TensorRefC const & ref_D, + typename EpilogueOutputOp::Params const & output_op, + SplitKMode const & split_k_mode = SplitKMode::kSerial + ): + problem_size(problem_size), + ref_A(ref_A), + ref_B(ref_B), + ref_C(ref_C), + ref_D(ref_D), + output_op(output_op), + split_k_mode(split_k_mode) + { + + } + + }; + + /// Parameters structure + struct Params { + ConvProblemSize problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + FastDivmod stride_h_divmod; + FastDivmod stride_w_divmod; + int gemm_k_iterations; + typename Mma::IteratorA::Params iterator_A; + typename Mma::IteratorA::Element const *ptr_A; + typename Mma::IteratorB::Params iterator_B; + typename Mma::IteratorB::Element const *ptr_B; + typename Epilogue::OutputTileIterator::Params iterator_C; + typename Epilogue::OutputTileIterator::Element *ptr_C; + typename Epilogue::OutputTileIterator::Params iterator_D; + typename Epilogue::OutputTileIterator::Element *ptr_D; + typename EpilogueOutputOp::Params output_op; + int *semaphore; + SplitKMode split_k_mode; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): gemm_k_iterations(0) { } + + /// + MCTLASS_HOST_DEVICE + Params( + Arguments const &args, + int *semaphore = nullptr + ): + problem_size(args.problem_size), + stride_h_divmod(args.problem_size.stride_h), + stride_w_divmod(args.problem_size.stride_w), + iterator_A(Mma::IteratorA::getParams(args.problem_size, args.ref_A.layout())), + ptr_A(args.ref_A.data()), + iterator_B(args.problem_size, args.ref_B.layout()), + ptr_B(args.ref_B.data()), + iterator_C(ConvOutputIteratorParameter::layout(args.ref_C), args.problem_size, ThreadblockShape::kM), + ptr_C(args.ref_C.data()), + iterator_D(ConvOutputIteratorParameter::layout(args.ref_D), args.problem_size, ThreadblockShape::kM), + ptr_D(args.ref_D.data()), + output_op(args.output_op), + semaphore(semaphore), + split_k_mode(args.split_k_mode) + { + gemm_k_iterations = implicit_gemm_k_iterations(kConvolutionalOperator, ThreadblockShape::kK, args.problem_size); + + ThreadblockSwizzle threadblock_swizzle; + + grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + kConvolutionalOperator, + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.problem_size.split_k_slices); + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + ImplicitGemmConvolutionStridedDgrad() { } + + /// Executes one ImplicitGEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_idx = + threadblock_swizzle.get_tile_offset(params.grid_tiled_shape); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_idx.m() || + params.grid_tiled_shape.n() <= threadblock_tile_idx.n()) { + + return; + } + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Compute starting filter position for strided dgrad + int tile_m_per_filter = strided_dgrad_tile_m_per_filter(params.problem_size, + ThreadblockShape::kM); + int filter_tile_m = (threadblock_tile_idx.m() / tile_m_per_filter); + + + // The subsequent fast_divmod() operations are equivalent to the following logical computation: + // + // int start_r = filter_tile_m / (params.problem_size.stride_w); + // int start_s = filter_tile_m % (params.problem_size.stride_w); + + int start_r, start_s; + params.stride_w_divmod(start_r, start_s, filter_tile_m); + + int filter_r = start_r; + int filter_s = start_s; + + if (params.problem_size.mode == Mode::kConvolution) { + filter_r = (params.problem_size.R - 1 - filter_r); + filter_s = (params.problem_size.S - 1 - filter_s); + } + + // Starting h, w positions for filter position in gemm_k=0 + int start_h, start_w; + strided_dgrad_starting_coords( + params.problem_size, + params.stride_h_divmod, params.stride_w_divmod, + filter_r, filter_s, + start_h, start_w); + + if (start_h >= params.problem_size.H || start_w >= params.problem_size.W) { + return; + } + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + int lane_idx = threadIdx.x % 64; + + // Check if CTA contributes valid MMA (Dy * w) and accumulator will be non-zero after MMA + if (start_r < params.problem_size.R && start_s < params.problem_size.S) { + // Scale gemm_k_iterations for strided dgrad + int gemm_k_iterations = (params.gemm_k_iterations / (params.problem_size.R * params.problem_size.S) + ) * params.problem_size.num_gemm_k_filter_positions(start_r, start_s); + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.iterator_A, + params.problem_size, + params.ptr_A, + thread_idx, + params.stride_h_divmod, params.stride_w_divmod, + start_r, start_s, + MatrixCoord( + threadblock_tile_idx.m() * Mma::Shape::kM, + threadblock_tile_idx.k() * Mma::Shape::kK + ) + ); + + typename Mma::IteratorB iterator_B( + params.iterator_B, + params.problem_size, + params.ptr_B, + thread_idx, + start_r, start_s, + MatrixCoord( + threadblock_tile_idx.k() * Mma::Shape::kK, + threadblock_tile_idx.n() * Mma::Shape::kN + ) + ); + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + // Compute threadblock-scoped matrix multiply-add + mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators); + } + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // Construct the semaphore. + int block_idx = threadblock_tile_idx.m() + threadblock_tile_idx.n() * params.grid_tiled_shape.m(); + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + // Compute logical position within grid + threadblock_tile_idx = + threadblock_swizzle.get_tile_offset(params.grid_tiled_shape); + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_idx.k(), params.grid_tiled_shape.k()); + } + + MatrixCoord threadblock_offset( + threadblock_tile_idx.m() * Mma::Shape::kM, + threadblock_tile_idx.n() * Mma::Shape::kN + ); + + // Tile iterator writing to destination tensor + typename Epilogue::OutputTileIterator iterator_D( + params.iterator_D, + params.ptr_D, + ConvOutputIteratorParameter::extent(params.problem_size), + thread_idx, + params.stride_h_divmod, params.stride_w_divmod, + start_r, start_s, + threadblock_offset + ); + + // Construct the epilogue + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + if (output_op.is_source_needed()) + { + // Tile iterator reading from source accumulator tensor + typename Epilogue::OutputTileIterator iterator_C( + params.iterator_C, + params.ptr_C, + ConvOutputIteratorParameter::extent(params.problem_size), + thread_idx, + params.stride_h_divmod, params.stride_w_divmod, + start_r, start_s, + threadblock_offset); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_idx.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_idx.k()); + } + + // Run epilogue with addend source iterator + epilogue(output_op, iterator_D, accumulators, iterator_C); + } + else + { + // Run epilogue without addend source iterator + epilogue(output_op, iterator_D, accumulators); + } + + // + // Release the semaphore + // + + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_idx.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_idx.k() + 1; + } + + semaphore.release(lock); + } + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution_with_fused_epilogue.h b/csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution_with_fused_epilogue.h new file mode 100644 index 0000000..ece45cf --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/kernel/implicit_gemm_convolution_with_fused_epilogue.h @@ -0,0 +1,499 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined Implicit GEMM kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/aligned_buffer.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/semaphore.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/epilogue/threadblock/output_iterator_parameter.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + conv::Operator ConvOperator, ///! Convolutional operator (Fprop, Dgrad, Wgrad) + typename ConvProblemSize_ = Conv2dProblemSize ///! Convolutional operator on 2D or 3D problem +> +struct ImplicitGemmConvolutionWithFusedEpilogue { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static Operator const kConvolutionalOperator = ConvOperator; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename EpilogueOutputOp::ElementOutput; + + /// Set output tensor C layout + using LayoutC = LayoutA; + + using ElementAccumulator = typename EpilogueOutputOp::ElementAccumulator; + using ElementCompute = typename EpilogueOutputOp::ElementCompute; + + using WarpMmaOperator = typename Mma::Policy::Operator; + + using ArchMmaOperator = typename WarpMmaOperator::ArchMmaOperator; + using MathOperator = typename ArchMmaOperator::Operator; + + using OperatorClass = typename WarpMmaOperator::OperatorClass; + using ArchTag = typename WarpMmaOperator::ArchTag; + + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename WarpMmaOperator::Shape; + using InstructionShape = typename ArchMmaOperator::Shape; + + static int const kStages = Mma::kStages; + static IteratorAlgorithm const kIteratorAlgorithm = Mma::IteratorA::kIteratorAlgorithm; + static StrideSupport const kStrideSupport = Mma::IteratorA::kStrideSupport; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + using TensorRefA = typename Mma::IteratorA::TensorRef; + using TensorRefB = typename Mma::IteratorB::TensorRef; + using TensorRefC = mctlass::TensorRef; + + /// Check iterator A and B convolution dimension are the same and + // set device::ImplicitGemmConvolution::kConvDim + static_assert(Mma::IteratorA::kConvDim == Mma::IteratorB::kConvDim, + "Convolution on different different dimensions is not supported"); + static int const kConvDim = Mma::IteratorA::kConvDim; + + /// Conv dimension and problem size structure (Conv2d or Conv3d) + using ConvProblemSize = ConvProblemSize_; + + static conv::GroupMode const kGroupMode = conv::GroupMode::kNone; + + /// Wgrad C stride idx for implicit gemm algorithm + // Conv2d row-major matrix C (KxRSC) + // Conv3d row-major matrix C (KxTRSC) + static int const kWgradCStrideIdx = + platform::is_same::value ? 2 : 3; + + /// This chooses the appropriate stride element of the C tensor. + static int const kTensorCStrideIdx = + (kConvolutionalOperator == conv::Operator::kWgrad ? kWgradCStrideIdx : 0); + + // + // + // + using ConvOutputIteratorParameter = epilogue::threadblock::ConvOutputIteratorParameter< + LayoutC, + typename Epilogue::OutputTileIterator::Layout, + TensorRefC, + ConvOperator, + ConvProblemSize + >; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + ConvProblemSize problem_size; + TensorRefA ref_A; + TensorRefB ref_B; + TensorRefC ref_C; + TensorRefC ref_D; + + typename EpilogueOutputOp::Params output_op; + SplitKMode split_k_mode; + + void * ptr_Vector; + void * ptr_Tensor; + + typename LayoutC::Stride::Index ldr; + typename LayoutC::Stride::Index ldt; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + MCTLASS_HOST_DEVICE + Arguments( + ConvProblemSize const & problem_size + ): + problem_size(problem_size) { } + + MCTLASS_HOST_DEVICE + Arguments( + ConvProblemSize const & problem_size, + TensorRefA const & ref_A, + TensorRefB const & ref_B, + TensorRefC const & ref_C, + TensorRefC const & ref_D, + typename EpilogueOutputOp::Params const & output_op, + SplitKMode const & split_k_mode = SplitKMode::kSerial, + void * ptr_Vector = nullptr, + void * ptr_Tensor = nullptr, + typename LayoutC::Stride::Index ldr = 0, + typename LayoutC::Stride::Index ldt = 0 + ): + problem_size(problem_size), + ref_A(ref_A), + ref_B(ref_B), + ref_C(ref_C), + ref_D(ref_D), + output_op(output_op), + split_k_mode(split_k_mode), + ptr_Vector(ptr_Vector), + ptr_Tensor(ptr_Tensor), + ldr(ldr), + ldt(ldt) + { + + } + + }; + + /// Parameters structure + struct Params { + ConvProblemSize problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + gemm::GemmCoord implicit_gemm_problem_size; + int swizzle_log_tile; + + int gemm_k_iterations; + typename Mma::IteratorA::Params iterator_A; + typename Mma::IteratorA::Element const *ptr_A; + typename Mma::IteratorB::Params iterator_B; + typename Mma::IteratorB::Element const *ptr_B; + typename Epilogue::OutputTileIterator::Params iterator_C; + typename Epilogue::OutputTileIterator::Element *ptr_C; + typename Epilogue::OutputTileIterator::Params iterator_D; + typename Epilogue::OutputTileIterator::Element *ptr_D; + typename EpilogueOutputOp::Params output_op; + int *semaphore; + SplitKMode split_k_mode; + + typename Epilogue::TensorTileIterator::Params params_Tensor; + void * ptr_Vector; + typename LayoutC::Stride::Index ldr; + void * ptr_Tensor; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + swizzle_log_tile(0), + gemm_k_iterations(0), + ptr_Vector(nullptr), + ldr(0), + ptr_Tensor(nullptr) + { } + + /// + MCTLASS_HOST_DEVICE + Params( + Arguments const &args, + int *semaphore = nullptr + ): + problem_size(args.problem_size), + implicit_gemm_problem_size(mctlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, args.problem_size)), + iterator_A(Mma::IteratorA::getParams(args.problem_size, args.ref_A.layout())), + ptr_A(args.ref_A.data()), + iterator_B(args.problem_size, args.ref_B.layout()), + ptr_B(args.ref_B.data()), + iterator_C(ConvOutputIteratorParameter::layout(args.ref_C)), + ptr_C(args.ref_C.data()), + iterator_D(ConvOutputIteratorParameter::layout(args.ref_D)), + ptr_D(args.ref_D.data()), + output_op(args.output_op), + semaphore(semaphore), + split_k_mode(args.split_k_mode), + params_Tensor(args.ldt), + ptr_Vector(args.ptr_Vector), + ldr(args.ldr), + ptr_Tensor(args.ptr_Tensor) + + { + gemm_k_iterations = implicit_gemm_k_iterations(kConvolutionalOperator, ThreadblockShape::kK, args.problem_size); + + ThreadblockSwizzle threadblock_swizzle; + + grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + implicit_gemm_problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.problem_size.split_k_slices); + + swizzle_log_tile = threadblock_swizzle.get_log_tile(grid_tiled_shape); + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + ImplicitGemmConvolutionWithFusedEpilogue() { } + + /// Executes one ImplicitGEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_idx = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_idx.m() || + params.grid_tiled_shape.n() <= threadblock_tile_idx.n()) { + + return; + } + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.iterator_A, + params.problem_size, + params.ptr_A, + thread_idx, + MatrixCoord( + threadblock_tile_idx.m() * Mma::Shape::kM, + threadblock_tile_idx.k() * Mma::Shape::kK + ) + ); + + typename Mma::IteratorB iterator_B( + params.iterator_B, + params.problem_size, + params.ptr_B, + thread_idx, + MatrixCoord( + threadblock_tile_idx.k() * Mma::Shape::kK, + threadblock_tile_idx.n() * Mma::Shape::kN + ) + ); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + mma(params.gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // Construct the semaphore. + int block_idx = threadblock_tile_idx.m() + threadblock_tile_idx.n() * params.grid_tiled_shape.m(); + + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + // Compute logical position within grid + threadblock_tile_idx = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_idx.k(), params.grid_tiled_shape.k()); + } + + MatrixCoord threadblock_offset( + threadblock_tile_idx.m() * Mma::Shape::kM, + threadblock_tile_idx.n() * Mma::Shape::kN + ); + + // Tile iterator writing to destination tensor + typename Epilogue::OutputTileIterator iterator_D( + params.iterator_D, + params.ptr_D, + ConvOutputIteratorParameter::extent(params.problem_size), + thread_idx, + threadblock_offset + ); + + // Tile iterator reading from source accumulator tensor + typename Epilogue::OutputTileIterator iterator_C( + params.iterator_C, + params.ptr_C, + ConvOutputIteratorParameter::extent(params.problem_size), + thread_idx, + threadblock_offset + ); + + typename Epilogue::ElementTensor *ptr_Tensor = + static_cast(params.ptr_Tensor); + + // Define the reduction output pointer and move to the appropriate place + typename Epilogue::ElementVector *ptr_Vector = + static_cast(params.ptr_Vector); + + // Additional tensor to load from + typename Epilogue::TensorTileIterator tensor_iterator( + params.params_Tensor, + // Only the final block outputs Tensor + ((params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) && + (params.grid_tiled_shape.k() != threadblock_tile_idx.k() + 1)) + ? nullptr + : ptr_Tensor, + ConvOutputIteratorParameter::extent(params.problem_size), + thread_idx, + threadblock_offset); + + // Construct the epilogue + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Move to appropriate location for this output tile + if (ptr_Vector) { + ptr_Vector += threadblock_offset.column() + threadblock_tile_idx.m() * params.ldr; + } + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_idx.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_idx.k()); + + } + // Each split-k-slice writes to a unique tensor location + else if (params.split_k_mode == SplitKMode::kParallel) { + iterator_D.add_pointer_offset(threadblock_tile_idx.k() * + mctlass::conv::implicit_gemm_tensor_c_size(ConvOperator, params.problem_size)); + } + + // Execute the epilogue operator to update the destination tensor. + epilogue(output_op, + // Only the final block uses Vector + ((params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) && + (params.grid_tiled_shape.k() != threadblock_tile_idx.k() + 1)) + ? nullptr + : ptr_Vector, + iterator_D, + accumulators, + iterator_C, + tensor_iterator, + ConvOutputIteratorParameter::extent(params.problem_size), + threadblock_offset); + + // + // Release the semaphore + // + + if (params.split_k_mode == SplitKMode::kSerial && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_idx.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_idx.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/thread/depthwise_mma.h b/csrc/mctlass/include/mctlass/conv/thread/depthwise_mma.h new file mode 100644 index 0000000..1900e95 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/thread/depthwise_mma.h @@ -0,0 +1,325 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates exposing architecture support for depthwise convolution +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/arch/mma.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/thread/mma.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// MMA operation +template < + /// Size of the matrix product (concept: GemmShape) + typename Shape_, + /// Number of threads participating + int kThreads_, + /// Data type of A elements + typename ElementA, + /// Data type of B elements + typename ElementB, + /// Element type of C matrix + typename ElementC, + /// Inner product operator + typename Operator +> +struct ElementwiseInnerProduct; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// General implementation +template < + /// Size of the matrix product (concept: GemmShape) + typename Shape_, + /// Data type of A elements + typename ElementA_, + /// Data type of B elements + typename ElementB_, + /// Element type of C matrix + typename ElementC_> +struct ElementwiseInnerProduct { + using Shape = Shape_; + using Operator = arch::OpMultiplyAdd; + using ElementC = ElementC_; + + MCTLASS_HOST_DEVICE + void operator()(Array &d, + Array const &a, + Array const &b, + Array const &c) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Shape::kN; ++i) { + d[i] = a[i] * b[i] + c[i]; + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Specialization of half_t +template <> +struct ElementwiseInnerProduct< + gemm::GemmShape<2, 2, 1>, + 1, + half_t, + half_t, + half_t, + arch::OpMultiplyAdd> { + + using Shape = gemm::GemmShape<2, 2, 1>; + using Operator = arch::OpMultiplyAdd; + using ElementC = half_t; + + MCTLASS_HOST_DEVICE + void operator()( + Array &d, + Array const &a, + Array const &b, + Array const &c + ) { + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 600)) + + __half2 const & A = reinterpret_cast<__half2 const &>(a); + __half2 const & B = reinterpret_cast<__half2 const &>(b); + __half2 const & C = reinterpret_cast<__half2 const &>(c); + + __half2 tmp_D = __hfma2(A, B, C); + + d = reinterpret_cast const &>(tmp_D); + +#else + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 2; ++i) { + d[i] = a[i] * b[i] + c[i]; + } +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape, + /// Data type of A elements + typename ElementA, + /// Data type of B elements + typename ElementB, + /// Element type of C matrix + typename ElementC, + /// Concept: arch::OpMultiplyAdd or arch::Mma<> + typename Operator = arch::OpMultiplyAdd, + /// Used for partial specialization + typename Enable = bool +> +struct DepthwiseDirectConvElementwiseInnerProduct; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Gemplate that handles all packed matrix layouts +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename ElementA_, + /// Data type of B elements + typename ElementB_, + /// Element type of C matrix + typename ElementC_, + /// Operator used to compute GEMM + typename Operator_ +> +struct DepthwiseDirectConvElementwiseInnerProductGeneric { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// Data type of operand A + using ElementA = ElementA_; + + /// Data type of operand B + using ElementB = ElementB_; + + /// Element type of operand C + using ElementC = ElementC_; + + /// Underlying mathematical operator + using Operator = Operator_; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Instruction + using MmaOp = mctlass::conv::thread::ElementwiseInnerProduct< + gemm::GemmShape, + 1, + ElementA, + ElementB, + ElementC, + Operator>; + + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + Array *ptr_D = reinterpret_cast *>(&D); + Array const *ptr_A = + reinterpret_cast const *>(&A); + Array const *ptr_B = + reinterpret_cast const *>(&B); + + MmaOp mma_op; + + // Copy accumulators + D = C; + + // Compute matrix product + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Shape::kN / MmaOp::Shape::kN; ++n) { + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Shape::kM; ++m) { + + Array tmpD = ptr_D[m * Shape::kN / MmaOp::Shape::kN + n]; + Array tmpA = ptr_A[m * Shape::kN / MmaOp::Shape::kN + n]; + Array tmpB = ptr_B[n]; + + mma_op(tmpD, tmpA, tmpB, tmpD); + + ptr_D[m * Shape::kN / MmaOp::Shape::kN + n] = tmpD; + + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename ElementA_, + /// Data type of B elements + typename ElementB_, + /// Element type of C matrix + typename ElementC_ +> +struct DepthwiseDirectConvElementwiseInnerProduct< + Shape_, + ElementA_, + ElementB_, + ElementC_, + arch::OpMultiplyAdd + > { + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// Data type of operand A + using ElementA = ElementA_; + + /// Data type of operand B + using ElementB = ElementB_; + + /// Element type of operand C + using ElementC = ElementC_; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + /// A operand storage + using FragmentA = + Array; // output_tile_size per thread * groups_per_thread + + /// B operand storage + using FragmentB = Array; // 1 * groups_per_thread + + /// C operand storage + using FragmentC = + Array; // output_tile_size per thread * groups_per_thread + + static bool const use_optimized = 0; + + using ArchMmaOperator = DepthwiseDirectConvElementwiseInnerProductGeneric; + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + ArchMmaOperator mma; + + mma(D, A, B, C); + + } +}; + +} // namespace thread +} // namespace conv +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_filter_tile_access_iterator_analytic.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_filter_tile_access_iterator_analytic.h new file mode 100644 index 0000000..f7b421b --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_filter_tile_access_iterator_analytic.h @@ -0,0 +1,485 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (filter tile) + matrix from memory. + + This iterator assumes TensorNHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + conv::StrideSupport StrideSupport_ = conv::StrideSupport::kUnity, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dDgradFilterTileAccessIteratorAnalytic; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Conv2dDgradFilterTileAccessIteratorAnalytic strided dgrad needs special handling to skip MMAs +// on non-contributing w positions +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + typename AccessType_ +> +class Conv2dDgradFilterTileAccessIteratorAnalytic < + Shape_, + Element_, + ThreadMap_, + conv::StrideSupport::kStrided, + AccessType_ +> { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNHWC; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + static_assert(sizeof_bits::value >= 8, + "DGRAD requires elements of size 8b or larger."); + + // + // Parameters structure + // + + using Params = Conv2dAnalyticParams; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + // For a fixed filter position (r,s) find and fill offset_k_, offset_c_ in strided and contiguous dimension + int filter_r_; + int filter_s_; + int start_r_; + int start_s_; + int offset_k_[ThreadMap::Iterations::kStrided]; + int offset_c_[ThreadMap::Iterations::kContiguous]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dDgradFilterTileAccessIteratorAnalytic( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + int start_r, int start_s, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + filter_r_(start_r), + filter_s_(start_s), + start_r_(start_r), + start_s_(start_s) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + offset_c_[c] = threadblock_offset.column() + thread_coord.contiguous() + + c * ThreadMap::Delta::kContiguous; + } + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_k_[s] = + threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + } + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + // Moves filter_s + filter_s_ += problem_size_.stride_w; + if (filter_s_ < problem_size_.S) { + return; + } + // Restore filter_s + filter_s_ = start_s_; + + // Move filter_r + filter_r_ += problem_size_.stride_h; + if (filter_r_ < problem_size_.R) { + return; + } + // Restore filter_r + filter_r_ = start_r_; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_k_[s] += Shape::kRow * problem_size_.split_k_slices; + } + } + + /// Returns the coordinate in the filter tensor w that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + int k = offset_k_[iteration_strided_]; + int c = offset_c_[iteration_contiguous_] + iteration_vector_ * AccessType::kElements; + + return TensorCoord(k, filter_r_, filter_s_, c); + } + + /// Returns true if the current coordinate is within the filter tensor w + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + return coord.n() < problem_size_.K && coord.c() < problem_size_.C; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dDgradFilterTileAccessIteratorAnalytic &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Conv2dDgradFilterTileAccessIteratorAnalytic unity strided dgrad is more performant for dgrad +// on problem sizes with stride = {1x1} +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + typename AccessType_ +> +class Conv2dDgradFilterTileAccessIteratorAnalytic < + Shape_, + Element_, + ThreadMap_, + conv::StrideSupport::kUnity, + AccessType_ +>{ +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNHWC; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kUnity; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + static_assert(sizeof_bits::value >= 8, + "DGRAD requires elements of size 8b or larger."); + + // + // Parameters structure + // + + using Params = Conv2dAnalyticParams; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + // For a fixed filter position (r,s) find and fill offset_k_, offset_c_ in strided and contiguous dimension + int filter_r_; + int filter_s_; + int offset_k_[ThreadMap::Iterations::kStrided]; + int offset_c_[ThreadMap::Iterations::kContiguous]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dDgradFilterTileAccessIteratorAnalytic( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + filter_r_(0), + filter_s_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + offset_c_[c] = threadblock_offset.column() + thread_coord.contiguous() + + c * ThreadMap::Delta::kContiguous; + } + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_k_[s] = + threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + } + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + // moves to the next tile + ++filter_s_; + if (filter_s_ < problem_size_.S) { + return; + } + filter_s_ = 0; + ++filter_r_; + if (filter_r_ < problem_size_.R) { + return; + } + filter_r_ = 0; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_k_[s] += Shape::kRow * problem_size_.split_k_slices; + } + } + + /// Returns the coordinate in the filter tensor w that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + int k = offset_k_[iteration_strided_]; + int c = offset_c_[iteration_contiguous_] + iteration_vector_ * AccessType::kElements; + + return TensorCoord(k, filter_r_, filter_s_, c); + } + + /// Returns true if the current coordinate is within the filter tensor w + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + return coord.n() < problem_size_.K && coord.c() < problem_size_.C; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dDgradFilterTileAccessIteratorAnalytic &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_filter_tile_access_iterator_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_filter_tile_access_iterator_optimized.h new file mode 100644 index 0000000..84c0533 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_filter_tile_access_iterator_optimized.h @@ -0,0 +1,619 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (filter tile) + matrix from memory. + + This iterator assumes TensorNHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" + +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + conv::StrideSupport StrideSupport_ = conv::StrideSupport::kUnity, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dDgradFilterTileAccessIteratorOptimized; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Conv2dDgradFilterTileAccessIteratorOptimized unity strided dgrad is more performant for dgrad +// on problem sizes with stride = {1x1} +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + typename AccessType_ +> +class Conv2dDgradFilterTileAccessIteratorOptimized < + Shape_, + Element_, + ThreadMap_, + conv::StrideSupport::kStrided, + AccessType_ + > { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNHWC; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + // + // Parameters structure + // + + struct Params : Conv2dStridedDgradFilterIteratorOptimizedParams { + + // + // Methods + // + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Conv2dStridedDgradFilterIteratorOptimizedParams const &base): + Conv2dStridedDgradFilterIteratorOptimizedParams(base) { } + + MCTLASS_HOST_DEVICE + Params( + Conv2dProblemSize const &problem_size, + Layout const &layout + ): + Conv2dStridedDgradFilterIteratorOptimizedParams( + problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}, + ThreadMap::kThreads, + ThreadMap::kElementsPerAccess, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided} + ) { } + + }; + +private: + + Conv2dStridedDgradFilterIteratorOptimizedParams const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + uint32_t predicates_[kAccessesPerVector]; + int filter_k_; + int filter_r_; + int filter_s_; + + int start_r_; + int start_s_; + + int64_t reset_bytes_s_; + int64_t reset_bytes_r_; + + // + // Assertions + // + + // We map predicates into bits packed in this uint32_t container + static_assert(ThreadMap::Iterations::kStrided * + ThreadMap::Iterations::kContiguous < sizeof(predicates_) * 8, + "Currently, the number of loads per iteration is limited by the size of the predicates container."); + +public: + + MCTLASS_HOST_DEVICE + Conv2dDgradFilterTileAccessIteratorOptimized( + Conv2dStridedDgradFilterIteratorOptimizedParams const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + int start_r, int start_s, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + predicates_{0}, + filter_r_(start_r), + filter_s_(start_s), + start_r_(start_r), + start_s_(start_s) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_k_ = threadblock_offset.row() + thread_coord.strided(); + Index column = threadblock_offset.column() + thread_coord.contiguous(); + + reset_bytes_s_ = (problem_size_.num_gemm_k_filter_s(start_s_) - 1) * params_.inc_next[0]; + reset_bytes_r_ = reset_bytes_s_ + + (problem_size_.num_gemm_k_filter_r(start_r_) - 1) * params_.inc_next[1]; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int filter_k = filter_k_ + s * ThreadMap::Delta::kStrided; + int filter_c = column + c * ThreadMap::Delta::kContiguous; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + + uint32_t pred = ((filter_k < problem_size_.K && (filter_c + v * AccessType::kElements) < problem_size_.C) ? 1u : 0); + + int pred_idx = c + s * ThreadMap::Iterations::kContiguous; + + predicates_[v] |= (pred << pred_idx); + } + } + } + + TensorCoord coord{filter_k_, filter_r_, filter_s_, column}; + + pointer_ += params_.layout(coord) * sizeof_bits::value / 8; + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_DEVICE + void advance() { + + int next_idx = 0; + LongIndex reset_bytes = params_.reset_bytes; + + // Move filter_s by stride_w + filter_s_ += problem_size_.stride_w; + if (filter_s_ >= problem_size_.S) { + + // Restore filter_s + filter_s_ = start_s_; + + // Move filter_r by stride_h + filter_r_ += problem_size_.stride_h; +#if 1 + bool check = (filter_r_ < problem_size_.R); + + filter_r_ = check ? filter_r_ : start_r_; + next_idx = check ? 1 : 2; + reset_bytes += (check ? reset_bytes_s_ : reset_bytes_r_); +#else + asm volatile( + "{\n\t" + " .reg .pred %%p;\n\t" + " .reg .s64 t1;\n\t" + " setp.lt.s32 %%p, %3, %4;\n\t" + " selp.s32 %0, %3, %5, %%p;\n\t" + " selp.s32 %1, 1, 2, %%p;\n\t" + " selp.s64 t1, %6, %7, %%p;\n\t" + " add.s64 %2, %8, t1;\n\t" + "}\n" + : "=r"(filter_r_), "=r"(next_idx), "=l"(reset_bytes) + : "r"(filter_r_), "r"(problem_size_.R), "r"(start_r_), + "l"(reset_bytes_s_), "l"(reset_bytes_r_), "l"(reset_bytes)); +#endif + } + + // offset pointers by offset_bytes + pointer_ += (params_.inc_next[next_idx] - reset_bytes); + + if (next_idx == 2) { + filter_k_ += params_.filter_k_delta; + } + + // Clear predicates if needed + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + if (filter_k_ + s * ThreadMap::Delta::kStrided >= problem_size_.K) { + uint32_t kClearMask = ((1u << ThreadMap::Iterations::kContiguous) - 1) << (s * ThreadMap::Iterations::kContiguous); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + predicates_[v] = (predicates_[v] & (~kClearMask)); + } + } + } + } + + /// Returns true if the current coordinate is within the filter tensor W + MCTLASS_HOST_DEVICE + bool valid() { + LongIndex pred_idx = iteration_contiguous_ + iteration_strided_ * ThreadMap::Iterations::kContiguous; + return (predicates_[iteration_vector_] & (1u << pred_idx)); + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + return reinterpret_cast(pointer_ + + iteration_contiguous_ * ThreadMap::Delta::kContiguous * sizeof_bits::value / 8) + iteration_vector_; + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dDgradFilterTileAccessIteratorOptimized &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + + // Move to the next K coordinate within the tile + pointer_ += params_.inc_next_strided; + + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Conv2dDgradFilterTileAccessIteratorOptimized unity strided dgrad is more performant for dgrad +// on problem sizes with stride = {1x1} +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + typename AccessType_ +> +class Conv2dDgradFilterTileAccessIteratorOptimized < + Shape_, + Element_, + ThreadMap_, + conv::StrideSupport::kUnity, + AccessType_ + > { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNHWC; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kUnity; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + // + // Parameters structure + // + + struct Params : Conv2dDgradFilterIteratorOptimizedParams { + + // + // Methods + // + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Conv2dDgradFilterIteratorOptimizedParams const &base): + Conv2dDgradFilterIteratorOptimizedParams(base) { } + + MCTLASS_HOST_DEVICE + Params( + Conv2dProblemSize const &problem_size, + Layout const &layout + ): + Conv2dDgradFilterIteratorOptimizedParams( + problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}, + ThreadMap::kThreads, + ThreadMap::kElementsPerAccess, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided} + ) { } + + }; + +private: + + Conv2dDgradFilterIteratorOptimizedParams const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + uint32_t predicates_[kAccessesPerVector]; + int filter_rs_; + int filter_k_; + + // + // Assertions + // + + // We map predicates into bits packed in this uint32_t container + static_assert(ThreadMap::Iterations::kStrided * + ThreadMap::Iterations::kContiguous < sizeof(predicates_) * 8, + "Currently, the number of loads per iteration is limited by the size of the predicates container."); + +public: + + MCTLASS_HOST_DEVICE + Conv2dDgradFilterTileAccessIteratorOptimized( + Conv2dDgradFilterIteratorOptimizedParams const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + predicates_{0}, + filter_rs_(0), + filter_k_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_k_ = threadblock_offset.row() + thread_coord.strided(); + Index column = threadblock_offset.column() + thread_coord.contiguous(); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int filter_k = filter_k_ + s * ThreadMap::Delta::kStrided; + int filter_c = column + c * ThreadMap::Delta::kContiguous; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + + uint32_t pred = ((filter_k < problem_size_.K && (filter_c + v * AccessType::kElements) < problem_size_.C) ? 1u : 0); + + int pred_idx = c + s * ThreadMap::Iterations::kContiguous; + + predicates_[v] |= (pred << pred_idx); + } + } + } + + pointer_ += ( + filter_k_ * params.layout.stride()[2] + column + ) * sizeof_bits::value / 8; + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + + LongIndex next = params_.inc_next_rs; + + // moves to the next tile + ++filter_rs_; + if (filter_rs_ == params_.RS) { + + filter_rs_ = 0; + next = params_.inc_next_k; + filter_k_ += params_.filter_k_delta; + } + + // Clear predicates if needed + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + if (filter_k_ + s * ThreadMap::Delta::kStrided >= problem_size_.K) { + uint32_t kClearMask = ((1u << ThreadMap::Iterations::kContiguous) - 1) << (s * ThreadMap::Iterations::kContiguous); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + predicates_[v] = (predicates_[v] & (~kClearMask)); + } + } + } + + pointer_ += next; + } + + /// Returns true if the current coordinate is within the filter tensor W + MCTLASS_HOST_DEVICE + bool valid() { + LongIndex pred_idx = iteration_contiguous_ + iteration_strided_ * ThreadMap::Iterations::kContiguous; + return (predicates_[iteration_vector_] & (1u << pred_idx)); + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + return reinterpret_cast(pointer_ + + iteration_contiguous_ * ThreadMap::Delta::kContiguous * sizeof_bits::value / 8) + iteration_vector_; + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dDgradFilterTileAccessIteratorOptimized &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + + // Move to the next K coordinate within the tile + pointer_ += params_.inc_next_strided; + + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_output_gradient_tile_access_iterator_analytic.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_output_gradient_tile_access_iterator_analytic.h new file mode 100644 index 0000000..5ef2feb --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_output_gradient_tile_access_iterator_analytic.h @@ -0,0 +1,606 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (output gradient tile) + matrix from memory. + + This iterator assumes TensorNHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/functional.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + conv::StrideSupport StrideSupport_ = conv::StrideSupport::kStrided, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dDgradOutputGradientTileAccessIteratorAnalytic; +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Conv2dDgradOutputGradientTileAccessIteratorAnalytic strided dgrad needs special handling using +// unscaled coordinations +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + typename AccessType_ +> +class Conv2dDgradOutputGradientTileAccessIteratorAnalytic < + Shape_, + Element_, + ThreadMap_, + conv::StrideSupport::kStrided, + AccessType_ +> { +public: + + // + // Types + // + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNHWC; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + static_assert(sizeof_bits::value >= 8, + "DGRAD requires elements of size 8b or greater."); + + // + // Simpligying assertions + // + + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv2dDgradOutputGradientTileAccessIteratorAnalyticParams; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + int filter_k_; + int filter_r_; + int filter_s_; + int start_r_; + int start_s_; + + int offset_n_[ThreadMap::Iterations::kStrided]; + int offset_p_[ThreadMap::Iterations::kStrided]; + int offset_q_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dDgradOutputGradientTileAccessIteratorAnalytic( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + FastDivmod const &stride_h_divmod, FastDivmod const &stride_w_divmod, + int start_r, int start_s, + MatrixCoord const &threadblock_offset = MatrixCoord() // threadblock offset - units are whole CTA tiles + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + filter_k_(0), + filter_r_(start_r), + filter_s_(start_s), + start_r_(start_r), + start_s_(start_s) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_k_ = threadblock_offset.column() + thread_coord.contiguous(); + + int filter_r = filter_r_; + int filter_s = filter_s_; + + if (problem_size_.mode == Mode::kConvolution) { + filter_r = (problem_size_.R - 1 - filter_r); + filter_s = (problem_size_.S - 1 - filter_s); + } + + // Starting h, w positions for filter position in gemm_k=0 + int start_h, start_w; + strided_dgrad_starting_coords( + problem_size_, + stride_h_divmod, stride_w_divmod, + filter_r, filter_s, + start_h, start_w); + + // Effective P and Q for filter position required for remapping NHW rows + int P = (problem_size_.H - start_h + problem_size_.stride_h - 1) / problem_size_.stride_h; + int Q = (problem_size_.W - start_w + problem_size_.stride_w - 1) / problem_size_.stride_w; + + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + int offset_npq = (threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided) % params_.tiled_rows_per_filter; + + // (STEP 1) [reorder NHW rows to start with same filter positions] + offset_n_[s] = offset_npq / (P * Q); + int residual = offset_npq % (P * Q); + + int p = (residual / Q); + int q = (residual % Q); + + int mapped_h = (start_h + p * problem_size_.stride_h); + int mapped_w = (start_w + q * problem_size_.stride_w); + + // Access (p, q) coordinates for Dy tensor and a filter position in gemm_k=0 + // note that (h + pad_h - filter_r) and (w + pad_w - filter_s) are divisible + // by stride_h and stride_w + offset_p_[s] = (mapped_h + problem_size_.pad_h - filter_r) / problem_size_.stride_h; + offset_q_[s] = (mapped_w + problem_size_.pad_w - filter_s) / problem_size_.stride_w; + } + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + + // Move filter_s by stride_w + filter_s_ += problem_size_.stride_w; + if (filter_s_ < problem_size_.S) { + return; + } + + // Restore filter_s + filter_s_ = start_s_; + + // Move filter_r by stride_h + filter_r_ += problem_size_.stride_h; + if (filter_r_ < problem_size_.R) { + return; + } + + // Restore filter_r + filter_r_ = start_r_; + + // Move filter_k + filter_k_ += Shape_::kColumn * problem_size_.split_k_slices; + } + + /// Returns the coordinate in the output tensor Dy that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + int n = offset_n_[iteration_strided_]; + int p = offset_p_[iteration_strided_]; + int q = offset_q_[iteration_strided_]; + + int conv_sign = (problem_size_.mode == Mode::kConvolution ? 1 : -1); + + p += (conv_sign * (filter_r_ / problem_size_.stride_h)); + q += (conv_sign * (filter_s_ / problem_size_.stride_w)); + + int k = filter_k_ + iteration_vector_ * AccessType::kElements; + + return TensorCoord( + n, + p, + q, + k); + } + + + /// Returns true if the current coordinate is within the output tensor Dy + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + return + coord.n() < problem_size_.N && + coord.h() >= 0 && coord.h() < problem_size_.P && + coord.w() >= 0 && coord.w() < problem_size_.Q && + coord.c() < problem_size_.K; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dDgradOutputGradientTileAccessIteratorAnalytic &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.K % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Conv2dDgradOutputGradientTileAccessIteratorAnalytic for unity strides can be optimized by +// eliminating modulo arithmetic to compute unscaled coordinates +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + typename AccessType_ +> +class Conv2dDgradOutputGradientTileAccessIteratorAnalytic < + Shape_, + Element_, + ThreadMap_, + conv::StrideSupport::kUnity, + AccessType_ +> { +public: + + // + // Types + // + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNHWC; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kUnity; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + static_assert(sizeof_bits::value >= 8, + "DGRAD requires elements of size 8b or greater."); + + // + // Simpligying assertions + // + + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + struct Params { + + Layout layout; + + // + // Methods + // + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params( + Conv2dProblemSize const &problem_size, + Layout const &layout + ): layout(layout) { + + } + }; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + int filter_k_; + int filter_r_; + int filter_s_; + + int offset_n_[ThreadMap::Iterations::kStrided]; + int offset_w_[ThreadMap::Iterations::kStrided]; + int offset_h_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dDgradOutputGradientTileAccessIteratorAnalytic( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() // threadblock offset - units are whole CTA tiles + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + filter_k_(0), + filter_r_(0), + filter_s_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_k_ = threadblock_offset.column() + thread_coord.contiguous(); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + int offset_nhw = threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + + offset_n_[s] = offset_nhw / (problem_size_.H * problem_size_.W); + int residual = offset_nhw % (problem_size_.H * problem_size_.W); + + offset_h_[s] = residual / problem_size_.W; + offset_w_[s] = residual % problem_size_.W; + } + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, layout); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + // move to the next tile + ++filter_s_; + if (filter_s_ < problem_size_.S) { + return; + } + filter_s_ = 0; + ++filter_r_; + if (filter_r_ < problem_size_.R) { + return; + } + filter_r_ = 0; + + filter_k_ += Shape_::kColumn * problem_size_.split_k_slices; + } + + /// Returns the coordinate in the output tensor Dy that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + int n = offset_n_[iteration_strided_]; + int h = offset_h_[iteration_strided_]; + int w = offset_w_[iteration_strided_]; + + int r = filter_r_; + int s = filter_s_; + + if (problem_size_.mode == Mode::kConvolution) { + r = (problem_size_.R - 1 - r); + s = (problem_size_.S - 1 - s); + } + + int p = (h + problem_size_.pad_h - r * problem_size_.dilation_h) / problem_size_.stride_h; + int q = (w + problem_size_.pad_w - s * problem_size_.dilation_w) / problem_size_.stride_w; + + int k = filter_k_ + iteration_vector_ * AccessType::kElements; + + return TensorCoord(n, p, q, k); + } + + /// Returns true if the current coordinate is within the output tensor Dy + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + return coord.n() < problem_size_.N && + coord.h() >= 0 && coord.h() < problem_size_.P && + coord.w() >= 0 && coord.w() < problem_size_.Q && + coord.c() < problem_size_.K; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dDgradOutputGradientTileAccessIteratorAnalytic &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // Conv2dDgradFilterTileAccessIteratorAnalytic unity stride specialization + // only supports (stride_h, stride_w) = (1, 1) + if (problem_size.stride() != MatrixCoord({1, 1})) { + return Status::kErrorNotSupported; + } + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.K % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + + diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_output_gradient_tile_access_iterator_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_output_gradient_tile_access_iterator_optimized.h new file mode 100644 index 0000000..6885be6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_dgrad_output_gradient_tile_access_iterator_optimized.h @@ -0,0 +1,821 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (output gradient tile) + matrix from memory. + + This iterator assumes TensorNHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + conv::StrideSupport StrideSupport_ = conv::StrideSupport::kUnity, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dDgradOutputGradientTileAccessIteratorOptimized; +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Conv2dDgradOutputGradientTileAccessIteratorOptimized strided dgrad needs special handling +// to skip MMAs (Dx = Dy * w) on invalid filter positions +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + typename AccessType_ +> +class Conv2dDgradOutputGradientTileAccessIteratorOptimized < + Shape_, + Element_, + ThreadMap_, + conv::StrideSupport::kStrided, + AccessType_ +> { +public: + + // + // Types + // + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNHWC; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + using Mask = uint64_t; + + static_assert(sizeof_bits::value >= 8, + "DGRAD requires elements of size 8b or greater."); + + // + // Simpligying assertions + // + + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv2dStridedDgradOutputGradientIteratorOptimizedParams; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + + // One pointer per access + char const *pointer_[ThreadMap::Iterations::kStrided]; + + int filter_k_; + int filter_r_; + int filter_s_; + int start_r_; + int start_s_; + int64_t reset_bytes_s_; + int64_t reset_bytes_r_; + + Index masks_[ThreadMap::Iterations::kStrided][kAccessesPerVector][2]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dDgradOutputGradientTileAccessIteratorOptimized( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + FastDivmod const &stride_h_divmod, FastDivmod const &stride_w_divmod, + int start_r, int start_s, + MatrixCoord const &threadblock_offset = MatrixCoord() // threadblock offset - units are whole CTA tiles + ): + params_(params), + problem_size_(problem_size), + filter_k_(0), + filter_r_(start_r), + filter_s_(start_s), + start_r_(start_r), + start_s_(start_s) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_k_ = threadblock_offset.column() + thread_coord.contiguous(); + + reset_bytes_s_ = (problem_size_.num_gemm_k_filter_s(start_s_) - 1) * params_.inc_next[0]; + + reset_bytes_r_ = (problem_size_.num_gemm_k_filter_s(start_s_) - 1) * params_.inc_next[0] + + (problem_size_.num_gemm_k_filter_r(start_r_) - 1) * params_.inc_next[1]; + + int offset_n[ThreadMap::Iterations::kStrided]; + int offset_p[ThreadMap::Iterations::kStrided]; + int offset_q[ThreadMap::Iterations::kStrided]; + + int filter_r = filter_r_; + int filter_s = filter_s_; + + if (problem_size_.mode == Mode::kConvolution) { + filter_r = (problem_size_.R - 1 - filter_r); + filter_s = (problem_size_.S - 1 - filter_s); + } + + // Starting h, w positions for filter position in gemm_k=0 + int start_h, start_w; + strided_dgrad_starting_coords( + problem_size_, + stride_h_divmod, stride_w_divmod, + filter_r, filter_s, + start_h, start_w); + + + // Effective starting P and Q for filter position required for remapping NHW rows + int P = (problem_size_.H - start_h + problem_size_.stride_h - 1) / problem_size_.stride_h; + int Q = (problem_size_.W - start_w + problem_size_.stride_w - 1) / problem_size_.stride_w; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + pointer_[s] = reinterpret_cast(ptr); + + int offset_npq = (threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided) % params_.tiled_rows_per_filter; + + // (STEP 1) [reorder NHW rows to start with same filter positions] + offset_n[s] = offset_npq / (P * Q); + int residual = offset_npq % (P * Q); + + int p = (residual / Q); + int q = (residual % Q); + + int mapped_h = (start_h + p * problem_size_.stride_h); + int mapped_w = (start_w + q * problem_size_.stride_w); + + // Access (p, q) coordinates for Dy tensor for filter position in gemm_k=0 + // note that (h + pad_h - filter_r) and (w + pad_w - filter_s) are ensured to be + // divisible by stride_h and stride_w + offset_p[s] = (mapped_h + problem_size_.pad_h - filter_r) / problem_size_.stride_h; + offset_q[s] = (mapped_w + problem_size_.pad_w - filter_s) / problem_size_.stride_w; + + // Initialize pointers for gemm_k=0 + TensorCoord coord{offset_n[s], offset_p[s], offset_q[s], filter_k_}; + + pointer_[s] += params_.layout(coord) * sizeof_bits::value / 8; + } + + // + // Precompute mask predicates + // + clear_mask(); + + MCTLASS_PRAGMA_NO_UNROLL + for (int r = start_r; r < problem_size_.R; r += problem_size_.stride_h) { + MCTLASS_PRAGMA_UNROLL + for (int s_idx = 0; s_idx < ThreadMap::Iterations::kStrided; ++s_idx) { + + int p = offset_p[s_idx] ; + + p += (params_.conv_sign * (r / problem_size_.stride_h)); + + bool pred = (offset_n[s_idx] < problem_size_.N && p >= 0 && p < problem_size_.P); + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + masks_[s_idx][v_idx][0] |= (pred << r); + } + } + } + + MCTLASS_PRAGMA_NO_UNROLL + for(int s = start_s; s < problem_size_.S; s += problem_size_.stride_w) { + MCTLASS_PRAGMA_UNROLL + for (int s_idx = 0; s_idx < ThreadMap::Iterations::kStrided; ++s_idx) { + + int q = offset_q[s_idx]; + q += (params_.conv_sign * (s / problem_size_.stride_w)); + + bool pred = (q >=0 && q < problem_size_.Q); + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + masks_[s_idx][v_idx][1] |= (pred << s); + } + } + } + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + clear_mask(v_idx, (filter_k_ + v_idx * AccessType::kElements) >= problem_size.K); + } + + set_iteration_index(0); + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}); + } + +private: + + /// Adds a pointer offset in units of element + MCTLASS_HOST_DEVICE + void add_byte_offset_(LongIndex byte_offset, LongIndex byte_reset = 0) { + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + pointer_[s] += byte_offset - byte_reset; + } + } + +public: + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + add_byte_offset_(pointer_offset * sizeof_bits::value / 8); + } + + MCTLASS_DEVICE + void advance() { + + int next_idx = 0; + int64_t reset_bytes = 0; + + // Move filter_s by stride_w + filter_s_ += problem_size_.stride_w; + if (filter_s_ >= problem_size_.S) { + + // Restore filter_s + filter_s_ = start_s_; + + // Move filter_r by stride_h + filter_r_ += problem_size_.stride_h; +#if 1 + if (filter_r_ < problem_size_.R) { + + next_idx = 1; + + // Restore bytes in q coordinate (Mma in filter s dimension) + reset_bytes = reset_bytes_s_; + + } else { + + // Restore filter_r + filter_r_ = start_r_; + + next_idx = 2; + + // Restore bytes in p and q coordinate (Mma in filter s and r dimension) + reset_bytes = reset_bytes_r_; + } +#else + asm volatile( + "{\n\t" + " .reg .pred %%p;\n\t" + " setp.lt.s32 %%p, %3, %4;\n\t" + " selp.s32 %0, %3, %5, %%p;\n\t" + " selp.s32 %1, 1, 2, %%p;\n\t" + " selp.s64 %2, %6, %7, %%p;\n\t" + "}\n" + : "=r"(filter_r_), "=r"(next_idx), "=l"(reset_bytes) + : "r"(filter_r_), "r"(problem_size_.R), "r"(start_r_), + "l"(reset_bytes_s_), "l"(reset_bytes_r_)); +#endif + } + + // offset pointers by offset_bytes + add_byte_offset_(params_.inc_next[next_idx] - reset_bytes); + + if (next_idx == 2) { + filter_k_ += params_.filter_k_delta; + } + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + clear_mask(v_idx, (filter_k_ + v_idx * AccessType::kElements) >= problem_size_.K); + } + } + + /// Clears the predicates + MCTLASS_HOST_DEVICE + void clear_mask(bool clear = true) { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + masks_[s][v][0] = clear ? Mask(0) : masks_[s][v][0]; + masks_[s][v][1] = clear ? Mask(0) : masks_[s][v][1]; + } + } + } + + /// Clears the predicates + MCTLASS_HOST_DEVICE + void clear_mask(int v, bool clear = true) { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + masks_[s][v][0] = clear ? Mask(0) : masks_[s][v][0]; + masks_[s][v][1] = clear ? Mask(0) : masks_[s][v][1]; + } + } + + /// Returns true if the current coordinate is within the output tensor Dy + MCTLASS_HOST_DEVICE + bool valid() const { + return + (masks_[iteration_strided_][iteration_vector_][0] & (Index(1) << filter_r_)) && + (masks_[iteration_strided_][iteration_vector_][1] & (Index(1) << filter_s_)); + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + return reinterpret_cast(pointer_[iteration_strided_]) + iteration_vector_; + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dDgradOutputGradientTileAccessIteratorOptimized &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.K % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + // Limit on filter size + if (problem_size.R > 32 || problem_size.S > 32) { + return Status::kErrorNotSupported; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Conv2dDgradOutputGradientTileAccessIteratorOptimized unity stride dgrad is optimized for dgrad +// with problem stride = {1x1} +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + typename AccessType_ +> +class Conv2dDgradOutputGradientTileAccessIteratorOptimized < + Shape_, + Element_, + ThreadMap_, + conv::StrideSupport::kUnity, + AccessType_ +> { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNHWC; + using TensorCoord = typename Layout::TensorCoord; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kUnity; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + using Mask = uint64_t; + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv2dDgradOutputGradientIteratorOptimizedParams; + +private: + + Conv2dDgradOutputGradientIteratorOptimizedParams const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + + // One pointer per access + char const *pointer_[ThreadMap::Iterations::kStrided]; + + // current filter position (r, s) + int filter_r_; + int filter_s_; + int filter_k_; + + Index masks_[ThreadMap::Iterations::kStrided][kAccessesPerVector][2]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dDgradOutputGradientTileAccessIteratorOptimized( + Conv2dDgradOutputGradientIteratorOptimizedParams const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() // tile index - units are threadblock-scoped tiles + ): + params_(params), + problem_size_(problem_size), + filter_k_(0), + filter_r_(0), + filter_s_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_k_ = threadblock_offset.column() + thread_coord.contiguous(); + + int offset_n[ThreadMap::Iterations::kStrided]; + int offset_h[ThreadMap::Iterations::kStrided]; + int offset_w[ThreadMap::Iterations::kStrided]; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + pointer_[s] = reinterpret_cast(ptr); + + int offset_nhw = threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + + // The subseqnet fast_divmod() operations are equivalent to the following logical computation: + // + // + // offset_n[s] = offset_nhw / (problem_size_.H * problem_size_.W); + // int residual = offset_nhw % (problem_size_.H * problem_size_.W); + // + // offset_h[s] = residual / problem_size_.W; + // offset_w[s] = residual % problem_size_.W; + // + + int residual; + + params_.hw_divmod(offset_n[s], residual, offset_nhw); + params_.w_divmod(offset_h[s], offset_w[s], residual); + + TensorCoord coord = at_(offset_n[s], offset_h[s], offset_w[s], 0, 0); + + pointer_[s] += params_.layout(coord) * sizeof_bits::value / 8; + } + + clear_mask(); + + MCTLASS_PRAGMA_NO_UNROLL + for (int r = 0; r < problem_size_.R; ++r) { + MCTLASS_PRAGMA_UNROLL + for (int s_idx = 0; s_idx < ThreadMap::Iterations::kStrided; ++s_idx) { + + int r_ = r; + if (problem_size_.mode == Mode::kConvolution) { + r_ = problem_size_.R - 1 - r; + } + + int p = offset_h[s_idx] + problem_size_.pad_h - r_ * problem_size_.dilation_h; + + bool pred = (offset_n[s_idx] < problem_size_.N && p >= 0 && p < problem_size_.P); + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + masks_[s_idx][v_idx][0] |= (pred << r); + } + } + } + + MCTLASS_PRAGMA_NO_UNROLL + for (int s = 0; s < problem_size_.S; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int s_idx = 0; s_idx < ThreadMap::Iterations::kStrided; ++s_idx) { + + int s_ = s; + if (problem_size_.mode == Mode::kConvolution) { + s_ = problem_size_.S - 1 - s; + } + + int q = offset_w[s_idx] + problem_size_.pad_w - s_ * problem_size_.dilation_w; + + bool pred = (q >= 0 && q < problem_size_.Q); + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + masks_[s_idx][v_idx][1] |= (pred << s); + } + } + } + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + clear_mask(v_idx, filter_k_ + v_idx * AccessType::kElements >= problem_size.K); + } + + set_iteration_index(0); + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}, + ThreadMap::kThreads, + ThreadMap::kElementsPerAccess, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided}); + } + +private: + + /// Returns the coordinate in the output gradient tensor dy that is correspoinding to + // activation nhw and filter position k, r, s + MCTLASS_HOST_DEVICE + TensorCoord at_(int n, int h, int w, int r, int s) const { + + if (problem_size_.mode == Mode::kConvolution) { + r = problem_size_.R - 1 - r; + s = problem_size_.S - 1 - s; + } + + int p = h + problem_size_.pad_h - r * problem_size_.dilation_h; + int q = w + problem_size_.pad_w - s * problem_size_.dilation_w; + + return TensorCoord(n, p, q, filter_k_); + } + + /// Adds a pointer offset in units of element + MCTLASS_HOST_DEVICE + void add_byte_offset_(LongIndex byte_offset) { + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + pointer_[s] += byte_offset; + } + } + +public: + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + add_byte_offset_(pointer_offset * sizeof_bits::value / 8); + } + + MCTLASS_HOST_DEVICE + void advance() { + + int next_idx = 0; + + // moves to the next tile + ++filter_s_; + if (filter_s_ == problem_size_.S) { + filter_s_ = 0; + ++filter_r_; + + if (filter_r_ < problem_size_.R) { + next_idx = 1; + } + else { + filter_r_ = 0; + next_idx = 2; + } + } + + add_byte_offset_(params_.inc_next[next_idx]); + + if (next_idx == 2) { + filter_k_ += params_.filter_k_delta; + } + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + clear_mask(v_idx, (filter_k_ + v_idx * AccessType::kElements) >= problem_size_.K); + } + } + + /// Clears the predicates + MCTLASS_HOST_DEVICE + void clear_mask(bool clear = true) { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + masks_[s][v][0] = clear ? Mask(0) : masks_[s][v][0]; + masks_[s][v][1] = clear ? Mask(0) : masks_[s][v][1]; + } + } + } + + /// Clears the predicates + MCTLASS_HOST_DEVICE + void clear_mask(int v, bool clear = true) { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + masks_[s][v][0] = clear ? Mask(0) : masks_[s][v][0]; + masks_[s][v][1] = clear ? Mask(0) : masks_[s][v][1]; + } + } + + MCTLASS_HOST_DEVICE + bool valid() { + + return + (masks_[iteration_strided_][iteration_vector_][0] & (Index(1) << filter_r_)) && + (masks_[iteration_strided_][iteration_vector_][1] & (Index(1) << filter_s_)); + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + return reinterpret_cast(pointer_[iteration_strided_]) + iteration_vector_; + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dDgradOutputGradientTileAccessIteratorOptimized &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // This is specialized for unit stride + if (problem_size.stride() != MatrixCoord({1, 1})) { + return Status::kErrorNotSupported; + } + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.K % AccessType::kElements) { + return Status::kErrorNotSupported; + } + + // Limit on filter size + if (problem_size.R > 32 || problem_size.S > 32) { + return Status::kErrorNotSupported; + } + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_analytic.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_analytic.h new file mode 100644 index 0000000..9428c99 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_analytic.h @@ -0,0 +1,332 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (activation tile) + matrix from memory. + + This iterator assumes TensorNHWC or TensorNCxHWx layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename Layout_, + typename ThreadMap_, + typename AccessType_ = mctlass::AlignedArray, + conv::GroupMode GroupMode_ = conv::GroupMode::kNone +> +class Conv2dFpropActivationTileAccessIteratorAnalytic { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using TensorCoord = typename Layout::TensorCoord; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + static conv::GroupMode const kGroupMode = GroupMode_; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv2dAnalyticParams; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + int filter_c_; + int filter_r_; + int filter_s_; + int filter_c_init_; + int group_idx_offset_; + int channels_per_group_; + int crs_cnt_; + int crs_per_group_; + + int offset_n_[ThreadMap::Iterations::kStrided]; + int offset_p_[ThreadMap::Iterations::kStrided]; + int offset_q_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dFpropActivationTileAccessIteratorAnalytic( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() // tile index - units are threadblock-scoped tiles + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + crs_cnt_(0), + group_idx_offset_(0), + filter_c_(0), + filter_r_(0), + filter_s_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_c_ = threadblock_offset.column() + thread_coord.contiguous(); + + if (kGroupMode != conv::GroupMode::kNone) { + filter_c_init_ = filter_c_; + channels_per_group_ = problem_size_.C / problem_size_.groups; + crs_per_group_ = problem_size_.S * problem_size_.R * ((channels_per_group_ + Shape::kColumn - 1) / Shape::kColumn); + } + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + int offset_npq = threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + + offset_n_[s] = offset_npq / (problem_size_.P * problem_size_.Q); + int residual = offset_npq % (problem_size_.P * problem_size_.Q); + + offset_p_[s] = residual / problem_size_.Q; + offset_q_[s] = residual % problem_size_.Q; + } + + set_iteration_index(0); + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, layout); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + // moves to the next tile + if (kGroupMode != conv::GroupMode::kNone) { + ++crs_cnt_; + } + + ++filter_s_; + if (filter_s_ < problem_size_.S) { + return; + } + filter_s_ = 0; + ++filter_r_; + if (filter_r_ < problem_size_.R) { + return; + } + filter_r_ = 0; + + if (kGroupMode == conv::GroupMode::kNone) { + filter_c_ += Shape::kColumn * problem_size_.split_k_slices; + } else { + if (crs_cnt_ == crs_per_group_) { + // moves to next group + crs_cnt_ = 0; + ++group_idx_offset_; + filter_c_ = group_idx_offset_ * channels_per_group_ + filter_c_init_; + } else { + filter_c_ += Shape::kColumn * problem_size_.split_k_slices; + } + } + } + + /// Returns the coordinate in the activations tensor X that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + int n = offset_n_[iteration_strided_]; + int p = offset_p_[iteration_strided_]; + int q = offset_q_[iteration_strided_]; + + int r = filter_r_; + int s = filter_s_; + + if (problem_size_.mode == Mode::kConvolution) { + r = (problem_size_.R - 1 - filter_r_); + s = (problem_size_.S - 1 - filter_s_); + } + + int h = p * problem_size_.stride_h - problem_size_.pad_h + r * problem_size_.dilation_h; + int w = q * problem_size_.stride_w - problem_size_.pad_w + s * problem_size_.dilation_w; + + int c = filter_c_ + iteration_vector_ * AccessType::kElements; + + return TensorCoord(n, h, w, c); + } + + /// Returns true if the current coordinate is within the activations tensor X + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + return coord.n() < problem_size_.N && + coord.h() >= 0 && coord.h() < problem_size_.H && + coord.w() >= 0 && coord.w() < problem_size_.W && + coord.c() < problem_size_.C; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + AccessType const *ptr = reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + + return ptr; + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dFpropActivationTileAccessIteratorAnalytic &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if ((problem_size.C / problem_size.groups) % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + if (platform::is_same>::value) { + if (problem_size.C % 32) { + return Status::kErrorInvalidProblem; + } + } + + if (platform::is_same>::value) { + if (problem_size.C % 64) { + return Status::kErrorInvalidProblem; + } + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_few_channels.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_few_channels.h new file mode 100644 index 0000000..32bfea7 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_few_channels.h @@ -0,0 +1,360 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (activation tile) + matrix from memory. + + This iterator assumes TensorNHWC or TensorNCxHWx layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename Layout_, + typename ThreadMap_, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dFpropActivationTileAccessIteratorFewChannels { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using TensorCoord = typename Layout::TensorCoord; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kFewChannels; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kPositionsPerTile = Shape::kColumn; + + static int const kAccessesPerVector = kElementsPerAccess / AccessType::kElements; + + static bool const kUseFastDivmodPrologue = true; + static bool const kUseFastDivmodMainloop = true; + + static int const kStrideH = 0; + static int const kStrideW = 0; + static int const kDilationH = 0; + static int const kDilationW = 0; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv2dFewChannelsParams; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + int rsc_index_; + int offset_n_[ThreadMap::Iterations::kStrided]; + int offset_p_[ThreadMap::Iterations::kStrided]; + int offset_q_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dFpropActivationTileAccessIteratorFewChannels( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() // tile index - units are threadblock-scoped tiles + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + rsc_index_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + rsc_index_ = (threadblock_offset.column() + thread_coord.contiguous()); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + int offset_npq = threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + + if (kUseFastDivmodPrologue) { + int residual = params_.divmod_Q.divmod(offset_q_[s], offset_npq); + offset_n_[s] = params_.divmod_P.divmod(offset_p_[s], residual); + } + else { + offset_n_[s] = offset_npq / (problem_size_.P * problem_size_.Q); + int residual = offset_npq % (problem_size_.P * problem_size_.Q); + + offset_p_[s] = residual / problem_size_.Q; + offset_q_[s] = residual % problem_size_.Q; + } + } + + set_iteration_index(0); + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, layout); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + + rsc_index_ += kPositionsPerTile * problem_size_.split_k_slices; + } + + /// Returns the coordinate in the activations tensor X that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + int n = offset_n_[iteration_strided_]; + int p = offset_p_[iteration_strided_]; + int q = offset_q_[iteration_strided_]; + + int rsc_index = rsc_index_ + iteration_vector_ * AccessType::kElements; + + int r = 0; + int s = 0; + int c = 0; + + if (kUseFastDivmodMainloop) { + int rs_index = params_.divmod_C.divmod(c, rsc_index); + r = params_.divmod_S.divmod(s, rs_index); + } + else { + c = (rsc_index % problem_size_.C); + + int rs_index = (rsc_index / problem_size_.C); + s = (rs_index % problem_size_.S); + r = (rs_index / problem_size_.S); + } + + if (problem_size_.mode == Mode::kConvolution) { + r = (problem_size_.R - 1 - r); + s = (problem_size_.S - 1 - s); + } + + int stride_h = kStrideH; + if (!kStrideH) { + stride_h = problem_size_.stride_h; + } + + int stride_w = kStrideW; + if (!kStrideW) { + stride_w = problem_size_.stride_w; + } + + int dilation_h = kDilationH; + if (!kDilationH) { + dilation_h = problem_size_.dilation_h; + } + + int dilation_w = kDilationW; + if (!kDilationW) { + dilation_w = problem_size_.dilation_w; + } + + int h = p * stride_h - problem_size_.pad_h + r * dilation_h; + int w = q * stride_w - problem_size_.pad_w + s * dilation_w; + + return TensorCoord(n, h, w, c); + } + + /// Returns true if the current coordinate is within the activations tensor X + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + bool in_bounds = + coord.n() < problem_size_.N && + coord.h() >= 0 && coord.h() < problem_size_.H && + coord.w() >= 0 && coord.w() < problem_size_.W && + coord.c() < problem_size_.C; + + return in_bounds; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + + int32_t offset = + coord.n() * params_.stride_n + + coord.h() * params_.stride_h + + coord.w() * params_.stride_w + + coord.c(); + + AccessType const *ptr = reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + + return ptr; + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dFpropActivationTileAccessIteratorFewChannels &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + if (kDilationH && problem_size.dilation_h != kDilationH) { + return Status::kErrorInvalidProblem; + } + + if (kDilationW && problem_size.dilation_w != kDilationW) { + return Status::kErrorInvalidProblem; + } + + if (kStrideH && problem_size.stride_h != kStrideH) { + return Status::kErrorInvalidProblem; + } + + if (kStrideW && problem_size.stride_w != kStrideW) { + return Status::kErrorInvalidProblem; + } + + if (platform::is_same>::value) { + if (problem_size.C % 32) { + return Status::kErrorInvalidProblem; + } + } + + if (platform::is_same>::value) { + if (problem_size.C % 64) { + return Status::kErrorInvalidProblem; + } + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_fixed_channels.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_fixed_channels.h new file mode 100644 index 0000000..b2641ab --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_fixed_channels.h @@ -0,0 +1,353 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (activation tile) + matrix from memory. + + This iterator assumes TensorNHWC or TensorNCxHWx layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename Layout_, + typename ThreadMap_, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dFpropActivationTileAccessIteratorFixedChannels { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using TensorCoord = typename Layout::TensorCoord; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kFixedChannels; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kFilterPositionsPerTile = Shape::kColumn / AccessType::kElements; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static bool const kUseFastDivmodPrologue = true; + static bool const kUseFastDivmodMainloop = true; + + static int const kStrideH = 0; + static int const kStrideW = 0; + static int const kDilationH = 0; + static int const kDilationW = 0; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv2dFewChannelsParams; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + int rs_index_; + int offset_n_[ThreadMap::Iterations::kStrided]; + int offset_p_[ThreadMap::Iterations::kStrided]; + int offset_q_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dFpropActivationTileAccessIteratorFixedChannels( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() // tile index - units are threadblock-scoped tiles + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + rs_index_(0) { + + // + // This requires problem_size.C == AccessType::kElements + // + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + rs_index_ = (threadblock_offset.column() + thread_coord.contiguous()) / AccessType::kElements; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + int offset_npq = threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + + if (kUseFastDivmodPrologue) { + int residual = params_.divmod_Q.divmod(offset_q_[s], offset_npq); + offset_n_[s] = params_.divmod_P.divmod(offset_p_[s], residual); + } + else { + offset_n_[s] = offset_npq / (problem_size_.P * problem_size_.Q); + int residual = offset_npq % (problem_size_.P * problem_size_.Q); + + offset_p_[s] = residual / problem_size_.Q; + offset_q_[s] = residual % problem_size_.Q; + } + } + + set_iteration_index(0); + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, layout); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + + rs_index_ += kFilterPositionsPerTile * problem_size_.split_k_slices; + } + + /// Returns the coordinate in the activations tensor X that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + int n = offset_n_[iteration_strided_]; + int p = offset_p_[iteration_strided_]; + int q = offset_q_[iteration_strided_]; + + int rs_index = rs_index_ + iteration_vector_; + + int r = 0; + int s = 0; + + if (kUseFastDivmodMainloop) { + r = params_.divmod_S.divmod(s, rs_index); + } + else { + s = (rs_index % problem_size_.S); + r = (rs_index / problem_size_.S); + } + + if (problem_size_.mode == Mode::kConvolution) { + r = (problem_size_.R - 1 - r); + s = (problem_size_.S - 1 - s); + } + + int stride_h = kStrideH; + if (!kStrideH) { + stride_h = problem_size_.stride_h; + } + + int stride_w = kStrideW; + if (!kStrideW) { + stride_w = problem_size_.stride_w; + } + + int dilation_h = kDilationH; + if (!kDilationH) { + dilation_h = problem_size_.dilation_h; + } + + int dilation_w = kDilationW; + if (!kDilationW) { + dilation_w = problem_size_.dilation_w; + } + + int h = p * stride_h - problem_size_.pad_h + r * dilation_h; + int w = q * stride_w - problem_size_.pad_w + s * dilation_w; + + return TensorCoord(n, h, w, 0); + } + + /// Returns true if the current coordinate is within the activations tensor X + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + return coord.n() < problem_size_.N && + coord.h() >= 0 && coord.h() < problem_size_.H && + coord.w() >= 0 && coord.w() < problem_size_.W; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + + int32_t offset = + coord.n() * params_.stride_n + + coord.h() * params_.stride_h + + coord.w() * params_.stride_w + coord.c(); + + AccessType const *ptr = reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + + return ptr; + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dFpropActivationTileAccessIteratorFixedChannels &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C != AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + if (kDilationH && problem_size.dilation_h != kDilationH) { + return Status::kErrorInvalidProblem; + } + + if (kDilationW && problem_size.dilation_w != kDilationW) { + return Status::kErrorInvalidProblem; + } + + if (kStrideH && problem_size.stride_h != kStrideH) { + return Status::kErrorInvalidProblem; + } + + if (kStrideW && problem_size.stride_w != kStrideW) { + return Status::kErrorInvalidProblem; + } + + if (platform::is_same>::value) { + if (problem_size.C % 32) { + return Status::kErrorInvalidProblem; + } + } + + if (platform::is_same>::value) { + if (problem_size.C % 64) { + return Status::kErrorInvalidProblem; + } + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_optimized.h new file mode 100644 index 0000000..aacad66 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_activation_tile_access_iterator_optimized.h @@ -0,0 +1,422 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (activation tile) + matrix from memory. + + This iterator assumes TensorNHWC or TensorNCxHWx layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename Layout_, + typename ThreadMap_, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dFpropActivationTileAccessIteratorOptimized { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using TensorCoord = typename Layout::TensorCoord; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + using Mask = uint64_t; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv2dFpropActivationIteratorOptimizedParams; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + + // One pointer per access + char const *pointer_[ThreadMap::Iterations::kStrided]; + + // current filter position (r, s) + int filter_r_; + int filter_s_; + int filter_c_; + + Index masks_[ThreadMap::Iterations::kStrided][kAccessesPerVector][2]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dFpropActivationTileAccessIteratorOptimized( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() // tile index - units are threadblock-scoped tiles + ): + params_(params), + problem_size_(problem_size), + filter_c_(0), + filter_r_(0), + filter_s_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_c_ = threadblock_offset.column() + thread_coord.contiguous(); + + int offset_n[ThreadMap::Iterations::kStrided]; + int offset_p[ThreadMap::Iterations::kStrided]; + int offset_q[ThreadMap::Iterations::kStrided]; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + pointer_[s] = reinterpret_cast(ptr); + + int offset_npq = threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + + // The subseqnet fast_divmod() operations are equivalent to the following logical computation: + // + // + // offset_n[s] = offset_npq / (problem_size_.P * problem_size_.Q); + // int residual = offset_npq % (problem_size_.P * problem_size_.Q); + // + // offset_p[s] = residual / problem_size_.Q; + // offset_q[s] = residual % problem_size_.Q; + // + + int residual; + + params.pq_divmod(offset_n[s], residual, offset_npq); + params.q_divmod(offset_p[s], offset_q[s], residual); + + TensorCoord coord = at_(offset_n[s], offset_p[s], offset_q[s], 0, 0); + + pointer_[s] += params_.layout(coord) * sizeof_bits::value / 8; + } + + clear_mask(); + + MCTLASS_PRAGMA_NO_UNROLL + for (int r = 0; r < problem_size_.R; ++r) { + MCTLASS_PRAGMA_UNROLL + for (int s_idx = 0; s_idx < ThreadMap::Iterations::kStrided; ++s_idx) { + + int r_ = r; + if (problem_size_.mode == Mode::kConvolution) { + r_ = problem_size_.R - 1 - r; + } + + int h = offset_p[s_idx] * problem_size_.stride_h - problem_size_.pad_h + r_ * problem_size_.dilation_h; + + bool pred = (offset_n[s_idx] < problem_size_.N && h >= 0 && h < problem_size_.H); + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + masks_[s_idx][v_idx][0] |= (pred << r); + } + } + } + + MCTLASS_PRAGMA_NO_UNROLL + for (int s = 0; s < problem_size_.S; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int s_idx = 0; s_idx < ThreadMap::Iterations::kStrided; ++s_idx) { + + int s_ = s; + if (problem_size_.mode == Mode::kConvolution) { + s_ = problem_size_.S - 1 - s; + } + + int w = offset_q[s_idx] * problem_size_.stride_w - problem_size_.pad_w + s_ * problem_size_.dilation_w; + + bool pred = (w >= 0 && w < problem_size_.W); + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + masks_[s_idx][v_idx][1] |= (pred << s); + } + } + } + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + clear_mask(v_idx, filter_c_ + v_idx * AccessType::kElements >= problem_size_.C); + } + + set_iteration_index(0); + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}, + ThreadMap::kThreads, + ThreadMap::kElementsPerAccess, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided}); + } + +private: + + /// Returns the coordinate in the activations tensor X that is correspoinding to + // output npq and filter position r, s + MCTLASS_HOST_DEVICE + TensorCoord at_(int n, int p, int q, int r, int s) const { + + if (problem_size_.mode == Mode::kConvolution) { + r = problem_size_.R - 1 - r; + s = problem_size_.S - 1 - s; + } + + int h = p * problem_size_.stride_h - problem_size_.pad_h + r * problem_size_.dilation_h; + int w = q * problem_size_.stride_w - problem_size_.pad_w + s * problem_size_.dilation_w; + + return TensorCoord(n, h, w, filter_c_); + } + + /// Adds a pointer offset in units of element + MCTLASS_HOST_DEVICE + void add_byte_offset_(LongIndex byte_offset) { + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + pointer_[s] += byte_offset; + } + } + +public: + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + add_byte_offset_(pointer_offset * sizeof_bits::value / 8); + } + + MCTLASS_HOST_DEVICE + void advance() { + + int next_idx = 0; + + // moves to the next tile + ++filter_s_; + if (filter_s_ == problem_size_.S) { + filter_s_ = 0; + ++filter_r_; + + if (filter_r_ < problem_size_.R) { + next_idx = 1; + } + else { + filter_r_ = 0; + next_idx = 2; + } + } + + add_byte_offset_(params_.inc_next[next_idx]); + + if (next_idx == 2) { + filter_c_ += params_.filter_c_delta; + } + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + clear_mask(v_idx, filter_c_ + v_idx * AccessType::kElements >= problem_size_.C); + } + } + + /// Clears the predicates + MCTLASS_HOST_DEVICE + void clear_mask(bool clear = true) { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + masks_[s][v][0] = clear ? 0 : masks_[s][v][0]; + masks_[s][v][1] = clear ? 0 : masks_[s][v][1]; + } + } + } + + /// Clears the predicates + MCTLASS_HOST_DEVICE + void clear_mask(int v, bool clear = true) { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + masks_[s][v][0] = clear ? 0 : masks_[s][v][0]; + masks_[s][v][1] = clear ? 0 : masks_[s][v][1]; + } + } + + MCTLASS_HOST_DEVICE + bool valid() { + + return + (masks_[iteration_strided_][iteration_vector_][0] & (Index(1) << filter_r_)) && + (masks_[iteration_strided_][iteration_vector_][1] & (Index(1) << filter_s_)); + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + return reinterpret_cast(pointer_[iteration_strided_]) + iteration_vector_; + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dFpropActivationTileAccessIteratorOptimized &operator++() { + + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if ((problem_size.C / problem_size.groups) % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + if (platform::is_same>::value) { + if (problem_size.C % 32) { + return Status::kErrorInvalidProblem; + } + } + + if (platform::is_same>::value) { + if (problem_size.C % 64) { + return Status::kErrorInvalidProblem; + } + } + + // Conv2dFpropActivationTileAccessIteratorOptimized has constraint on filter positions + // due to the number of mask bits. + if (problem_size.R > 32 || problem_size.S > 32) { + return Status::kErrorNotSupported; + } + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h new file mode 100644 index 0000000..75a45ff --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h @@ -0,0 +1,319 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (filter tile) + matrix from memory. + + This iterator assumes TensorNHWC or TensorCxRSKx layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename Layout_, + typename ThreadMap_, + typename AccessType_ = mctlass::AlignedArray, + conv::GroupMode GroupMode_ = conv::GroupMode::kNone +> +class Conv2dFpropFilterTileAccessIteratorAnalytic { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + static conv::GroupMode const kGroupMode = GroupMode_; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv2dAnalyticParams; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + int filter_r_; + int filter_s_; + int filter_c_; + int filter_c_init_; + int crs_cnt_; + int crs_per_group_; + int group_idx_offset_c_; + int channels_per_group_; + + int offset_k_[ThreadMap::Iterations::kStrided]; + int group_idx_offset_k_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dFpropFilterTileAccessIteratorAnalytic( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + crs_cnt_(0), + group_idx_offset_c_(0), + filter_r_(0), + filter_s_(0), + filter_c_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_c_ = threadblock_offset.row() + thread_coord.contiguous(); + + if (kGroupMode != conv::GroupMode::kNone) { + filter_c_init_ = filter_c_; + if (kGroupMode == conv::GroupMode::kDepthwise){ + channels_per_group_ = 1; + crs_per_group_ = problem_size_.S * problem_size_.R; + } else { + channels_per_group_ = problem_size_.C / problem_size_.groups; + crs_per_group_ = problem_size_.S * problem_size_.R * ((channels_per_group_ + Shape::kRow - 1) / Shape::kRow); + } + } + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_k_[s] = threadblock_offset.column() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + if (kGroupMode != conv::GroupMode::kNone && kGroupMode != conv::GroupMode::kDepthwise) { + group_idx_offset_k_[s] = (thread_coord.strided() + s * ThreadMap::Delta::kStrided) / (problem_size_.K / problem_size_.groups); + } + } + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * 8 / sizeof_bits::value; + } + + MCTLASS_HOST_DEVICE + void advance() { + // moves to the next tile + if (kGroupMode != conv::GroupMode::kNone) { + ++crs_cnt_; + } + + ++filter_s_; + if (filter_s_ < problem_size_.S) { + return; + } + filter_s_ = 0; + + ++filter_r_; + if (filter_r_ < problem_size_.R) { + return; + } + filter_r_ = 0; + + if (kGroupMode == conv::GroupMode::kNone) { + filter_c_ += Shape::kRow * problem_size_.split_k_slices; + } else { + if (crs_cnt_ == crs_per_group_) { + crs_cnt_ = 0; + filter_c_ = filter_c_init_; + if (kGroupMode != conv::GroupMode::kDepthwise) { + // moves to next group + ++group_idx_offset_c_; + } + } else { + filter_c_ += Shape::kRow * problem_size_.split_k_slices; + } + } + } + + /// Returns the coordinate in the filter tensor W that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + int k = offset_k_[iteration_strided_]; + int c = filter_c_ + iteration_vector_ * AccessType::kElements; + + return TensorCoord(k, filter_r_, filter_s_, c); + } + + /// Returns true if the current coordinate is within the activations tensor W + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + if (kGroupMode == conv::GroupMode::kNone) { + return coord.n() < problem_size_.K && coord.c() < problem_size_.C; + } else if (kGroupMode == conv::GroupMode::kDepthwise) { + return coord.n() < problem_size_.K && coord.c() < 1; // channels_per_group_ is always equal to ONE. + } else { + return coord.n() < problem_size_.K && coord.c() < channels_per_group_ && + group_idx_offset_c_ == group_idx_offset_k_[iteration_strided_]; + } + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dFpropFilterTileAccessIteratorAnalytic &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if ((problem_size.C / problem_size.groups) % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + if (platform::is_same>::value) { + if (problem_size.K % 32) { + return Status::kErrorInvalidProblem; + } + } + + if (platform::is_same>::value) { + if (problem_size.K % 64) { + return Status::kErrorInvalidProblem; + } + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_few_channels.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_few_channels.h new file mode 100644 index 0000000..4a1970c --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_few_channels.h @@ -0,0 +1,289 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (filter tile) + matrix from memory. + + This iterator assumes TensorNHWC or TensorCxRSKx layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename Layout_, + typename ThreadMap_, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dFpropFilterTileAccessIteratorFewChannels { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kFewChannels; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kPositionsPerTile = Shape::kRow; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static bool const kUseFastDivmodPrologue = true; + static bool const kUseFastDivmodMainloop = true; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv2dFewChannelsParams; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + int rsc_index_; + + int offset_k_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dFpropFilterTileAccessIteratorFewChannels( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + rsc_index_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + rsc_index_ = (threadblock_offset.row() + thread_coord.contiguous()); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_k_[s] = threadblock_offset.column() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + } + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * 8 / sizeof_bits::value; + } + + MCTLASS_HOST_DEVICE + void advance() { + // moves to the next tile + rsc_index_ += kPositionsPerTile * problem_size_.split_k_slices; + } + + /// Returns the coordinate in the filter tensor W that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + int rsc_index = rsc_index_ + iteration_vector_ * AccessType::kElements; + + int c = 0; + int s = 0; + int r = 0; + + if (kUseFastDivmodMainloop) { + int rs_index = params_.divmod_C.divmod(c, rsc_index); + r = params_.divmod_S.divmod(s, rs_index); + } + else { + c = (rsc_index % problem_size_.C); + int rs_index = (rsc_index / problem_size_.C); + + s = (rs_index % problem_size_.S); + r = (rs_index / problem_size_.S); + } + + int k = offset_k_[iteration_strided_]; + + return TensorCoord(k, r, s, c); + } + + /// Returns true if the current coordinate is within the activations tensor W + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + bool in_bounds = + coord.n() < problem_size_.K && + coord.h() >= 0 && + coord.h() < problem_size_.R && + coord.c() < problem_size_.C; + + return in_bounds; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + + int32_t offset = + coord.n() * params_.stride_n + + coord.h() * params_.stride_h + + coord.w() * params_.stride_w + + coord.c(); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dFpropFilterTileAccessIteratorFewChannels &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + if (platform::is_same>::value) { + if (problem_size.K % 32) { + return Status::kErrorInvalidProblem; + } + } + + if (platform::is_same>::value) { + if (problem_size.K % 64) { + return Status::kErrorInvalidProblem; + } + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_fixed_channels.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_fixed_channels.h new file mode 100644 index 0000000..8d95856 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_fixed_channels.h @@ -0,0 +1,275 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (filter tile) + matrix from memory. + + This iterator assumes TensorNHWC or TensorCxRSKx layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename Layout_, + typename ThreadMap_, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dFpropFilterTileAccessIteratorFixedChannels { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kFixedChannels; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kFilterPositionsPerTile = Shape::kRow / AccessType::kElements; + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static bool const kUseFastDivmodPrologue = true; + static bool const kUseFastDivmodMainloop = true; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv2dFewChannelsParams; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + int rs_index_; + + int offset_k_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dFpropFilterTileAccessIteratorFixedChannels( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + rs_index_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + rs_index_ = (threadblock_offset.row() + thread_coord.contiguous()) / AccessType::kElements; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_k_[s] = threadblock_offset.column() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + } + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * 8 / sizeof_bits::value; + } + + MCTLASS_HOST_DEVICE + void advance() { + // moves to the next tile + rs_index_ += kFilterPositionsPerTile * problem_size_.split_k_slices; + } + + /// Returns the coordinate in the filter tensor W that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + int rs_index = rs_index_ + iteration_vector_; + + int r = 0; + int s = 0; + + if (kUseFastDivmodMainloop) { + r = params_.divmod_S.divmod(s, rs_index); + } + else { + s = (rs_index % problem_size_.S); + r = (rs_index / problem_size_.S); + } + + int k = offset_k_[iteration_strided_]; + + return TensorCoord(k, r, s, 0); + } + + /// Returns true if the current coordinate is within the activations tensor W + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + return coord.n() < problem_size_.K && coord.h() >= 0 && coord.h() < problem_size_.R; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + + int32_t offset = + coord.n() * params_.stride_n + + coord.h() * params_.stride_h + + coord.w() * params_.stride_w + coord.c(); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dFpropFilterTileAccessIteratorFixedChannels &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C != AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + if (platform::is_same>::value) { + if (problem_size.K % 32) { + return Status::kErrorInvalidProblem; + } + } + + if (platform::is_same>::value) { + if (problem_size.K % 64) { + return Status::kErrorInvalidProblem; + } + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_optimized.h new file mode 100644 index 0000000..bf2094d --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_optimized.h @@ -0,0 +1,317 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (filter tile) + matrix from memory. + + This iterator assumes TensorNHWC or TensorCxRSKx layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" + +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename Layout_, + typename ThreadMap_, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dFpropFilterTileAccessIteratorOptimized{ +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + struct Params : Conv2dFpropFilterIteratorOptimizedParams { + + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Conv2dFpropFilterIteratorOptimizedParams const &base): + Conv2dFpropFilterIteratorOptimizedParams(base) { } + + MCTLASS_HOST_DEVICE + Params( + Conv2dProblemSize const &problem_size, + Layout const &layout + ): + Conv2dFpropFilterIteratorOptimizedParams( + problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}, + ThreadMap::kThreads, + ThreadMap::kElementsPerAccess, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided} + ) { + + } + }; + +private: + + Conv2dFpropFilterIteratorOptimizedParams const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + uint32_t predicates_[kAccessesPerVector]; + int filter_rs_; + int filter_c_; + int channels_per_group_; + + // + // Assertions + // + + // We map predicates into bits packed in this uint32_t container + static_assert(ThreadMap::Iterations::kStrided < sizeof(predicates_) * 8, + "Currently, the number of loads per iteration is limited by the size of the predicates container."); + +public: + + MCTLASS_HOST_DEVICE + Conv2dFpropFilterTileAccessIteratorOptimized( + Conv2dFpropFilterIteratorOptimizedParams const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + predicates_{0}, + filter_rs_(0), + filter_c_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_c_ = threadblock_offset.row() + thread_coord.contiguous(); + Index column = threadblock_offset.column() + thread_coord.strided(); + channels_per_group_ = problem_size_.C / problem_size_.groups; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + uint32_t pred = ((column + s * ThreadMap::Delta::kStrided < problem_size_.K) ? 1u : 0); + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + predicates_[v_idx] |= (pred << s); + } + } + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + clear_mask(v_idx, filter_c_ + v_idx * AccessType::kElements >= channels_per_group_); + } + + pointer_ += ( + params_.layout({filter_c_, column}) + ) * sizeof_bits::value / 8; + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + + LongIndex next = params_.inc_next_rs; + + // moves to the next tile + ++filter_rs_; + if (filter_rs_ == params_.RS) { + + filter_rs_ = 0; + next = params_.inc_next_c; + filter_c_ += params_.filter_c_delta; + } + + MCTLASS_PRAGMA_UNROLL + for (int v_idx = 0; v_idx < kAccessesPerVector; ++v_idx) { + clear_mask(v_idx, filter_c_ + v_idx * AccessType::kElements >= channels_per_group_); + } + + pointer_ += next; + } + + /// Clears the predicates + MCTLASS_HOST_DEVICE + void clear_mask(int v, bool clear = true) { + predicates_[v] = clear ? 0u : predicates_[v]; + } + + /// Returns true if the current coordinate is within the filter tensor W + MCTLASS_HOST_DEVICE + bool valid() { + return (predicates_[iteration_vector_] & (1u << iteration_strided_)); + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + return reinterpret_cast(pointer_) + iteration_vector_; + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dFpropFilterTileAccessIteratorOptimized &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + + // Move to the next K coordinate within the tile + pointer_ += params_.inc_next_k; + + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if ((problem_size.C / problem_size.groups) % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + if (platform::is_same>::value) { + if (problem_size.K % 32) { + return Status::kErrorInvalidProblem; + } + } + + if (platform::is_same>::value) { + if (problem_size.K % 64) { + return Status::kErrorInvalidProblem; + } + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_params.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_params.h new file mode 100644 index 0000000..cda3a2c --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_params.h @@ -0,0 +1,893 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Extracts the host-params objects into non-template code. +*/ + +#pragma once + +#define TRACE_CONV_PARAMS_INITIALIZERS_ENABLED 0 + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" + +#if TRACE_CONV_PARAMS_INITIALIZERS_ENABLED +#include +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Params structure used for all Conv2d analytic tile iterators +template< typename Layout_ = layout::TensorNHWC > +struct Conv2dAnalyticParams { + + using Layout = Layout_; + + Layout layout; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Conv2dAnalyticParams() { } + + MCTLASS_HOST_DEVICE + Conv2dAnalyticParams( + Conv2dProblemSize const &, // unused; placeholder to match other Params interfaces. + Layout const &layout + ): layout(layout) { + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Params structure used for all Conv2d analytic tile iterators +template< typename Layout_ = layout::TensorNHWC > +struct Conv2dFewChannelsParams { + + using Layout = Layout_; + + + int32_t stride_w; + int32_t stride_h; + int32_t stride_n; + + FastDivmod divmod_P; + FastDivmod divmod_Q; + FastDivmod divmod_S; + FastDivmod divmod_C; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Conv2dFewChannelsParams() { } + + MCTLASS_HOST_DEVICE + Conv2dFewChannelsParams( + Conv2dProblemSize const &problem_size, // unused; placeholder to match other Params interfaces. + Layout const &layout + ): + stride_w(int32_t(layout.stride()[0])), + stride_h(int32_t(layout.stride()[1])), + stride_n(int32_t(layout.stride()[2])), + divmod_P(problem_size.P), + divmod_Q(problem_size.Q), + divmod_S(problem_size.S), + divmod_C(problem_size.C) + { + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Parameters structure used for Conv2dDgradOutputGradientTileAccessIteratorAnalyticParams +struct Conv2dDgradOutputGradientTileAccessIteratorAnalyticParams { + + using Layout = layout::TensorNHWC; + + Layout layout; + int tiled_rows_per_filter; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Conv2dDgradOutputGradientTileAccessIteratorAnalyticParams() { } + + MCTLASS_HOST_DEVICE + Conv2dDgradOutputGradientTileAccessIteratorAnalyticParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, ///< layout object + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape + ): layout(layout) { + + int tile_m_per_filter = strided_dgrad_tile_m_per_filter(problem_size, threadblock_shape.row()); + + tiled_rows_per_filter = tile_m_per_filter * threadblock_shape.row(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#if TRACE_CONV_PARAMS_INITIALIZERS_ENABLED + +MCTLASS_HOST_DEVICE +void TraceIteratorParams( + char const *conv_operator, + char const *operand, + int element_size_bits, + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta +) { + +#if !defined(__MACA_ARCH__) + + char const *fname = "conv_iterator_params.csv"; + + std::ifstream test(fname); + bool file_exists = test.is_open(); + + if (file_exists) { + test.close(); + } + + std::ofstream trace("conv_iterator_params.csv", std::ofstream::app); + + if (!file_exists) { + trace + << "Operator,Operand,ElementSize,CtaRows,CtaColumns,ThreadCount,AccessSize," + << "IterationsContiguous,IterationsStrided,DeltaContiguous,DeltaStrided\n"; + } + + trace << conv_operator << "," << operand << "," << element_size_bits << "," + << threadblock_shape.row() << "," << threadblock_shape.column() + << "," << thread_count << "," << access_size + << "," << threadmap_iterations.contiguous() << "," << threadmap_iterations.strided() + << "," << threadmap_delta.contiguous() << "," << threadmap_delta.strided() << "\n"; +#endif +} + +#define TRACE_CONV_INITIALIZERS(conv_op, operand, element_size, cta_shape, thread_count, access_size, iterations, delta) \ + TraceIteratorParams(conv_op, operand, element_size, cta_shape, thread_count, access_size, iterations, delta); + +#else + +#define TRACE_CONV_INITIALIZERS(conv_op, operand, element_size, cta_shape, thread_count, access_size, iterations, delta) {} + +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Parameters structure used for Conv2dFpropActivationTileIteratorOptimized +template< typename Layout_ = layout::TensorNHWC > +struct Conv2dFpropActivationIteratorOptimizedParams; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Parameters structure used for Conv2dFpropActivationTileIteratorOptimized +template<> +struct Conv2dFpropActivationIteratorOptimizedParams { + + using Layout = layout::TensorNHWC; + + Layout layout; + + int64_t inc_next[3]; // {next S, next R, next C} + int filter_c_delta; // number of logical elements to add to filter_c_ + int PQ; // product of P*Q + + FastDivmod pq_divmod; + FastDivmod q_divmod; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Conv2dFpropActivationIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv2dFpropActivationIteratorOptimizedParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, ///< layout object + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): + layout(layout), + PQ(problem_size.P * problem_size.Q), + pq_divmod(PQ), + q_divmod(problem_size.Q) { + + TRACE_CONV_INITIALIZERS("conv2d_fprop", "activation", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + int conv_sign = (problem_size.mode == Mode::kConvolution ? -1 : 1); + + // next S + inc_next[0] = conv_sign * ( + int64_t(layout.stride()[0]) * problem_size.dilation_w + ) * element_size_bits / 8; + + // next R + inc_next[1] = conv_sign * ( + int64_t(layout.stride()[1]) * problem_size.dilation_h + - (problem_size.S - 1) * layout.stride()[0] * problem_size.dilation_w + ) * element_size_bits / 8; + + // next C + inc_next[2] = ( + threadblock_shape.column() * problem_size.split_k_slices + - conv_sign * int64_t(problem_size.R - 1) * layout.stride()[1] * problem_size.dilation_h + - conv_sign * int64_t(problem_size.S - 1) * layout.stride()[0] * problem_size.dilation_w + ) * element_size_bits / 8; + + // logical offset added to internal channel counter - units are elements, not bytes + filter_c_delta = threadblock_shape.column() * problem_size.split_k_slices; + } + +#if ENABLE_CONV2D_PARAMS_PRINT + /// Prints internal state. + MCTLASS_HOST_DEVICE + void print() { + auto stride = layout.stride(); + printf( + "Conv2dFpropActivationIteratorOptimizedParams:\n" + " layout(w: %d, h: %d, n: %d)\n" + " inc_next[%ld, %ld, %ld]\n" + " filter_c_delta(%d) - PQ(%d)\n" + " pq_divmod(divisor: %d, multiplier: %u, shift_right: %u)\n" + " q_divmod(divisor: %d, multiplier: %u, shift_right: %u)\n", + stride[0], stride[1], stride[2], + inc_next[0], inc_next[1], inc_next[2], + filter_c_delta, + PQ, + pq_divmod.divisor, + pq_divmod.multiplier, + pq_divmod.shift_right, + q_divmod.divisor, + q_divmod.multiplier, + q_divmod.shift_right + ); + } +#endif +}; + +/// Parameters structure used for Conv2dFpropActivationTileIteratorOptimized +template +struct Conv2dFpropActivationIteratorOptimizedParams> { + static int const kInterleaved = Interleaved_; + + using Layout = layout::TensorNCxHWx; + + Layout layout; + + int64_t inc_next[3]; // {next S, next R, next C} + int filter_c_delta; // number of logical elements to add to filter_c_ + int PQ; // product of P*Q + + FastDivmod pq_divmod; + FastDivmod q_divmod; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Conv2dFpropActivationIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv2dFpropActivationIteratorOptimizedParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, ///< layout object + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): + layout(layout), PQ(problem_size.P * problem_size.Q), pq_divmod(PQ), q_divmod(problem_size.Q) { + + TRACE_CONV_INITIALIZERS("conv2d_fprop", "activation", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + int conv_sign = (problem_size.mode == Mode::kConvolution ? -1 : 1); + + // next S + inc_next[0] = conv_sign * (kInterleaved * problem_size.dilation_w) * element_size_bits / 8; + + // next R + inc_next[1] = conv_sign * ( + int64_t(layout.stride()[0]) * problem_size.dilation_h + - (problem_size.S - 1) * kInterleaved * problem_size.dilation_w + ) * element_size_bits / 8; + + // next C + inc_next[2] = ( + threadblock_shape.column() * problem_size.split_k_slices / kInterleaved * int64_t(layout.stride()[1]) + - conv_sign * int64_t(problem_size.R - 1) * layout.stride()[0] * problem_size.dilation_h + - conv_sign * int64_t(problem_size.S - 1) * kInterleaved * problem_size.dilation_w + ) * element_size_bits / 8; + + // logical offset added to internal channel counter - units are elements, not bytes + filter_c_delta = threadblock_shape.column() * problem_size.split_k_slices; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template< typename Layout_ = layout::TensorNHWC > +struct Conv2dFpropFilterIteratorOptimizedParams; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +struct Conv2dFpropFilterIteratorOptimizedParams +{ + + using Layout = layout::TensorNHWC; + + Layout layout; + int RS; + int filter_c_delta; + + int64_t inc_next_k; // offset in units of bytes to next K position + int64_t inc_next_rs; // offset in units of bytes to next RS position + int64_t inc_next_c; // offset in units of bytes to next C position + + // + // Methods + // + MCTLASS_HOST_DEVICE + Conv2dFpropFilterIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv2dFpropFilterIteratorOptimizedParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): + layout(layout) { + + TRACE_CONV_INITIALIZERS("conv2d_fprop", "filter", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + RS = problem_size.R * problem_size.S; + + inc_next_k = (int64_t(layout.stride()[2]) * threadmap_delta.strided() * element_size_bits) / 8; + + inc_next_rs = + ( int64_t(layout.stride()[0]) + - int64_t(layout.stride()[2]) * (threadmap_iterations.strided() - 1) * threadmap_delta.strided() + ) * element_size_bits / 8; + + inc_next_c = + ( + threadblock_shape.row() * problem_size.split_k_slices + - int64_t(RS - 1) * layout.stride()[0] + - int64_t(threadmap_iterations.strided() - 1) * threadmap_delta.strided() * layout.stride()[2] + ) * element_size_bits / 8; + + filter_c_delta = threadblock_shape.row() * problem_size.split_k_slices; + } + +#if ENABLE_CONV2D_PARAMS_PRINT + /// Prints internal state. + MCTLASS_HOST_DEVICE + void print() { + auto stride = layout.stride(); + printf( + "Conv2dFpropFilterIteratorOptimizedParams:\n" + " layout[%d, %d, %d]\n" + " RS(%d), filter_c_delta(%d), inc_next(k: %ld, rs: %ld, c: %ld)\n", + stride[0], stride[1], stride[2], + RS, + filter_c_delta, + inc_next_k, inc_next_rs, inc_next_c + ); + } +#endif +}; + +template +struct Conv2dFpropFilterIteratorOptimizedParams> +{ + static int const kInterleaved = Interleaved_; + using Layout = layout::TensorCxRSKx; + + Layout layout; + int RS; + int filter_c_delta; + + int64_t inc_next_k; // offset in units of bytes to next K position + int64_t inc_next_rs; // offset in units of bytes to next RS position + int64_t inc_next_c; // offset in units of bytes to next C position + + // + // Methods + // + MCTLASS_HOST_DEVICE + Conv2dFpropFilterIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv2dFpropFilterIteratorOptimizedParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): + layout(layout) { + + TRACE_CONV_INITIALIZERS("conv2d_fprop", "filter", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + RS = problem_size.R * problem_size.S; + + inc_next_k = (kInterleaved * threadmap_delta.strided() * element_size_bits) / 8; + + inc_next_rs = + ( int64_t(layout.stride()[0]) + - kInterleaved * (threadmap_iterations.strided() - 1) * threadmap_delta.strided() + ) * element_size_bits / 8; + + inc_next_c = + ( + threadblock_shape.row() * problem_size.split_k_slices / kInterleaved * int64_t(layout.stride()[2]) + - int64_t(RS - 1) * layout.stride()[0] + - int64_t(threadmap_iterations.strided() - 1) * threadmap_delta.strided() * kInterleaved + ) * element_size_bits / 8; + + filter_c_delta = threadblock_shape.row() * problem_size.split_k_slices; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Dgrad Optimized Dy params (layout::TensorNHWC) +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Parameters object for Conv2d DGRAD OutputGradient (dy) iterator +struct Conv2dDgradOutputGradientIteratorOptimizedParams { + + using Layout = layout::TensorNHWC; + + Layout layout; + + int64_t inc_next[3]; // {next S, next R, next K} + + int filter_k_delta; // number of logical elements to add to filter_k_ + + int HW; // product of H*W + + FastDivmod hw_divmod; + FastDivmod w_divmod; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Conv2dDgradOutputGradientIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv2dDgradOutputGradientIteratorOptimizedParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): + layout(layout), + HW(problem_size.H *problem_size.W), + hw_divmod(HW), + w_divmod(problem_size.W) { + + TRACE_CONV_INITIALIZERS("conv2d_dgrad", "output_gradient", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + int conv_sign = (problem_size.mode == Mode::kConvolution ? 1 : -1); + + // next S + inc_next[0] = conv_sign * ( + (int64_t)layout.stride()[0] * problem_size.dilation_w + ) * element_size_bits / 8; + + // next R + inc_next[1] = conv_sign * ( + (int64_t)layout.stride()[1] * problem_size.dilation_h + - (problem_size.S - 1) * (int64_t)layout.stride()[0] * problem_size.dilation_w + ) * element_size_bits / 8; + + // next K + inc_next[2] = ( + threadblock_shape.column() * problem_size.split_k_slices + - conv_sign * (problem_size.R - 1) * (int64_t)layout.stride()[1] * problem_size.dilation_h + - conv_sign * (problem_size.S - 1) * (int64_t)layout.stride()[0] * problem_size.dilation_w + ) * element_size_bits / 8; + + // logical offset added to internal channel counter - units are elements, not bytes + filter_k_delta = threadblock_shape.column() * problem_size.split_k_slices; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Strided Dgrad Optimized Dy params (layout::TensorNHWC) +///////////////////////////////////////////////////////////////////////////////////////////////// +struct Conv2dStridedDgradOutputGradientIteratorOptimizedParams { + + using Layout = layout::TensorNHWC; + + Layout layout; + + int64_t inc_next[3]; // {next S, next R, next K} + + int filter_k_delta; // number of logical elements to add to filter_k_ + + int tiled_rows_per_filter; + + int conv_sign; + // + // Methods + // + + MCTLASS_HOST_DEVICE + Conv2dStridedDgradOutputGradientIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv2dStridedDgradOutputGradientIteratorOptimizedParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, ///< layout object + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape + ): layout(layout) { + + int tile_m_per_filter = strided_dgrad_tile_m_per_filter(problem_size, threadblock_shape.row()); + + tiled_rows_per_filter = tile_m_per_filter * threadblock_shape.row(); + + conv_sign = (problem_size.mode == Mode::kConvolution ? 1 : -1); + + // next S + inc_next[0] = conv_sign * ( + (int64_t)layout.stride()[0] * problem_size.dilation_w + ) * element_size_bits / 8; + + // next R + inc_next[1] = conv_sign * ( + (int64_t)layout.stride()[1] * problem_size.dilation_h + ) * element_size_bits / 8; + + // next K + inc_next[2] = ( + threadblock_shape.column() * problem_size.split_k_slices + ) * element_size_bits / 8; + + // logical offset added to internal channel counter - units are elements, not bytes + filter_k_delta = threadblock_shape.column() * problem_size.split_k_slices; + } +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////////////////////// +// Dgrad Optimized w params (layout::TensorNHWC) +///////////////////////////////////////////////////////////////////////////////////////////////// +struct Conv2dDgradFilterIteratorOptimizedParams { + + using Layout = layout::TensorNHWC; + + Layout layout; + int RS; + int filter_k_delta; + + int64_t inc_next_strided; // offset in units of bytes to next K coordinate within tile + int64_t inc_next_rs; // offset in units of bytes to next RS position + int64_t inc_next_k; // offset in units of bytes to next K position in subsequent tile + + // + // Methods + // + MCTLASS_HOST_DEVICE + Conv2dDgradFilterIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv2dDgradFilterIteratorOptimizedParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): + layout(layout), RS(problem_size.R * problem_size.S) { + + TRACE_CONV_INITIALIZERS("conv2d_dgrad", "filter", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + inc_next_strided = ((int64_t)layout.stride()[2] * threadmap_delta.strided() * element_size_bits) / 8; + + inc_next_rs = + ( (int64_t)layout.stride()[0] + - (threadmap_iterations.strided() - 1) * threadmap_delta.strided() * (int64_t)layout.stride()[2] + ) * element_size_bits / 8; + + inc_next_k = + ( + threadblock_shape.row() * problem_size.split_k_slices * (int64_t)layout.stride()[2] + - (problem_size.R * problem_size.S - 1) * (int64_t)layout.stride()[0] + - (threadmap_iterations.strided() - 1) * threadmap_delta.strided() * (int64_t)layout.stride()[2] + ) * element_size_bits / 8; + + filter_k_delta = threadblock_shape.row() * problem_size.split_k_slices; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////////////////////// +// StridedDgrad Optimized w params (layout::TensorNHWC) +///////////////////////////////////////////////////////////////////////////////////////////////// +struct Conv2dStridedDgradFilterIteratorOptimizedParams { + + using Layout = layout::TensorNHWC; + + Layout layout; + int RS; + int filter_k_delta; + + int64_t inc_next_strided; // offset in units of bytes to next K coordinate within tile + int64_t inc_next[3]; // {next S, next R, next K} + int64_t reset_bytes; // offset in units of bytes to move back the pointer + // + // Methods + // + MCTLASS_HOST_DEVICE + Conv2dStridedDgradFilterIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv2dStridedDgradFilterIteratorOptimizedParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): + layout(layout), RS(problem_size.R * problem_size.S) { + + TRACE_CONV_INITIALIZERS("conv2d_dgrad", "filter", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + inc_next_strided = (layout.stride()[2] * threadmap_delta.strided() * element_size_bits) / 8; + + // next S + inc_next[0] = + ( (int64_t)layout.stride()[0] * problem_size.stride_w + //- (threadmap_iterations.strided() - 1) * threadmap_delta.strided() * layout.stride()[2] + ) * element_size_bits / 8; + + // next R + inc_next[1] = + ( (int64_t)layout.stride()[1] * problem_size.stride_h + //- (threadmap_iterations.strided() - 1) * threadmap_delta.strided() * layout.stride()[2] + ) * element_size_bits / 8; + + // next K + inc_next[2] = + ( + threadblock_shape.row() * problem_size.split_k_slices * (int64_t)layout.stride()[2] + //- (problem_size.R * problem_size.S - 1) * layout.stride()[0] + //- (threadmap_iterations.strided() - 1) * threadmap_delta.strided() * layout.stride()[2] + ) * element_size_bits / 8; + + // offset in units of bytes to move the pointer in backward direction + reset_bytes = (threadmap_iterations.strided() - 1) * threadmap_delta.strided() * (int64_t)layout.stride()[2] + * element_size_bits / 8; + + filter_k_delta = threadblock_shape.row() * problem_size.split_k_slices; + } +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Parameters object for Conv2d WGRAD Output Gradient (dy) iterator +struct Conv2dWgradOutputGradientIteratorOptimizedParams { + + using Layout = layout::TensorNHWC; + + Layout layout; + + int NPQ; // precomputd product of N*P*Q for clearing predicates + + FastDivmod pq_divmod; + FastDivmod q_divmod; + + int64_t offset_next_strided; // offset in units of bytes to next npq coordinate within tile + int64_t offset_next_contiguous; // offset in units of bytes to next k coordinate within tile + int64_t inc_next_npq; // offset in units of bytes to next npq position in subsequent tile + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Conv2dWgradOutputGradientIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv2dWgradOutputGradientIteratorOptimizedParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): + layout(layout), + NPQ(problem_size.N * problem_size.P * problem_size.Q), + pq_divmod(problem_size.P * problem_size.Q), + q_divmod(problem_size.Q) { + + TRACE_CONV_INITIALIZERS("conv2d_wgrad", "output_gradient", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + // Incremental offsets in unites of bytes (number of elements) * sizeof_bits::value / 8 + offset_next_strided = (threadmap_delta.strided() * (int64_t)layout.stride()[0]) + * element_size_bits / 8; + + offset_next_contiguous = (threadmap_delta.contiguous()) + * element_size_bits / 8; + + inc_next_npq = (threadblock_shape.column() * problem_size.split_k_slices * (int64_t)layout.stride()[0]) + * element_size_bits / 8; + } +}; + +struct Conv2dWgradActivationIteratorOptimizedParams { + + using Layout = layout::TensorNHWC; + + Layout layout; + + FastDivmod sc_divmod; + FastDivmod pq_divmod; + FastDivmod q_divmod; + FastDivmod c_divmod; + FastDivmod s_divmod; + int small_channel_conv_s_offset; + + // + // Methods + // + MCTLASS_HOST_DEVICE + Conv2dWgradActivationIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv2dWgradActivationIteratorOptimizedParams( + Conv2dProblemSize const &problem_size, + Layout const &layout + ): + layout(layout), + sc_divmod(problem_size.S * problem_size.C), + pq_divmod(problem_size.P * problem_size.Q), + q_divmod(problem_size.Q), + c_divmod(problem_size.C), + s_divmod(problem_size.S * problem_size.dilation_w), + small_channel_conv_s_offset((problem_size.S - 1) * problem_size.dilation_w - problem_size.pad_w) { + } + + MCTLASS_HOST_DEVICE + Conv2dWgradActivationIteratorOptimizedParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): + Conv2dWgradActivationIteratorOptimizedParams( + problem_size, + layout + ) { + + TRACE_CONV_INITIALIZERS("conv2d_wgrad", "activation", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + } +}; + +struct PredicatedScaleBiasVectorAccessIteratorParams { + public: + /// Default ctor + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIteratorParams() { } + + // Default ctor + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIteratorParams( + Conv2dProblemSize const &problem_size, + layout::PitchLinear const &layout) {} + + // Default ctor + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIteratorParams( + Conv2dProblemSize const &problem_size, + layout::RowMajor const &layout) {} +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_tile_iterator.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_tile_iterator.h new file mode 100644 index 0000000..8194d10 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_tile_iterator.h @@ -0,0 +1,337 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template wraps the tile access iterator concept to load whole tiles from tensors in + memory used for implicit GEMM convolution. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +class TileIterator { +public: + using TileAccessIterator = TileAccessIterator_; + + using Shape = typename TileAccessIterator::Shape; + using Element = typename TileAccessIterator::Element; + using Layout = typename TileAccessIterator::Layout; + using TensorCoord = typename Layout::TensorCoord; + using ThreadMap = typename TileAccessIterator::ThreadMap; + using AccessType = typename TileAccessIterator::AccessType; + using TensorRef = typename TileAccessIterator::TensorRef; + using Index = typename TileAccessIterator::Index; + using LongIndex = typename TileAccessIterator::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = TileAccessIterator::kIteratorAlgorithm; + static StrideSupport const kStrideSupport = TileAccessIterator::kStrideSupport; + using Params = typename TileAccessIterator::Params; + static int const kConvDim = TileAccessIterator::kConvDim; + using ConvProblemSize = typename TileAccessIterator::ConvProblemSize; + static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + +private: + + /// Internal state + TileAccessIterator tile_access_iterator_; + +public: + + /// Constructor + MCTLASS_HOST_DEVICE + TileIterator( + Params const ¶ms, + ConvProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + tile_access_iterator_(params, problem_size, ptr, thread_idx, threadblock_offset) { } + + MCTLASS_HOST_DEVICE + static Params getParams(ConvProblemSize const &problem_size, Layout const &layout) { + return TileAccessIterator::getParams(problem_size, layout); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + tile_access_iterator_.set_iteration_index(index); + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + tile_access_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + TileIterator &operator++() { + tile_access_iterator_.advance(); + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + TileIterator operator++(int) { + TileIterator self(*this); + operator++(); + return self; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + frag.clear(); + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + mctlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + frag_ptr[idx], + tile_access_iterator_.get() + pointer_offset, + tile_access_iterator_.valid() + ); + + ++tile_access_iterator_; + } + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + tile_access_iterator_.set_iteration_index(0); + load_with_pointer_offset(frag, 0); + } + + MCTLASS_DEVICE + void advance() { + tile_access_iterator_.advance(); + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(ConvProblemSize const &problem_size) { + + // dispatch to iterator implementation + return TileAccessIterator::can_implement(problem_size); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Strided Dgrad Tile Iterator +template +class TileIteratorStridedDgrad { +public: + using TileAccessIterator = TileAccessIterator_; + + using Shape = typename TileAccessIterator::Shape; + using Element = typename TileAccessIterator::Element; + using Layout = typename TileAccessIterator::Layout; + using TensorCoord = typename Layout::TensorCoord; + using ThreadMap = typename TileAccessIterator::ThreadMap; + using AccessType = typename TileAccessIterator::AccessType; + using TensorRef = typename TileAccessIterator::TensorRef; + using Index = typename TileAccessIterator::Index; + using LongIndex = typename TileAccessIterator::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = TileAccessIterator::kIteratorAlgorithm; + static StrideSupport const kStrideSupport = TileAccessIterator::kStrideSupport; + using Params = typename TileAccessIterator::Params; + static int const kConvDim = TileAccessIterator::kConvDim; + using ConvProblemSize = typename TileAccessIterator::ConvProblemSize; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array< + Element, + ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess>; + +private: + + /// Internal state + TileAccessIterator tile_access_iterator_; + +public: + + /// Constructor (output gradient (Dy) OperandA ctor) + MCTLASS_HOST_DEVICE + TileIteratorStridedDgrad( + Params const ¶ms, + ConvProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + FastDivmod const &stride_h_divmod, FastDivmod const &stride_w_divmod, + int start_r, int start_s, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + tile_access_iterator_( + params, + problem_size, + ptr, + thread_idx, + stride_h_divmod, stride_w_divmod, + start_r, start_s, + threadblock_offset) { } + + /// Constructor (filter (w) OperandB ctor) + MCTLASS_HOST_DEVICE + TileIteratorStridedDgrad( + Params const ¶ms, + ConvProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + int start_r, int start_s, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + tile_access_iterator_(params, + problem_size, + ptr, + thread_idx, + start_r, start_s, + threadblock_offset) { } + + MCTLASS_HOST_DEVICE + static Params getParams(ConvProblemSize const &problem_size, Layout const &layout) { + return TileAccessIterator::getParams(problem_size, layout); + } + + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + tile_access_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + TileIteratorStridedDgrad &operator++() { + tile_access_iterator_.advance(); + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + TileIteratorStridedDgrad operator++(int) { + TileIteratorStridedDgrad self(*this); + operator++(); + return self; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + frag.clear(); + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + mctlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + frag_ptr[c + s * ThreadMap::Iterations::kContiguous], + tile_access_iterator_.get() + pointer_offset, + tile_access_iterator_.valid() + ); + + ++tile_access_iterator_; + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + tile_access_iterator_.set_iteration_index(0); + load_with_pointer_offset(frag, 0); + } + + MCTLASS_DEVICE + void advance() { + tile_access_iterator_.advance(); + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(ConvProblemSize const &problem_size) { + + // dispatch to iterator implementation + return TileAccessIterator::can_implement(problem_size); + } +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_analytic.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_analytic.h new file mode 100644 index 0000000..077d28a --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_analytic.h @@ -0,0 +1,285 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (activation tile) + matrix from memory. + + This iterator assumes TensorNHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dWgradActivationTileAccessIteratorAnalytic { +public: + + // + // Types + // + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNHWC; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + static_assert(sizeof_bits::value >= 8, + "WGRAD requires elements of size 8b or greater."); + + // + // Parameters structure + // + + using Params = Conv2dAnalyticParams; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + // Filter postion (r,s,c) in contiguous dimension stays constant for each gemm_iteration_k + int filter_r_[ThreadMap::Iterations::kContiguous]; + int filter_s_[ThreadMap::Iterations::kContiguous]; + int filter_c_[ThreadMap::Iterations::kContiguous]; + + int offset_npq_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dWgradActivationTileAccessIteratorAnalytic( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)) + { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + // initialize r,s,c filter position for every contiguous iteration + MCTLASS_PRAGMA_UNROLL + for(int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int rsc_offset = threadblock_offset.column() + thread_coord.contiguous() + + c * ThreadMap::Delta::kContiguous; + + filter_r_[c] = rsc_offset / (problem_size_.S * problem_size_.C); + int residual = rsc_offset % (problem_size_.S * problem_size_.C); + + filter_s_[c] = residual / problem_size_.C; + filter_c_[c] = residual % problem_size_.C; + } + + // initialize n, p, q offset for every strided iteration + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + offset_npq_[s] = threadblock_offset.row() + thread_coord.strided() + + s * ThreadMap::Delta::kStrided; + } + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + + // moves to the next GEMM-K offset (offset_npq_) in GEMM-B by a CTA-K tile + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_npq_[s] += Shape::kRow * problem_size_.split_k_slices; + } + } + + /// Returns the coordinate in the activation tensor x that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + int r, s, c; + + if (kAccessesPerVector == 1) { + /// One 128b aligned access fetching more than one element + c = filter_c_[iteration_contiguous_]; + r = filter_r_[iteration_contiguous_]; + s = filter_s_[iteration_contiguous_]; + } + else { + /// Multiple access to support non-128b alignment in contiguous dimension + c = (filter_c_[iteration_contiguous_] + iteration_vector_ * AccessType::kElements) % problem_size_.C; + int wrap_c = (filter_c_[iteration_contiguous_] + iteration_vector_ * AccessType::kElements) / problem_size_.C; + s = (filter_s_[iteration_contiguous_] + wrap_c) % problem_size_.S; + int wrap_s = (filter_s_[iteration_contiguous_] + wrap_c) / problem_size_.S; + r = filter_r_[iteration_contiguous_] + wrap_s; + } + + if (problem_size_.mode == Mode::kConvolution) { + r = (problem_size_.R - 1 - r); + s = (problem_size_.S - 1 - s); + } + + int n = offset_npq_[iteration_strided_] / (problem_size_.P * problem_size_.Q); + int residual = offset_npq_[iteration_strided_] % (problem_size_.P * problem_size_.Q); + + int p = residual / problem_size_.Q; + int q = residual % problem_size_.Q; + + int h = p * problem_size_.stride_h - problem_size_.pad_h + r * problem_size_.dilation_h; + int w = q * problem_size_.stride_w - problem_size_.pad_w + s * problem_size_.dilation_w; + + return TensorCoord(n, h, w, c); + } + + /// Returns true if the current coordinate is within the activation tensor x + MCTLASS_HOST_DEVICE + bool valid() const { + TensorCoord coord = at(); + + return coord.n() < problem_size_.N && + coord.h() >= 0 && coord.h() < problem_size_.H && + coord.w() >= 0 && coord.w() < problem_size_.W; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dWgradActivationTileAccessIteratorAnalytic &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_optimized.h new file mode 100644 index 0000000..48642ea --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_activation_tile_access_iterator_optimized.h @@ -0,0 +1,321 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (activation tile) + matrix from memory. + + This iterator assumes TensorNHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dWgradActivationTileAccessIteratorOptimized { +public: + + // + // Types + // + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNHWC; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + static_assert(sizeof_bits::value >= 8, + "WGRAD requires elements of size 8b or greater."); + + // + // Parameters structure + // + + using Params = Conv2dWgradActivationIteratorOptimizedParams; + +private: + + Conv2dWgradActivationIteratorOptimizedParams const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + // Precomputed effective filter postion (r,s) in contiguous dimension stays constant for each gemm_iteration_k + // required for npq -> nhw translation + int precomputed_filter_r_[ThreadMap::Iterations::kContiguous]; + int precomputed_filter_s_[ThreadMap::Iterations::kContiguous]; + + // Channel dimension in contiguous dimension stays constant for each gemm_iteration_k + int filter_c_[ThreadMap::Iterations::kContiguous]; + + int offset_npq_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dWgradActivationTileAccessIteratorOptimized( + Conv2dWgradActivationIteratorOptimizedParams const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)) + { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + // initialize r,s,c filter position for every contiguous iteration + MCTLASS_PRAGMA_UNROLL + for(int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int rsc_offset = threadblock_offset.column() + thread_coord.contiguous() + + c * ThreadMap::Delta::kContiguous; + + // The subseqnet fast_divmod() operations are equivalent to the following logical computation: + // + // + // filter_r_[c] = rsc_offset / (problem_size_.S * problem_size_.C); + // int residual = rsc_offset % (problem_size_.S * problem_size_.C); + // + // filter_s_[c] = residual / problem_size_.C; + // filter_c_[c] = residual % problem_size_.C; + + int residual; + params_.sc_divmod(precomputed_filter_r_[c], residual, rsc_offset); + params_.c_divmod(precomputed_filter_s_[c], filter_c_[c], residual); + + int r = precomputed_filter_r_[c]; + int s = precomputed_filter_s_[c]; + + if (problem_size_.mode == Mode::kConvolution) { + r = (problem_size_.R - 1 - r); + s = (problem_size_.S - 1 - s); + } + + precomputed_filter_r_[c] = -problem_size_.pad_h + r * problem_size_.dilation_h; + precomputed_filter_s_[c] = -problem_size_.pad_w + s * problem_size_.dilation_w; + } + + // initialize n, p, q offset for every strided iteration + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + offset_npq_[s] = threadblock_offset.row() + thread_coord.strided() + + s * ThreadMap::Delta::kStrided; + } + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + + // moves to the next GEMM-K offset (offset_npq_) in GEMM-B by a CTA-K tile + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_npq_[s] += Shape::kRow * problem_size_.split_k_slices; + } + } + + /// Returns the coordinate in the activation tensor x that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + int r = precomputed_filter_r_[iteration_contiguous_]; + int s = precomputed_filter_s_[iteration_contiguous_]; + int c = filter_c_[iteration_contiguous_]; + + if (kAccessesPerVector > 1) { + // This code section is only to support non-128b alignment + // Multiple access to support non-128b alignment in contiguous dimension + int wrap_c; + params_.c_divmod(wrap_c, c, c + iteration_vector_ * AccessType::kElements); + + if (problem_size_.mode == Mode::kConvolution) { + s -= (problem_size_.dilation_w * wrap_c); + + int wrap_s; + params_.s_divmod(wrap_s, s, params_.small_channel_conv_s_offset - s); + s = params_.small_channel_conv_s_offset - s; + + r -= (problem_size_.dilation_h * wrap_s); + + } else { + s += (problem_size_.dilation_w * wrap_c); + + int wrap_s; + params_.s_divmod(wrap_s, s, s + problem_size_.pad_w); + s -= problem_size_.pad_w; + + r += (problem_size_.dilation_h * wrap_s); + } + } + + // The subseqnet fast_divmod() operations are equivalent to the following logical computation: + // + // + // int n = offset_npq_[iteration_strided_] / (problem_size_.P * problem_size_.Q); + // int residual = offset_npq_[iteration_strided_] % (problem_size_.P * problem_size_.Q); + // + // int p = residual / problem_size_.Q; + // int q = residual % problem_size_.Q; + + int residual, n, p, q; + + params_.pq_divmod(n, residual, offset_npq_[iteration_strided_]); + params_.q_divmod(p, q, residual); + + int h = p * problem_size_.stride_h + r; + int w = q * problem_size_.stride_w + s; + + return TensorCoord(n, h, w, c); + } + + /// Returns true if the current coordinate is within the activation tensor x + MCTLASS_HOST_DEVICE + bool valid() const { + TensorCoord coord = at(); + + return coord.n() < problem_size_.N && + coord.h() >= 0 && coord.h() < problem_size_.H && + coord.w() >= 0 && coord.w() < problem_size_.W; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dWgradActivationTileAccessIteratorOptimized &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_analytic.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_analytic.h new file mode 100644 index 0000000..ca49047 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_analytic.h @@ -0,0 +1,260 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (output gradient tile) + matrix from memory. + + This iterator assumes TensorNHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dWgradOutputGradientTileAccessIteratorAnalytic { +public: + + // + // Types + // + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNHWC; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + static_assert(sizeof_bits::value >= 8, + "WGRAD requires elements of size 8b or greater."); + + // + // Parameters structure + // + + using Params = Conv2dAnalyticParams; + +private: + + Params const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + int filter_k_[ThreadMap::Iterations::kContiguous]; + + int offset_npq_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv2dWgradOutputGradientTileAccessIteratorAnalytic( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + // initialize filter_k for every contiguous iteration + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + filter_k_[c] = threadblock_offset.row() + thread_coord.contiguous() + + c * ThreadMap::Delta::kContiguous; + } + + // initialize n, p, q offset for every strided iteration + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_npq_[s] = threadblock_offset.column() + thread_coord.strided() + + s * ThreadMap::Delta::kStrided; + + } + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, layout); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + // moves to the next GEMM-K offset (offset_npq_) in GEMM-A by a CTA-K tile + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_npq_[s] += Shape::kColumn * problem_size_.split_k_slices; + } + } + + /// Returns the coordinate in the output gradient tensor Dy that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + int npq = offset_npq_[iteration_strided_]; + + int n = npq / (problem_size_.P * problem_size_.Q); + int residual = npq % (problem_size_.P * problem_size_.Q); + + int p = residual / problem_size_.Q; + int q = residual % problem_size_.Q; + + int k = filter_k_[iteration_contiguous_] + iteration_vector_ * AccessType::kElements; + + return TensorCoord(n, p, q, k); + } + + + /// Returns true if the current coordinate is within the output gradient tensor Dy + MCTLASS_HOST_DEVICE + bool valid() const { + TensorCoord coord = at(); + + return coord.n() < problem_size_.N && + coord.h() < problem_size_.P && + coord.w() < problem_size_.Q && + coord.c() < problem_size_.K; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dWgradOutputGradientTileAccessIteratorAnalytic &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.K % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_optimized.h new file mode 100644 index 0000000..96cc9d0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv2d_wgrad_output_gradient_tile_access_iterator_optimized.h @@ -0,0 +1,310 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (output gradient tile) + matrix from memory. + + This iterator assumes TensorNHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + typename AccessType_ = mctlass::AlignedArray +> +class Conv2dWgradOutputGradientTileAccessIteratorOptimized { +public: + + // + // Types + // + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNHWC; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + static_assert(sizeof_bits::value >= 8, + "WGRAD requires elements of size 8b or greater."); + + // + // Parameters structure + // + + using Params = Conv2dWgradOutputGradientIteratorOptimizedParams; + +private: + + Conv2dWgradOutputGradientIteratorOptimizedParams const ¶ms_; + Conv2dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + uint32_t predicates_[kAccessesPerVector]; + int filter_k_; + int offset_npq_; + +public: + + MCTLASS_HOST_DEVICE + Conv2dWgradOutputGradientTileAccessIteratorOptimized( + Conv2dWgradOutputGradientIteratorOptimizedParams const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + predicates_{0}, + filter_k_(0), + offset_npq_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_k_ = threadblock_offset.row() + thread_coord.contiguous(); + offset_npq_ = threadblock_offset.column() + thread_coord.strided(); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int filter_k = filter_k_ + c * ThreadMap::Delta::kContiguous; + int offset_npq = offset_npq_ + s * ThreadMap::Delta::kStrided; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + bool predicate = valid_(at_(offset_npq, filter_k + v * AccessType::kElements)); + + uint32_t pred = (predicate ? 1u : 0); + + int pred_idx = c + s * ThreadMap::Iterations::kContiguous; + + predicates_[v] |= (pred << pred_idx); + } + } + } + + // Offset pointer to (iteration_strided_, iteration_contiguous_) = (0, 0) + pointer_ += ( + offset_npq_ * params.layout.stride()[0] + filter_k_ + ) * sizeof_bits::value / 8; + + set_iteration_index(0); + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}, + ThreadMap::kThreads, + ThreadMap::kElementsPerAccess, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided}); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + // moves to the next GEMM-K offset (offset_npq_) in GEMM-A by a CTA-K tile + offset_npq_ += Shape::kColumn * problem_size_.split_k_slices; + + // Clear predicates if needed + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + if (offset_npq_ + s * ThreadMap::Delta::kStrided >= params_.NPQ) { + uint32_t kClearMask = ((1u << ThreadMap::Iterations::kContiguous) - 1) << (s * ThreadMap::Iterations::kContiguous); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + predicates_[v] = (predicates_[v] & (~kClearMask)); + } + } + } + + pointer_ += params_.inc_next_npq; + } + +private: + /// Returns the coordinate in the output gradient tensor Dy that is pointed to + /// by offset_npq and k. + MCTLASS_HOST_DEVICE + TensorCoord at_(int offset_npq, int k) const { + + // The subsequent fast_divmod() operations are equivalent to the following logical computation: + // + // + // int npq = offset_npq; + // int n = npq / (problem_size_.P * problem_size_.Q); + // int residual = npq % (problem_size_.P * problem_size_.Q); + // + // int p = residual / problem_size_.Q; + // int q = residual % problem_size_.Q; + + int residual, n, p, q; + + params_.pq_divmod(n, residual, offset_npq); + params_.q_divmod(p, q, residual); + + return TensorCoord(n, p, q, k); + } + + /// Returns true if the coord is within the output gradient tensor Dy + MCTLASS_HOST_DEVICE + bool valid_(TensorCoord coord) const { + + return coord.n() < problem_size_.N && + coord.c() < problem_size_.K; + } + +public: + + /// Returns true if the current coordinate is within the output gradient tensor Dy + MCTLASS_HOST_DEVICE + bool valid() const { + + LongIndex pred_idx = iteration_contiguous_ + iteration_strided_ * ThreadMap::Iterations::kContiguous; + return (predicates_[iteration_vector_] & (1u << pred_idx)); + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + return reinterpret_cast( + pointer_ + + iteration_strided_ * params_.offset_next_strided + + iteration_contiguous_ * params_.offset_next_contiguous + ) + iteration_vector_; + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv2dWgradOutputGradientTileAccessIteratorOptimized &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.K % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_filter_tile_access_iterator_analytic.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_filter_tile_access_iterator_analytic.h new file mode 100644 index 0000000..699a09e --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_filter_tile_access_iterator_analytic.h @@ -0,0 +1,268 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (filter tile) + matrix from memory. + + This iterator assumes TensorNDHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv3d_problem_size.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_ +> +class Conv3dDgradFilterTileAccessIteratorAnalytic { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNDHWC; + using ThreadMap = ThreadMap_; + using AccessType = AlignedArray; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 3; + using ConvProblemSize = typename conv::Conv3dProblemSize; + static int const kAccessesPerVector = 1; + + static_assert(sizeof_bits::value >= 8, + "DGRAD requires elements of size 8b or larger."); + + // + // Parameters structure + // + + struct Params { + + Layout layout; + + // + // Methods + // + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params( + Conv3dProblemSize const &problem_size, + Layout const &layout + ): layout(layout) { + + } + }; + +private: + + Params const ¶ms_; + Conv3dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + char const *pointer_; + + // For a fixed filter position (t,r,s) find and fill offset_k_, offset_c_ in strided and contiguous dimension + int filter_t_; + int filter_r_; + int filter_s_; + int offset_k_[ThreadMap::Iterations::kStrided]; + int offset_c_[ThreadMap::Iterations::kContiguous]; + +public: + + MCTLASS_HOST_DEVICE + Conv3dDgradFilterTileAccessIteratorAnalytic( + Params const ¶ms, + Conv3dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + filter_t_(0), + filter_r_(0), + filter_s_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + offset_c_[c] = threadblock_offset.column() + thread_coord.contiguous() + + c * ThreadMap::Delta::kContiguous; + } + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_k_[s] = + threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + } + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + // moves to the next tile + ++filter_s_; + if (filter_s_ < problem_size_.S) { + return; + } + filter_s_ = 0; + ++filter_r_; + if (filter_r_ < problem_size_.R) { + return; + } + filter_r_ = 0; + ++filter_t_; + if (filter_t_ < problem_size_.T) { + return; + } + filter_t_ = 0; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_k_[s] += Shape::kRow * problem_size_.split_k_slices; + } + } + + /// Returns the coordinate in the filter tensor w that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + int c = offset_c_[iteration_contiguous_]; + int k = offset_k_[iteration_strided_]; + + return TensorCoord(k, filter_t_, filter_r_, filter_s_, c); + } + + /// Returns true if the current coordinate is within the filter tensor w + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + return coord.n() < problem_size_.K && coord.c() < problem_size_.C; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv3dDgradFilterTileAccessIteratorAnalytic &operator++() { + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv3dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % (128/sizeof_bits::value)) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_filter_tile_access_iterator_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_filter_tile_access_iterator_optimized.h new file mode 100644 index 0000000..70448e5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_filter_tile_access_iterator_optimized.h @@ -0,0 +1,289 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (filter tile) + matrix from memory. + + This iterator assumes TensorNHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv3d_problem_size.h" + +#include "mctlass/conv/threadblock/conv3d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + conv::StrideSupport StrideSupport_ = conv::StrideSupport::kUnity +> +class Conv3dDgradFilterTileAccessIteratorOptimized { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNDHWC; + using ThreadMap = ThreadMap_; + using AccessType = AlignedArray; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = StrideSupport_; + static int const kConvDim = 3; + using ConvProblemSize = typename conv::Conv3dProblemSize; + static int const kAccessesPerVector = 1; + + // + // Parameters structure + // + + struct Params : Conv3dDgradFilterIteratorOptimizedParams { + + // + // Methods + // + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Conv3dDgradFilterIteratorOptimizedParams const &base): + Conv3dDgradFilterIteratorOptimizedParams(base) { } + + MCTLASS_HOST_DEVICE + Params( + Conv3dProblemSize const &problem_size, + Layout const &layout + ): + Conv3dDgradFilterIteratorOptimizedParams( + problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}, + ThreadMap::kThreads, + ThreadMap::kElementsPerAccess, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided} + ) { } + + }; + +private: + + Conv3dDgradFilterIteratorOptimizedParams const ¶ms_; + Conv3dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + char const *pointer_; + + uint32_t predicates_; + int filter_trs_; + int filter_k_; + + // + // Assertions + // + + // We map predicates into bits packed in this uint32_t container + static_assert(ThreadMap::Iterations::kStrided * + ThreadMap::Iterations::kContiguous < sizeof(predicates_) * 8, + "Currently, the number of loads per iteration is limited by the size of the predicates container."); + +public: + + MCTLASS_HOST_DEVICE + Conv3dDgradFilterTileAccessIteratorOptimized( + Conv3dDgradFilterIteratorOptimizedParams const ¶ms, + Conv3dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + predicates_(0), + filter_trs_(0), + filter_k_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_k_ = threadblock_offset.row() + thread_coord.strided(); + Index column = threadblock_offset.column() + thread_coord.contiguous(); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int filter_k = filter_k_ + s * ThreadMap::Delta::kStrided; + int filter_c = column + c * ThreadMap::Delta::kContiguous; + + uint32_t pred = ((filter_k < problem_size_.K && filter_c < problem_size_.C) ? 1u : 0); + + int pred_idx = c + s * ThreadMap::Iterations::kContiguous; + + predicates_ |= (pred << pred_idx); + } + } + + pointer_ += ( + filter_k_ * params.layout.stride()[3] + column + ) * sizeof_bits::value / 8; + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + + LongIndex next = params_.inc_next_trs; + + // moves to the next tile + ++filter_trs_; + if (filter_trs_ == params_.TRS) { + + filter_trs_ = 0; + next = params_.inc_next_k; + filter_k_ += params_.filter_k_delta; + } + + // Clear predicates if needed + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + if (filter_k_ + s * ThreadMap::Delta::kStrided >= problem_size_.K) { + uint32_t kClearMask = ((1u << ThreadMap::Iterations::kContiguous) - 1) << (s * ThreadMap::Iterations::kContiguous); + + predicates_ = (predicates_ & (~kClearMask)); + } + } + + pointer_ += next; + } + + /// Returns true if the current coordinate is within the filter tensor W + MCTLASS_HOST_DEVICE + bool valid() { + LongIndex pred_idx = iteration_contiguous_ + iteration_strided_ * ThreadMap::Iterations::kContiguous; + return (predicates_ & (1u << pred_idx)); + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + return reinterpret_cast(pointer_ + + iteration_contiguous_ * ThreadMap::Delta::kContiguous * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv3dDgradFilterTileAccessIteratorOptimized &operator++() { + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + + // Move to the next K coordinate within the tile + pointer_ += params_.inc_next_strided; + + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv3dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % (128/sizeof_bits::value)) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_output_gradient_tile_access_iterator_analytic.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_output_gradient_tile_access_iterator_analytic.h new file mode 100644 index 0000000..ef0410d --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_output_gradient_tile_access_iterator_analytic.h @@ -0,0 +1,343 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (output gradient tile) + matrix from memory. + + This iterator assumes TensorNDHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv3d_problem_size.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + conv::StrideSupport StrideSupport_ = conv::StrideSupport::kStrided +> +class Conv3dDgradOutputGradientTileAccessIteratorAnalytic; +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Conv3dDgradOutputGradientTileAccessIteratorAnalytic strided dgrad needs special handling using +// unscaled coordinations +template < + typename Shape_, + typename Element_, + typename ThreadMap_ +> +class Conv3dDgradOutputGradientTileAccessIteratorAnalytic < + Shape_, + Element_, + ThreadMap_, + conv::StrideSupport::kStrided +> { +public: + + // + // Types + // + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNDHWC; + using ThreadMap = ThreadMap_; + using AccessType = AlignedArray; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 3; + using ConvProblemSize = typename conv::Conv3dProblemSize; + static int const kAccessesPerVector = 1; + + static_assert(sizeof_bits::value >= 8, + "DGRAD requires elements of size 8b or greater."); + + // + // Simpligying assertions + // + + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + struct Params { + + Layout layout; + + // + // Methods + // + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params( + ConvProblemSize const &problem_size, + Layout const &layout + ): layout(layout) { + + } + }; + +private: + + Params const ¶ms_; + ConvProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + char const *pointer_; + + int filter_k_; + int filter_t_; + int filter_r_; + int filter_s_; + + int offset_n_[ThreadMap::Iterations::kStrided]; + int offset_d_[ThreadMap::Iterations::kStrided]; + int offset_w_[ThreadMap::Iterations::kStrided]; + int offset_h_[ThreadMap::Iterations::kStrided]; + +private: + + /// Returns the coordinate in the output tensor Dy that is currently pointed to + /// by the iterator but DOES NOT scale by the convolution stride. This is needed + /// to compute predicates in the valid() method. The return value of the public at() + /// method is correctly scaled. + MCTLASS_HOST_DEVICE + TensorCoord unscaled_at_() const { + int n = offset_n_[iteration_strided_]; + int d = offset_d_[iteration_strided_]; + int h = offset_h_[iteration_strided_]; + int w = offset_w_[iteration_strided_]; + + int t = filter_t_; + int r = filter_r_; + int s = filter_s_; + + if (problem_size_.mode == Mode::kConvolution) { + t = (problem_size_.T - 1 - t); + r = (problem_size_.R - 1 - r); + s = (problem_size_.S - 1 - s); + } + + int z = (d + problem_size_.pad_d - t * problem_size_.dilation_d); + int p = (h + problem_size_.pad_h - r * problem_size_.dilation_h); + int q = (w + problem_size_.pad_w - s * problem_size_.dilation_w); + + return TensorCoord(n, z, p, q, filter_k_); + } + +public: + + MCTLASS_HOST_DEVICE + Conv3dDgradOutputGradientTileAccessIteratorAnalytic( + Params const ¶ms, + ConvProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() // threadblock offset - units are whole CTA tiles + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + filter_k_(0), + filter_t_(0), + filter_r_(0), + filter_s_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_k_ = threadblock_offset.column() + thread_coord.contiguous(); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + int offset_ndhw = threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + + offset_n_[s] = offset_ndhw / (problem_size_.D * problem_size_.H * problem_size_.W); + int residual = offset_ndhw % (problem_size_.D * problem_size_.H * problem_size_.W); + + offset_d_[s] = residual / (problem_size_.H * problem_size_.W); + residual = residual % (problem_size_.H * problem_size_.W); + + offset_h_[s] = residual / problem_size_.W; + offset_w_[s] = residual % problem_size_.W; + } + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv3dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, layout); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + // move to the next tile + ++filter_s_; + if (filter_s_ < problem_size_.S) { + return; + } + filter_s_ = 0; + ++filter_r_; + if (filter_r_ < problem_size_.R) { + return; + } + filter_r_ = 0; + ++filter_t_; + if (filter_t_ < problem_size_.T) { + return; + } + filter_t_ = 0; + + filter_k_ += Shape_::kColumn * problem_size_.split_k_slices; + } + + /// Returns the coordinate in the output tensor Dy that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + TensorCoord coord = unscaled_at_(); + + return TensorCoord( + coord.n(), + coord.d() / problem_size_.stride_d, + coord.h() / problem_size_.stride_h, + coord.w() / problem_size_.stride_w, + coord.c()); + } + + + /// Returns true if the current coordinate is within the output tensor Dy + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord unscaled_coord = unscaled_at_(); + TensorCoord coord = at(); + + return + !(unscaled_coord.d() % problem_size_.stride_d) && + !(unscaled_coord.h() % problem_size_.stride_h) && + !(unscaled_coord.w() % problem_size_.stride_w) && + coord.n() < problem_size_.N && + coord.d() >= 0 && coord.d() < problem_size_.Z && + coord.h() >= 0 && coord.h() < problem_size_.P && + coord.w() >= 0 && coord.w() < problem_size_.Q && + coord.c() < problem_size_.K; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv3dDgradOutputGradientTileAccessIteratorAnalytic &operator++() { + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(ConvProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.K % (128/sizeof_bits::value)) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_output_gradient_tile_access_iterator_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_output_gradient_tile_access_iterator_optimized.h new file mode 100644 index 0000000..de58de6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_dgrad_output_gradient_tile_access_iterator_optimized.h @@ -0,0 +1,490 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (output gradient tile) + matrix from memory. + + This iterator assumes TensorNDHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/conv/threadblock/conv3d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_, + conv::StrideSupport StrideSupport_ = conv::StrideSupport::kUnity +> +class Conv3dDgradOutputGradientTileAccessIteratorOptimized { +public: + + static_assert(StrideSupport_ == conv::StrideSupport::kUnity, + "Only unit-stride dgrad is supported at this time."); + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNDHWC; + using TensorCoord = typename Layout::TensorCoord; + using ThreadMap = ThreadMap_; + using AccessType = AlignedArray; + using TensorRef = mctlass::TensorRef; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kUnity; + static int const kConvDim = 3; + using ConvProblemSize = typename conv::Conv3dProblemSize; + using Coord3D = Coord<3>; + static int const kAccessesPerVector = 1; + using Mask = uint64_t; + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv3dDgradOutputGradientIteratorOptimizedParams; + +private: + + Params const ¶ms_; + ConvProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + + + // One pointer per access + char const *pointer_[ThreadMap::Iterations::kStrided]; + + // current filter position (t, r, s) + int filter_t_; + int filter_r_; + int filter_s_; + int filter_k_; + + Index masks_[ThreadMap::Iterations::kStrided][3]; + +public: + + MCTLASS_HOST_DEVICE + Conv3dDgradOutputGradientTileAccessIteratorOptimized( + Params const ¶ms, + ConvProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() // tile index - units are threadblock-scoped tiles + ): + params_(params), + problem_size_(problem_size), + filter_k_(0), + filter_t_(0), + filter_r_(0), + filter_s_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_k_ = threadblock_offset.column() + thread_coord.contiguous(); + + int offset_n[ThreadMap::Iterations::kStrided]; + int offset_d[ThreadMap::Iterations::kStrided]; + int offset_h[ThreadMap::Iterations::kStrided]; + int offset_w[ThreadMap::Iterations::kStrided]; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + pointer_[s] = reinterpret_cast(ptr); + + int offset_ndhw = threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + + // The subseqnet fast_divmod() operations are equivalent to the following logical computation: + // + // + // offset_n[s] = offset_ndhw / (problem_size_.D * problem_size_.H * problem_size_.W); + // int residual = offset_ndhw % (problem_size_.D * problem_size_.H * problem_size_.W); + // + // + // offset_d[s] = residual / (problem_size_.H * problem_size_.W); + // residual = residual % (problem_size_.H * problem_size_.W); + // + // offset_h[s] = residual / problem_size_.W; + // offset_w[s] = residual % problem_size_.W; + // + + int residual; + + // input: (ndhw offset) output: (n offset and resudial (dhw offset)) + params_.dhw_divmod(offset_n[s], residual, offset_ndhw); + // input: (dhw offset) output: (d offset and resudial (hw)) + params_.hw_divmod(offset_d[s], residual, residual); + // input: (hw offset) output: (h offset and resudial (w offset)) + params_.w_divmod(offset_h[s], offset_w[s], residual); + + TensorCoord coord = at_(offset_n[s], offset_d[s], offset_h[s], offset_w[s], 0, 0, 0); + + pointer_[s] += params_.layout(coord) * sizeof_bits::value / 8; + } + + clear_mask(); + + MCTLASS_PRAGMA_NO_UNROLL + for (int t = 0; t < problem_size_.T; ++t) { + MCTLASS_PRAGMA_UNROLL + for (int s_idx = 0; s_idx < ThreadMap::Iterations::kStrided; ++s_idx) { + + int t_ = t; + if (problem_size_.mode == Mode::kConvolution) { + t_ = problem_size_.T - 1 - t; + } + + int z = offset_d[s_idx] + problem_size_.pad_d - t_ * problem_size_.dilation_d; + + bool pred = (offset_n[s_idx] < problem_size_.N && z >= 0 && z < problem_size_.Z); + masks_[s_idx][0] |= (pred << t); + } + } + + MCTLASS_PRAGMA_NO_UNROLL + for (int r = 0; r < problem_size_.R; ++r) { + MCTLASS_PRAGMA_UNROLL + for (int s_idx = 0; s_idx < ThreadMap::Iterations::kStrided; ++s_idx) { + + int r_ = r; + if (problem_size_.mode == Mode::kConvolution) { + r_ = problem_size_.R - 1 - r; + } + + int p = offset_h[s_idx] + problem_size_.pad_h - r_ * problem_size_.dilation_h; + + bool pred = (p >= 0 && p < problem_size_.P); + masks_[s_idx][1] |= (pred << r); + } + } + + MCTLASS_PRAGMA_NO_UNROLL + for (int s = 0; s < problem_size_.S; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int s_idx = 0; s_idx < ThreadMap::Iterations::kStrided; ++s_idx) { + + int s_ = s; + if (problem_size_.mode == Mode::kConvolution) { + s_ = problem_size_.S - 1 - s; + } + + int q = offset_w[s_idx] + problem_size_.pad_w - s_ * problem_size_.dilation_w; + + bool pred = (q >= 0 && q < problem_size_.Q); + masks_[s_idx][2] |= (pred << s); + } + } + + if (filter_k_ >= problem_size.K) { + clear_mask(); + } + + set_iteration_index(0); + + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv3dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}, + ThreadMap::kThreads, + ThreadMap::kElementsPerAccess, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided}); + } + +private: + + + /// Returns the coordinate in the output gradient tensor dy that is correspoinding to + // activation ndhw and filter position k, t, r, s + MCTLASS_HOST_DEVICE + TensorCoord at_(int n, int d, int h, int w, int t, int r, int s) const { + + if (problem_size_.mode == Mode::kConvolution) { + t = problem_size_.T - 1 - t; + r = problem_size_.R - 1 - r; + s = problem_size_.S - 1 - s; + } + + int z = d + problem_size_.pad_d - t * problem_size_.dilation_d; + int p = h + problem_size_.pad_h - r * problem_size_.dilation_h; + int q = w + problem_size_.pad_w - s * problem_size_.dilation_w; + + return TensorCoord(n, z, p, q, filter_k_); + } + + + /// Adds a pointer offset in units of element + MCTLASS_HOST_DEVICE + void add_byte_offset_(LongIndex byte_offset) { + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + pointer_[s] += byte_offset; + } + } + + /// Clears the predicates + MCTLASS_HOST_DEVICE + void clear_mask_(bool clear) { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + // We are using inline PTX assembly here to avoid an CUDA C++ compilation + // artifact in which control flow instructions are generated. Instead, our + // intent is to predicate the mov instructions. + //#if defined(__MACA_ARCH__) + #if 0 + asm volatile( + "{\n" + " .reg .pred p;\n" + " .reg .u32 m;" + " mov.u32 m, %2;" + " setp.ne.b32 p, %1, 0;\n" + " @p mov.u32 m, 0;\n" + " mov.u32 %0, m;\n" + "}\n" + : + "=r"(masks_[s][0]) + : + "r"((int)clear), + "r"(masks_[s][0]) + ); + asm volatile( + "{\n" + " .reg .pred p;\n" + " .reg .u32 m;" + " mov.u32 m, %2;" + " setp.ne.b32 p, %1, 0;\n" + " @p mov.u32 m, 0;\n" + " mov.u32 %0, m;\n" + "}\n" + : + "=r"(masks_[s][1]) + : + "r"((int)clear), + "r"(masks_[s][1]) + ); + asm volatile( + "{\n" + " .reg .pred p;\n" + " .reg .u32 m;" + " mov.u32 m, %2;" + " setp.ne.b32 p, %1, 0;\n" + " @p mov.u32 m, 0;\n" + " mov.u32 %0, m;\n" + "}\n" + : + "=r"(masks_[s][2]) + : + "r"((int)clear), + "r"(masks_[s][2]) + ); + #else + if (clear) { + masks_[s][0] = 0; + masks_[s][1] = 0; + masks_[s][2] = 0; + } + #endif + } + } + +public: + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + add_byte_offset_(pointer_offset * sizeof_bits::value / 8); + } + + + MCTLASS_HOST_DEVICE + void advance() { + + int next_idx = 0; + + // moves to the next tile + ++filter_s_; + if (filter_s_ == problem_size_.S) { + + filter_s_ = 0; + ++filter_r_; + next_idx = 1; + + if (filter_r_ == problem_size_.R) { + filter_r_ = 0; + ++filter_t_; + + if (filter_t_ < problem_size_.T) { + next_idx = 2; + } + else { + filter_t_ = 0; + next_idx = 3; + } + } + } + + add_byte_offset_(params_.inc_next[next_idx]); + + if (next_idx == 3) { + filter_k_ += params_.filter_k_delta; + } + + clear_mask_(filter_k_ >= problem_size_.K); + } + + + /// Clears the predicates + MCTLASS_HOST_DEVICE + void clear_mask() { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + masks_[s][0] = Mask(0); + masks_[s][1] = Mask(0); + masks_[s][2] = Mask(0); + } + } + + MCTLASS_HOST_DEVICE + bool valid() { + + return + (masks_[iteration_strided_][0] & (Index(1) << filter_t_)) && + (masks_[iteration_strided_][1] & (Index(1) << filter_r_)) && + (masks_[iteration_strided_][2] & (Index(1) << filter_s_)); + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + return reinterpret_cast(pointer_[iteration_strided_]); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv3dDgradOutputGradientTileAccessIteratorOptimized &operator++() { + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(ConvProblemSize const &problem_size) { + + // This is specialized for unit stride + if (problem_size.stride() != Coord3D({1, 1, 1})) { + return Status::kErrorNotSupported; + } + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.K % (128/sizeof_bits::value)) { + return Status::kErrorNotSupported; + } + + // Limit on filter size + if (problem_size.T > 32 || problem_size.R > 32 || problem_size.S > 32) { + return Status::kErrorNotSupported; + } + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + + diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_analytic.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_analytic.h new file mode 100644 index 0000000..2407f70 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_analytic.h @@ -0,0 +1,291 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (activation tile) + matrix from memory. + + This iterator assumes TensorNDHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/conv/threadblock/conv3d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_ +> +class Conv3dFpropActivationTileAccessIteratorAnalytic { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNDHWC; + using TensorCoord = typename Layout::TensorCoord; + using ThreadMap = ThreadMap_; + using AccessType = AlignedArray; + using TensorRef = mctlass::TensorRef; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 3; + using ConvProblemSize = typename conv::Conv3dProblemSize; + static int const kAccessesPerVector = 1; + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv3dAnalyticParams; + +private: + + Params const ¶ms_; + ConvProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + char const *pointer_; + + int filter_t_; + int filter_r_; + int filter_s_; + int filter_c_; + + int offset_n_[ThreadMap::Iterations::kStrided]; + int offset_z_[ThreadMap::Iterations::kStrided]; + int offset_p_[ThreadMap::Iterations::kStrided]; + int offset_q_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv3dFpropActivationTileAccessIteratorAnalytic( + Params const ¶ms, + ConvProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() // tile index - units are threadblock-scoped tiles + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + filter_t_(0), + filter_r_(0), + filter_s_(0), + filter_c_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_c_ = threadblock_offset.column() + thread_coord.contiguous(); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + int offset_nzpq = threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + + offset_n_[s] = offset_nzpq / (problem_size_.Z * problem_size_.P * problem_size_.Q); + int residual = offset_nzpq % (problem_size_.Z * problem_size_.P * problem_size_.Q); + + offset_z_[s] = residual / (problem_size_.P * problem_size_.Q); + residual = residual % (problem_size_.P * problem_size_.Q); + + offset_p_[s] = residual / problem_size_.Q; + offset_q_[s] = residual % problem_size_.Q; + } + + set_iteration_index(0); + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv3dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, layout); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + // moves to the next tile + ++filter_s_; + if (filter_s_ < problem_size_.S) { + return; + } + filter_s_ = 0; + ++filter_r_; + if (filter_r_ < problem_size_.R) { + return; + } + filter_r_ = 0; + ++filter_t_; + if (filter_t_ < problem_size_.T) { + return; + } + filter_t_ = 0; + + filter_c_ += Shape::kColumn * problem_size_.split_k_slices; + } + + /// Returns the coordinate in the activations tensor X that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + int n = offset_n_[iteration_strided_]; + int z = offset_z_[iteration_strided_]; + int p = offset_p_[iteration_strided_]; + int q = offset_q_[iteration_strided_]; + + int t = filter_t_; + int r = filter_r_; + int s = filter_s_; + + if (problem_size_.mode == Mode::kConvolution) { + t = (problem_size_.T - 1 - filter_t_); + r = (problem_size_.R - 1 - filter_r_); + s = (problem_size_.S - 1 - filter_s_); + } + + int d = z * problem_size_.stride_d - problem_size_.pad_d + t * problem_size_.dilation_d; + int h = p * problem_size_.stride_h - problem_size_.pad_h + r * problem_size_.dilation_h; + int w = q * problem_size_.stride_w - problem_size_.pad_w + s * problem_size_.dilation_w; + + return TensorCoord(n, d, h, w, filter_c_); + } + + /// Returns true if the current coordinate is within the activations tensor X + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + return coord.n() < problem_size_.N && + coord.d() >= 0 && coord.d() < problem_size_.D && + coord.h() >= 0 && coord.h() < problem_size_.H && + coord.w() >= 0 && coord.w() < problem_size_.W && + coord.c() < problem_size_.C; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + AccessType const *ptr = reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + + return ptr; + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv3dFpropActivationTileAccessIteratorAnalytic &operator++() { + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(ConvProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % (128/sizeof_bits::value)) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + + diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_optimized.h new file mode 100644 index 0000000..04ee175 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_activation_tile_access_iterator_optimized.h @@ -0,0 +1,479 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (activation tile) + matrix from memory. + + This iterator assumes TensorNDHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/conv/threadblock/conv3d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename Layout_, + typename ThreadMap_ +> +class Conv3dFpropActivationTileAccessIteratorOptimized { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using TensorCoord = typename Layout::TensorCoord; + using ThreadMap = ThreadMap_; + using AccessType = AlignedArray; + using TensorRef = mctlass::TensorRef; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 3; + using ConvProblemSize = typename conv::Conv3dProblemSize; + static int const kAccessesPerVector = 1; + using Mask = uint64_t; + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv3dFpropActivationIteratorOptimizedParams; + +private: + + Conv3dFpropActivationIteratorOptimizedParams const ¶ms_; + Conv3dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + + // One pointer per access + char const *pointer_[ThreadMap::Iterations::kStrided]; + + // current filter position (t, r, s) + int filter_t_; + int filter_r_; + int filter_s_; + int filter_c_; + + // mask for t, r, and s + Index masks_[ThreadMap::Iterations::kStrided][3]; + +public: + + MCTLASS_HOST_DEVICE + Conv3dFpropActivationTileAccessIteratorOptimized( + Conv3dFpropActivationIteratorOptimizedParams const ¶ms, + Conv3dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() // tile index - units are threadblock-scoped tiles + ) : + params_(params), + problem_size_(problem_size), + filter_t_(0), + filter_r_(0), + filter_s_(0), + filter_c_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_c_ = threadblock_offset.column() + thread_coord.contiguous(); + + int offset_n[ThreadMap::Iterations::kStrided]; + int offset_z[ThreadMap::Iterations::kStrided]; + int offset_p[ThreadMap::Iterations::kStrided]; + int offset_q[ThreadMap::Iterations::kStrided]; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + pointer_[s] = reinterpret_cast(ptr); + + int offset_nzpq = threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + + // The subseqnet fast_divmod() operations are equivalent to the following logical computation: + // + // + // offset_n[s] = offset_nzpq / (problem_size_.Z * problem_size_.P * problem_size_.Q); + // int residual = offset_nzpq % (problem_size_.Z * problem_size_.P * problem_size_.Q); + // + // offset_z[s] = residual / (problem_size_.P * problem_size_.Q); + // residual = residual % (problem_size_.P * problem_size_.Q); + // + // offset_p[s] = residual / problem_size_.Q; + // offset_q[s] = residual % problem_size_.Q; + // + + int residual; + + // input: (nzpq offset) output: (n offset and resudial (zpq offset)) + params.zpq_divmod(offset_n[s], residual, offset_nzpq); + // input: (zpq offset) output: (z offset and resudial (pq)) + params.pq_divmod(offset_z[s], residual, residual); + // input: (pq offset) output: (p offset and resudial (q offset)) + params.q_divmod(offset_p[s], offset_q[s], residual); + + TensorCoord coord = at_(offset_n[s], offset_z[s], offset_p[s], offset_q[s], 0, 0, 0); + + pointer_[s] += params_.layout(coord) * sizeof_bits::value / 8; + } + + clear_mask(); + + // mask predicates for filter position T + MCTLASS_PRAGMA_NO_UNROLL + for (int t = 0; t < problem_size_.T; ++t) { + MCTLASS_PRAGMA_UNROLL + for (int s_idx = 0; s_idx < ThreadMap::Iterations::kStrided; ++s_idx) { + + int t_ = t; + if (problem_size_.mode == Mode::kConvolution) { + t_ = problem_size_.T - 1 - t; + } + + int d = offset_z[s_idx] * problem_size_.stride_d - problem_size_.pad_d + t_ * problem_size_.dilation_d; + + bool pred = (offset_n[s_idx] < problem_size_.N && d >= 0 && d < problem_size_.D); + masks_[s_idx][0] |= (pred << t); + } + } + + // mask predicates for filter position R + MCTLASS_PRAGMA_NO_UNROLL + for (int r = 0; r < problem_size_.R; ++r) { + MCTLASS_PRAGMA_UNROLL + for (int s_idx = 0; s_idx < ThreadMap::Iterations::kStrided; ++s_idx) { + + int r_ = r; + if (problem_size_.mode == Mode::kConvolution) { + r_ = problem_size_.R - 1 - r; + } + + int h = offset_p[s_idx] * problem_size_.stride_h - problem_size_.pad_h + r_ * problem_size_.dilation_h; + + bool pred = (h >= 0 && h < problem_size_.H); + masks_[s_idx][1] |= (pred << r); + } + } + + // mask predicates for filter position S + MCTLASS_PRAGMA_NO_UNROLL + for (int s = 0; s < problem_size_.S; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int s_idx = 0; s_idx < ThreadMap::Iterations::kStrided; ++s_idx) { + + int s_ = s; + if (problem_size_.mode == Mode::kConvolution) { + s_ = problem_size_.S - 1 - s; + } + + int w = offset_q[s_idx] * problem_size_.stride_w - problem_size_.pad_w + s_ * problem_size_.dilation_w; + + bool pred = (w >= 0 && w < problem_size_.W); + masks_[s_idx][2] |= (pred << s); + } + } + + if (filter_c_ >= problem_size.C) { + clear_mask(); + } + + set_iteration_index(0); + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv3dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}, + ThreadMap::kThreads, + ThreadMap::kElementsPerAccess, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided}); + } + +private: + + /// Returns the coordinate in the activations tensor X that is correspoinding to + // output nzpq and filter position t, r, s + MCTLASS_HOST_DEVICE + TensorCoord at_(int n, int z, int p, int q, int t, int r, int s) const { + + if (problem_size_.mode == Mode::kConvolution) { + t = problem_size_.T - 1 - t; + r = problem_size_.R - 1 - r; + s = problem_size_.S - 1 - s; + } + + int d = z * problem_size_.stride_d - problem_size_.pad_d + t * problem_size_.dilation_d; + int h = p * problem_size_.stride_h - problem_size_.pad_h + r * problem_size_.dilation_h; + int w = q * problem_size_.stride_w - problem_size_.pad_w + s * problem_size_.dilation_w; + + return TensorCoord(n, d, h, w, filter_c_); + } + + /// Adds a pointer offset in units of element + MCTLASS_HOST_DEVICE + void add_byte_offset_(LongIndex byte_offset) { + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + pointer_[s] += byte_offset; + } + } + + + /// Clears the predicates + MCTLASS_HOST_DEVICE + void clear_mask_(bool clear) { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + // We are using inline PTX assembly here to avoid an CUDA C++ compilation + // artifact in which control flow instructions are generated. Instead, our + // intent is to predicate the mov instructions. + //#if defined(__MACA_ARCH__) + #if 0 + asm volatile( + "{\n" + " .reg .pred p;\n" + " .reg .u32 m;" + " mov.u32 m, %2;" + " setp.ne.b32 p, %1, 0;\n" + " @p mov.u32 m, 0;\n" + " mov.u32 %0, m;\n" + "}\n" + : + "=r"(masks_[s][0]) + : + "r"((int)clear), + "r"(masks_[s][0]) + ); + asm volatile( + "{\n" + " .reg .pred p;\n" + " .reg .u32 m;" + " mov.u32 m, %2;" + " setp.ne.b32 p, %1, 0;\n" + " @p mov.u32 m, 0;\n" + " mov.u32 %0, m;\n" + "}\n" + : + "=r"(masks_[s][1]) + : + "r"((int)clear), + "r"(masks_[s][1]) + ); + asm volatile( + "{\n" + " .reg .pred p;\n" + " .reg .u32 m;" + " mov.u32 m, %2;" + " setp.ne.b32 p, %1, 0;\n" + " @p mov.u32 m, 0;\n" + " mov.u32 %0, m;\n" + "}\n" + : + "=r"(masks_[s][2]) + : + "r"((int)clear), + "r"(masks_[s][2]) + ); + #else + if (clear) { + masks_[s][0] = 0; + masks_[s][1] = 0; + masks_[s][2] = 0; + } + #endif + } + } + +public: + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + add_byte_offset_(pointer_offset * sizeof_bits::value / 8); + } + + MCTLASS_HOST_DEVICE + void advance() { + + int next_idx = 0; + + // moves to the next tile + ++filter_s_; + if (filter_s_ == problem_size_.S) { + + filter_s_ = 0; + ++filter_r_; + next_idx = 1; + + if (filter_r_ == problem_size_.R) { + filter_r_ = 0; + ++filter_t_; + + if (filter_t_ < problem_size_.T) { + next_idx = 2; + } + else { + filter_t_ = 0; + next_idx = 3; + } + } + } + + add_byte_offset_(params_.inc_next[next_idx]); + + if (next_idx == 3) { + filter_c_ += params_.filter_c_delta; + } + + clear_mask_(filter_c_ >= problem_size_.C); + } + + /// Clears the predicates + MCTLASS_HOST_DEVICE + void clear_mask() { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + masks_[s][0] = Mask(0); + masks_[s][1] = Mask(0); + masks_[s][2] = Mask(0); + } + } + + MCTLASS_HOST_DEVICE + bool valid() { + + return + (masks_[iteration_strided_][0] & (Index(1) << filter_t_)) && + (masks_[iteration_strided_][1] & (Index(1) << filter_r_)) && + (masks_[iteration_strided_][2] & (Index(1) << filter_s_)); + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + return reinterpret_cast(pointer_[iteration_strided_]); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv3dFpropActivationTileAccessIteratorOptimized &operator++() { + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv3dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % (128/sizeof_bits::value)) { + return Status::kErrorInvalidProblem; + } + + // Conv3dFpropActivationTileAccessIteratorOptimized has constraint on filter positions + // due to the number of mask bits. + if (problem_size.T > 32 || problem_size.R > 32 || problem_size.S > 32) { + return Status::kErrorNotSupported; + } + return Status::kSuccess; + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_analytic.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_analytic.h new file mode 100644 index 0000000..0ca5d8b --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_analytic.h @@ -0,0 +1,253 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (filter tile) + matrix from memory. + + This iterator assumes TensorNDHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/conv/threadblock/conv3d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_ +> +class Conv3dFpropFilterTileAccessIteratorAnalytic { +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNDHWC; + using ThreadMap = ThreadMap_; + using AccessType = AlignedArray; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 3; + using ConvProblemSize = typename conv::Conv3dProblemSize; + static int const kAccessesPerVector = 1; + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + using Params = Conv3dAnalyticParams; + +private: + + Params const ¶ms_; + ConvProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + char const *pointer_; + + int filter_t_; + int filter_r_; + int filter_s_; + int filter_c_; + + int offset_k_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv3dFpropFilterTileAccessIteratorAnalytic( + Params const ¶ms, + ConvProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + filter_t_(0), + filter_r_(0), + filter_s_(0), + filter_c_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_c_ = threadblock_offset.row() + thread_coord.contiguous(); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_k_[s] = threadblock_offset.column() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + } + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * 8 / sizeof_bits::value; + } + + MCTLASS_HOST_DEVICE + void advance() { + // moves to the next tile + ++filter_s_; + if (filter_s_ < problem_size_.S) { + return; + } + filter_s_ = 0; + + ++filter_r_; + if (filter_r_ < problem_size_.R) { + return; + } + filter_r_ = 0; + + ++filter_t_; + if (filter_t_ < problem_size_.T) { + return; + } + filter_t_ = 0; + + filter_c_ += Shape::kRow * problem_size_.split_k_slices; + } + + /// Returns the coordinate in the filter tensor W that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + int k = offset_k_[iteration_strided_]; + + return TensorCoord(k, filter_t_, filter_r_, filter_s_, filter_c_); + } + + /// Returns true if the current coordinate is within the activations tensor W + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + return coord.n() < problem_size_.K && + coord.c() < problem_size_.C; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv3dFpropFilterTileAccessIteratorAnalytic &operator++() { + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(ConvProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.K % (128/sizeof_bits::value)) { + return Status::kErrorInvalidProblem; + } + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + + diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_optimized.h new file mode 100644 index 0000000..102f54a --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_fprop_filter_tile_access_iterator_optimized.h @@ -0,0 +1,277 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (filter tile) + matrix from memory. + + This iterator assumes TensorNHWC or TensorCxRSKx layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv3d_problem_size.h" + +#include "mctlass/conv/threadblock/conv3d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename Layout_, + typename ThreadMap_ +> +class Conv3dFpropFilterTileAccessIteratorOptimized{ +public: + + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using ThreadMap = ThreadMap_; + using AccessType = AlignedArray; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 3; + using ConvProblemSize = typename conv::Conv3dProblemSize; + static int const kAccessesPerVector = 1; + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + + struct Params : Conv3dFpropFilterIteratorOptimizedParams { + + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Conv3dFpropFilterIteratorOptimizedParams const &base): + Conv3dFpropFilterIteratorOptimizedParams(base) { } + + MCTLASS_HOST_DEVICE + Params( + Conv3dProblemSize const &problem_size, + Layout const &layout + ): + Conv3dFpropFilterIteratorOptimizedParams( + problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}, + ThreadMap::kThreads, + ThreadMap::kElementsPerAccess, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided} + ) { + + } + }; + +private: + + Conv3dFpropFilterIteratorOptimizedParams const ¶ms_; + Conv3dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + char const *pointer_; + + uint32_t predicates_; + int filter_trs_; + int filter_c_; + + // + // Assertions + // + + // We map predicates into bits packed in this uint32_t container + static_assert(ThreadMap::Iterations::kStrided < sizeof(predicates_) * 8, + "Currently, the number of loads per iteration is limited by the size of the predicates container."); + +public: + + MCTLASS_HOST_DEVICE + Conv3dFpropFilterTileAccessIteratorOptimized( + Conv3dFpropFilterIteratorOptimizedParams const ¶ms, + Conv3dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + predicates_{0}, + filter_trs_(0), + filter_c_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_c_ = threadblock_offset.row() + thread_coord.contiguous(); + Index column = threadblock_offset.column() + thread_coord.strided(); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + uint32_t pred = ((column + s * ThreadMap::Delta::kStrided < problem_size_.K) ? 1u : 0); + predicates_ |= (pred << s); + } + + if (filter_c_ >= problem_size.C) { + predicates_ = 0u; + } + + pointer_ += ( + params_.layout({filter_c_, column}) + ) * sizeof_bits::value / 8; + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + + LongIndex next = params_.inc_next_trs; + + // moves to the next tile + ++filter_trs_; + if (filter_trs_ == params_.TRS) { + + filter_trs_ = 0; + next = params_.inc_next_c; + filter_c_ += params_.filter_c_delta; + } + + if (filter_c_ >= problem_size_.C) { + predicates_ = 0; + } + + pointer_ += next; + } + + /// Returns true if the current coordinate is within the filter tensor W + MCTLASS_HOST_DEVICE + bool valid() { + return (predicates_ & (1u << iteration_strided_)); + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + return reinterpret_cast(pointer_); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv3dFpropFilterTileAccessIteratorOptimized &operator++() { + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + + // Move to the next K coordinate within the tile + pointer_ += params_.inc_next_k; + + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv3dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % (128/sizeof_bits::value)) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_params.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_params.h new file mode 100644 index 0000000..1d6536a --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_params.h @@ -0,0 +1,508 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Extracts the host-params objects into non-template code. +*/ + +#pragma once + +#define TRACE_CONV_PARAMS_INITIALIZERS_ENABLED 0 + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/threadblock/conv2d_params.h" +#include "mctlass/conv/conv3d_problem_size.h" + +#if TRACE_CONV_PARAMS_INITIALIZERS_ENABLED +#include +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Params structure used for all Conv3d analytic tile iterators +template< typename Layout_ = layout::TensorNDHWC > +struct Conv3dAnalyticParams { + + using Layout = Layout_; + + Layout layout; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Conv3dAnalyticParams() { } + + MCTLASS_HOST_DEVICE + Conv3dAnalyticParams( + Conv3dProblemSize const &, // unused; placeholder to match other Params interfaces. + Layout const &layout + ): layout(layout) { + + } +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Parameters structure used for Conv3dFpropActivationTileIteratorOptimized +template< typename Layout_ = layout::TensorNDHWC > +struct Conv3dFpropActivationIteratorOptimizedParams; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Parameters structure used for Conv3dFpropActivationTileIteratorOptimized +template<> +struct Conv3dFpropActivationIteratorOptimizedParams { + + using Layout = layout::TensorNDHWC; + + Layout layout; + + int64_t inc_next[4]; // {next S, next R, next T, next C} + int filter_c_delta; // number of logical elements to add to filter_c_ + int ZPQ; // product of Z*P*Q + int PQ; // product of P*Q + + FastDivmod zpq_divmod; + FastDivmod pq_divmod; + FastDivmod q_divmod; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Conv3dFpropActivationIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv3dFpropActivationIteratorOptimizedParams( + Conv3dProblemSize const &problem_size, + Layout const &layout, ///< layout object + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): + layout(layout), + PQ(problem_size.P * problem_size.Q), + ZPQ(problem_size.Z * problem_size.P * problem_size.Q), + zpq_divmod(ZPQ), + pq_divmod(PQ), + q_divmod(problem_size.Q) { + + TRACE_CONV_INITIALIZERS("conv3d_fprop", "activation", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + + int conv_sign = (problem_size.mode == Mode::kConvolution ? -1 : 1); + + // next S + inc_next[0] = conv_sign * ( + int64_t(layout.stride()[0]) * problem_size.dilation_w + ) * element_size_bits / 8; + + // next R + inc_next[1] = conv_sign * ( + int64_t(layout.stride()[1]) * problem_size.dilation_h + - (problem_size.S - 1) * layout.stride()[0] * problem_size.dilation_w + ) * element_size_bits / 8; + + // next T + inc_next[2] = conv_sign * ( + int64_t(layout.stride()[2]) * problem_size.dilation_d + - (problem_size.R - 1) * layout.stride()[1] * problem_size.dilation_h + - (problem_size.S - 1) * layout.stride()[0] * problem_size.dilation_w + ) * element_size_bits / 8; + + // next C + inc_next[3] = ( + threadblock_shape.column() * problem_size.split_k_slices + - conv_sign * int64_t(problem_size.T - 1) * layout.stride()[2] * problem_size.dilation_d + - conv_sign * int64_t(problem_size.R - 1) * layout.stride()[1] * problem_size.dilation_h + - conv_sign * int64_t(problem_size.S - 1) * layout.stride()[0] * problem_size.dilation_w + ) * element_size_bits / 8; + + // logical offset added to internal channel counter - units are elements, not bytes + filter_c_delta = threadblock_shape.column() * problem_size.split_k_slices; + } +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + + +template< typename Layout_ = layout::TensorNDHWC > +struct Conv3dFpropFilterIteratorOptimizedParams; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template<> +struct Conv3dFpropFilterIteratorOptimizedParams +{ + + using Layout = layout::TensorNDHWC; + + Layout layout; + int TRS; + int filter_c_delta; + + int64_t inc_next_k; // offset in units of bytes to next K position + int64_t inc_next_trs; // offset in units of bytes to next TRS position + int64_t inc_next_c; // offset in units of bytes to next C position + + // + // Methods + // + MCTLASS_HOST_DEVICE + Conv3dFpropFilterIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv3dFpropFilterIteratorOptimizedParams( + Conv3dProblemSize const &problem_size, + Layout const &layout, + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): + layout(layout) { + + TRACE_CONV_INITIALIZERS("conv3d_fprop", "filter", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + TRS = problem_size.T * problem_size.R * problem_size.S; + + inc_next_k = (int64_t(layout.stride()[3]) * threadmap_delta.strided() * element_size_bits) / 8; + + inc_next_trs = + ( int64_t(layout.stride()[0]) + - int64_t(layout.stride()[3]) * (threadmap_iterations.strided() - 1) * threadmap_delta.strided() + ) * element_size_bits / 8; + + inc_next_c = + ( + threadblock_shape.row() * problem_size.split_k_slices + - int64_t(TRS - 1) * layout.stride()[0] + - int64_t(threadmap_iterations.strided() - 1) * threadmap_delta.strided() * layout.stride()[3] + ) * element_size_bits / 8; + + filter_c_delta = threadblock_shape.row() * problem_size.split_k_slices; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Parameters object for Conv3d DGRAD OutputGradient (dy) iterator +struct Conv3dDgradOutputGradientIteratorOptimizedParams { + + using Layout = layout::TensorNDHWC; + + Layout layout; + + int64_t inc_next[4]; // {next S, next R, next T, next K} + int filter_k_delta; // number of logical elements to add to filter_k_ + + FastDivmod dhw_divmod; + FastDivmod hw_divmod; + FastDivmod w_divmod; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Conv3dDgradOutputGradientIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv3dDgradOutputGradientIteratorOptimizedParams( + Conv3dProblemSize const &problem_size, + Layout const &layout, ///< layout object + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): + layout(layout), + dhw_divmod(problem_size.D * problem_size.H * problem_size.W), + hw_divmod(problem_size.H * problem_size.W), + w_divmod(problem_size.W) { + + TRACE_CONV_INITIALIZERS("conv3d_dgrad", "output_gradient", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + int conv_sign = (problem_size.mode == Mode::kConvolution ? 1 : -1); + + // next S + inc_next[0] = conv_sign * ( + int64_t(layout.stride()[0]) * problem_size.dilation_w + ) * element_size_bits / 8; + + // next R + inc_next[1] = conv_sign * ( + int64_t(layout.stride()[1]) * problem_size.dilation_h + - (problem_size.S - 1) * layout.stride()[0] * problem_size.dilation_w + ) * element_size_bits / 8; + + // next T + inc_next[2] = conv_sign * ( + int64_t(layout.stride()[2]) * problem_size.dilation_d + - (problem_size.R - 1) * layout.stride()[1] * problem_size.dilation_h + - (problem_size.S - 1) * layout.stride()[0] * problem_size.dilation_w + ) * element_size_bits / 8; + + // next K + inc_next[3] = ( + threadblock_shape.column() * problem_size.split_k_slices + - conv_sign * int64_t(problem_size.T - 1) * layout.stride()[2] * problem_size.dilation_d + - conv_sign * int64_t(problem_size.R - 1) * layout.stride()[1] * problem_size.dilation_h + - conv_sign * int64_t(problem_size.S - 1) * layout.stride()[0] * problem_size.dilation_w + ) * element_size_bits / 8; + + // logical offset added to internal channel counter - units are elements, not bytes + filter_k_delta = threadblock_shape.column() * problem_size.split_k_slices; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Parameters object for Conv2d DGRAD Filter (w) iterator +struct Conv3dDgradFilterIteratorOptimizedParams { + + using Layout = layout::TensorNDHWC; + + Layout layout; + int TRS; + int filter_k_delta; + + int64_t inc_next_strided; // offset in units of bytes to next K coordinate within tile + int64_t inc_next_trs; // offset in units of bytes to next TRS position + int64_t inc_next_k; // offset in units of bytes to next K position in subsequent tile + + // + // Methods + // + MCTLASS_HOST_DEVICE + Conv3dDgradFilterIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv3dDgradFilterIteratorOptimizedParams( + Conv3dProblemSize const &problem_size, + Layout const &layout, + int element_size_bits, ///< size of each element in bits + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): + layout(layout), TRS(problem_size.T * problem_size.R * problem_size.S) { + + TRACE_CONV_INITIALIZERS("conv3d_dgrad", "filter", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + inc_next_strided = ((int64_t)layout.stride()[3] * threadmap_delta.strided() * element_size_bits) / 8; + + inc_next_trs = + ( (int64_t)layout.stride()[0] + - (threadmap_iterations.strided() - 1) * threadmap_delta.strided() * (int64_t)layout.stride()[3] + ) * element_size_bits / 8; + + inc_next_k = + ( + threadblock_shape.row() * problem_size.split_k_slices * (int64_t)layout.stride()[3] + - (problem_size.T * problem_size.R * problem_size.S - 1) * (int64_t)layout.stride()[0] + - (threadmap_iterations.strided() - 1) * threadmap_delta.strided() * (int64_t)layout.stride()[3] + ) * element_size_bits / 8; + + filter_k_delta = threadblock_shape.row() * problem_size.split_k_slices; + } +}; + +/// Parameters object for Conv3d WGRAD OutputGradient iterator +struct Conv3dWgradOutputGradientIteratorOptimizedParams { + + using Layout = layout::TensorNDHWC; + using LongIndex = typename Layout::LongIndex; + + Layout layout; + + int NZPQ; // precomputd product of N*Z*P*Q for clearing predicates + int ZPQ; // product of Z*P*Q + unsigned zpq_mul; // precomputed quantities for fast computation of div/% by ZPQ + unsigned zpq_shr; // in device code. + + int PQ; // product of P*Q + unsigned pq_mul; // precomputed quantities for fast computation of div/% by PQ + unsigned pq_shr; // in device code. + + unsigned q_mul; // precomputed quantities for fast computation of div/% by Q + unsigned q_shr; // in device code. + + LongIndex offset_next_strided; // offset in units of bytes to next nzpq coordinate within tile + LongIndex offset_next_contiguous; // offset in units of bytes to next k coordinate within tile + LongIndex inc_next_nzpq; // offset in units of bytes to next nzpq position in subsequent tile + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Conv3dWgradOutputGradientIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv3dWgradOutputGradientIteratorOptimizedParams( + Conv3dProblemSize const &problem_size, + Layout const &layout, + int element_size_bits, + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): layout(layout) { + + TRACE_CONV_INITIALIZERS("conv3d_wgrad", "output_gradient", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + // Incremental offsets in unites of bytes (number of elements) * element_size_bits / 8 + offset_next_strided = (threadmap_delta.strided() * (int64_t)layout.stride()[0]) + * element_size_bits / 8; + + offset_next_contiguous = (threadmap_delta.contiguous()) + * element_size_bits / 8; + + inc_next_nzpq = (threadblock_shape.column() * problem_size.split_k_slices * (int64_t)layout.stride()[0]) + * element_size_bits / 8; + + // Precompute several quantities for fast modulo arithmetic. + NZPQ = problem_size.N * problem_size.Z * problem_size.P * problem_size.Q; + ZPQ = problem_size.Z * problem_size.P * problem_size.Q; + find_divisor(zpq_mul, zpq_shr, ZPQ); + + PQ = problem_size.P * problem_size.Q; + find_divisor(pq_mul, pq_shr, PQ); + + find_divisor(q_mul, q_shr, problem_size.Q); + + } +}; + +/// Parameters object for Conv3d WGRAD Activation Tile Access Iterator +struct Conv3dWgradActivationIteratorOptimizedParams { + + using Layout = layout::TensorNDHWC; + + Layout layout; + + int RSC; // product of R*S*C + unsigned rsc_mul; // precomputed quantities for fast computation of div/% by RSC + unsigned rsc_shr; // in device code. + + int SC; // product of S*C + unsigned sc_mul; // precomputed quantities for fast computation of div/% by SC + unsigned sc_shr; // in device code. + + unsigned c_mul; // precomputed quantities for fast computation of div/% by C + unsigned c_shr; // in device code. + + int ZPQ; // product of Z*P*Q + unsigned zpq_mul; // precomputed quantities for fast computation of div/% by ZPQ + unsigned zpq_shr; // in device code. + + int PQ; // product of P*Q + unsigned pq_mul; // precomputed quantities for fast computation of div/% by PQ + unsigned pq_shr; // in device code. + + unsigned q_mul; // precomputed quantities for fast computation of div/% by Q + unsigned q_shr; // in device code. + + // + // Methods + // + MCTLASS_HOST_DEVICE + Conv3dWgradActivationIteratorOptimizedParams() { } + + MCTLASS_HOST_DEVICE + Conv3dWgradActivationIteratorOptimizedParams( + Conv3dProblemSize const &problem_size, + Layout const &layout, + int element_size_bits, + MatrixCoord threadblock_shape, + int thread_count, + int access_size, + layout::PitchLinearCoord threadmap_iterations, + layout::PitchLinearCoord threadmap_delta + ): layout(layout) { + + TRACE_CONV_INITIALIZERS("conv3d_wgrad", "activation", + element_size_bits, threadblock_shape, thread_count, access_size, threadmap_iterations, threadmap_delta); + + // Precompute several quantities for fast modulo arithmetic. + RSC = problem_size.R * problem_size.S * problem_size.C; + find_divisor(rsc_mul, rsc_shr, RSC); + + SC = problem_size.S * problem_size.C; + find_divisor(sc_mul, sc_shr, SC); + + find_divisor(c_mul, c_shr, problem_size.C); + + ZPQ = problem_size.Z * problem_size.P * problem_size.Q; + find_divisor(zpq_mul, zpq_shr, ZPQ); + + PQ = problem_size.P * problem_size.Q; + find_divisor(pq_mul, pq_shr, PQ); + + find_divisor(q_mul, q_shr, problem_size.Q); + + } +}; + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_analytic.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_analytic.h new file mode 100644 index 0000000..cd3ffce --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_analytic.h @@ -0,0 +1,289 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (activation tile) + matrix from memory. + + This iterator assumes TensorNDHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv3d_problem_size.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_ +> +class Conv3dWgradActivationTileAccessIteratorAnalytic { +public: + + // + // Types + // + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNDHWC; + using ThreadMap = ThreadMap_; + using AccessType = AlignedArray; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 3; + using ConvProblemSize = typename conv::Conv3dProblemSize; + + static int const kAccessesPerVector = 1; + + static_assert(sizeof_bits::value >= 8, + "WGRAD requires elements of size 8b or greater."); + + // + // Parameters structure + // + + struct Params { + + Layout layout; + + // + // Methods + // + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params( + Conv3dProblemSize const &problem_size, + Layout const &layout + ): layout(layout) { + + } + }; + +private: + + Params const ¶ms_; + Conv3dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + char const *pointer_; + + // Filter postion (t,r,s,c) in contiguous dimension stays constant for each gemm_iteration_k + int filter_t_[ThreadMap::Iterations::kContiguous]; + int filter_r_[ThreadMap::Iterations::kContiguous]; + int filter_s_[ThreadMap::Iterations::kContiguous]; + int filter_c_[ThreadMap::Iterations::kContiguous]; + + int offset_nzpq_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv3dWgradActivationTileAccessIteratorAnalytic( + Params const ¶ms, + Conv3dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + // initialize t,r,s,c filter position for every contiguous iteration + MCTLASS_PRAGMA_UNROLL + for(int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int trsc_offset = threadblock_offset.column() + thread_coord.contiguous() + + c * ThreadMap::Delta::kContiguous; + + filter_t_[c] = trsc_offset / (problem_size_.R * problem_size_.S * problem_size_.C); + int residual = trsc_offset % (problem_size_.R * problem_size_.S * problem_size_.C); + + filter_r_[c] = residual / (problem_size_.S * problem_size_.C); + residual = residual % (problem_size_.S * problem_size_.C); + + filter_s_[c] = residual / problem_size_.C; + filter_c_[c] = residual % problem_size_.C; + + } + + // initialize n, z, p, q offset for every strided iteration + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + offset_nzpq_[s] = threadblock_offset.row() + thread_coord.strided() + + s * ThreadMap::Delta::kStrided; + } + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + + // moves to the next GEMM-K offset (offset_nzpq_) in GEMM-B by a CTA-K tile + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_nzpq_[s] += Shape::kRow * problem_size_.split_k_slices; + } + } + + /// Returns the coordinate in the activation tensor x that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + int t = filter_t_[iteration_contiguous_]; + int r = filter_r_[iteration_contiguous_]; + int s = filter_s_[iteration_contiguous_]; + + if (problem_size_.mode == Mode::kConvolution) { + t = (problem_size_.T - 1 - t); + r = (problem_size_.R - 1 - r); + s = (problem_size_.S - 1 - s); + } + + int n = offset_nzpq_[iteration_strided_] / (problem_size_.Z * problem_size_.P * problem_size_.Q); + int residual = offset_nzpq_[iteration_strided_] % (problem_size_.Z * problem_size_.P * problem_size_.Q); + + int z = residual / (problem_size_.P * problem_size_.Q); + residual = residual % (problem_size_.P * problem_size_.Q); + + int p = residual / problem_size_.Q; + int q = residual % problem_size_.Q; + + int d = z * problem_size_.stride_d - problem_size_.pad_d + t * problem_size_.dilation_d; + int h = p * problem_size_.stride_h - problem_size_.pad_h + r * problem_size_.dilation_h; + int w = q * problem_size_.stride_w - problem_size_.pad_w + s * problem_size_.dilation_w; + + return TensorCoord(n, d, h, w, filter_c_[iteration_contiguous_]); + } + + /// Returns true if the current coordinate is within the activation tensor x + MCTLASS_HOST_DEVICE + bool valid() const { + TensorCoord coord = at(); + + return coord.n() < problem_size_.N && + coord.d() >= 0 && coord.d() < problem_size_.D && + coord.h() >= 0 && coord.h() < problem_size_.H && + coord.w() >= 0 && coord.w() < problem_size_.W && + coord.c() < problem_size_.C; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv3dWgradActivationTileAccessIteratorAnalytic &operator++() { + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv3dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % (128/sizeof_bits::value)) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } + +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + + diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_optimized.h new file mode 100644 index 0000000..f36d5b5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_activation_tile_access_iterator_optimized.h @@ -0,0 +1,319 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (activation tile) + matrix from memory. + + This iterator assumes TensorNDHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/conv/threadblock/conv3d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_ +> +class Conv3dWgradActivationTileAccessIteratorOptimized { +public: + + // + // Types + // + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNDHWC; + using ThreadMap = ThreadMap_; + using AccessType = AlignedArray; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 3; + using ConvProblemSize = typename conv::Conv3dProblemSize; + static int const kAccessesPerVector = 1; + static_assert(sizeof_bits::value >= 8, + "WGRAD requires elements of size 8b or greater."); + + // + // Parameters structure + // + + struct Params : Conv3dWgradActivationIteratorOptimizedParams { + // + // Methods + // + MCTLASS_HOST_DEVICE + Params() {} + + MCTLASS_HOST_DEVICE + Params(Conv3dWgradActivationIteratorOptimizedParams const &base) + : Conv3dWgradActivationIteratorOptimizedParams(base) {} + + MCTLASS_HOST_DEVICE + Params(Conv3dProblemSize const &problem_size, Layout const &layout) + : Conv3dWgradActivationIteratorOptimizedParams( + problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}, + ThreadMap::kThreads, + ThreadMap::kElementsPerAccess, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided}) {} + }; + +private: + + Params const ¶ms_; + Conv3dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + char const *pointer_; + + // Precomputed effective filter postion (t,r,s) in contiguous dimension stays constant for each gemm_iteration_k + // required for nzpq -> ndhw translation + int precomputed_filter_t_[ThreadMap::Iterations::kContiguous]; + int precomputed_filter_r_[ThreadMap::Iterations::kContiguous]; + int precomputed_filter_s_[ThreadMap::Iterations::kContiguous]; + + // Channel dimension in contiguous dimension stays constant for each gemm_iteration_k + int filter_c_[ThreadMap::Iterations::kContiguous]; + + int offset_nzpq_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv3dWgradActivationTileAccessIteratorOptimized( + Params const ¶ms, + Conv3dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + // initialize t,r,s,c filter position for every contiguous iteration + MCTLASS_PRAGMA_UNROLL + for(int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int trsc_offset = threadblock_offset.column() + thread_coord.contiguous() + + c * ThreadMap::Delta::kContiguous; + + // The subseqnet fast_divmod() operations are equivalent to the following logical computation: + // + // + // filter_t_[c] = trsc_offset / (problem_size_.R * problem_size_.S * problem_size_.C); + // int residual = trsc_offset % (problem_size_.R * problem_size_.S * problem_size_.C); + // + // filter_r_[c] = residual / (problem_size_.S * problem_size_.C); + // residual = residual % (problem_size_.S * problem_size_.C); + // + // filter_s_[c] = residual / problem_size_.C; + // filter_c_[c] = residual % problem_size_.C; + + int residual; + fast_divmod(precomputed_filter_t_[c], residual, trsc_offset, params_.RSC, params_.rsc_mul, params_.rsc_shr); + fast_divmod(precomputed_filter_r_[c], residual, residual, params_.SC, params_.sc_mul, params_.sc_shr); + fast_divmod(precomputed_filter_s_[c], filter_c_[c], residual, problem_size_.C, params_.c_mul, params_.c_shr); + + int t = precomputed_filter_t_[c]; + int r = precomputed_filter_r_[c]; + int s = precomputed_filter_s_[c]; + + if (problem_size_.mode == Mode::kConvolution) { + t = (problem_size_.T - 1 - t); + r = (problem_size_.R - 1 - r); + s = (problem_size_.S - 1 - s); + } + + // efective t,r,s for every contiguous dimension + precomputed_filter_t_[c] = - problem_size_.pad_d + t * problem_size_.dilation_d; + precomputed_filter_r_[c] = - problem_size_.pad_h + r * problem_size_.dilation_h; + precomputed_filter_s_[c] = - problem_size_.pad_w + s * problem_size_.dilation_w; + + + } + + // initialize n, z, p, q offset for every strided iteration + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + offset_nzpq_[s] = threadblock_offset.row() + thread_coord.strided() + + s * ThreadMap::Delta::kStrided; + } + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + + // moves to the next GEMM-K offset (offset_nzpq_) in GEMM-B by a CTA-K tile + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_nzpq_[s] += Shape::kRow * problem_size_.split_k_slices; + } + } + + /// Returns the coordinate in the activation tensor x that is currently pointed to + /// by the iterator. + + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + // The subseqnet fast_divmod() operations are equivalent to the following logical computation: + // + // + // int n = offset_nzpq_[iteration_strided_] / (problem_size_.Z * problem_size_.P * problem_size_.Q); + // int residual = offset_nzpq_[iteration_strided_] % (problem_size_.Z * problem_size_.P * problem_size_.Q); + // + // int z = residual / (problem_size_.P * problem_size_.Q); + // residual = residual % (problem_size_.P * problem_size_.Q); + // + // int p = residual / problem_size_.Q; + // int q = residual % problem_size_.Q; + + int residual, n, z, p, q; + fast_divmod(n, residual, offset_nzpq_[iteration_strided_], params_.ZPQ, params_.zpq_mul, params_.zpq_shr); + fast_divmod(z, residual, residual, params_.PQ, params_.pq_mul, params_.pq_shr); + fast_divmod(p, q, residual, problem_size_.Q, params_.q_mul, params_.q_shr); + + int d = z * problem_size_.stride_d + precomputed_filter_t_[iteration_contiguous_]; + int h = p * problem_size_.stride_h + precomputed_filter_r_[iteration_contiguous_];; + int w = q * problem_size_.stride_w + precomputed_filter_s_[iteration_contiguous_]; + + return TensorCoord(n, d, h, w, filter_c_[iteration_contiguous_]); + } + + /// Returns true if the current coordinate is within the activation tensor x + MCTLASS_HOST_DEVICE + bool valid() const { + TensorCoord coord = at(); + + return coord.n() < problem_size_.N && + coord.d() >= 0 && coord.d() < problem_size_.D && + coord.h() >= 0 && coord.h() < problem_size_.H && + coord.w() >= 0 && coord.w() < problem_size_.W && + coord.c() < problem_size_.C; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv3dWgradActivationTileAccessIteratorOptimized &operator++() { + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv3dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % (128/sizeof_bits::value)) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } + +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + + diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_analytic.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_analytic.h new file mode 100644 index 0000000..a9040f2 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_analytic.h @@ -0,0 +1,267 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (output gradient tile) + matrix from memory. + + This iterator assumes TensorNDHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv3d_problem_size.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_ +> +class Conv3dWgradOutputGradientTileAccessIteratorAnalytic { +public: + + // + // Types + // + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNDHWC; + using ThreadMap = ThreadMap_; + using AccessType = AlignedArray; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kAnalytic; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 3; + using ConvProblemSize = typename conv::Conv3dProblemSize; + static int const kAccessesPerVector = 1; + static_assert(sizeof_bits::value >= 8, + "WGRAD requires elements of size 8b or greater."); + + // + // Parameters structure + // + + struct Params { + + Layout layout; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params( + Conv3dProblemSize const &problem_size, + Layout const &layout + ): layout(layout) { + + } + }; + +private: + + Params const ¶ms_; + Conv3dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + char const *pointer_; + + int filter_k_[ThreadMap::Iterations::kContiguous]; + + int offset_nzpq_[ThreadMap::Iterations::kStrided]; + +public: + + MCTLASS_HOST_DEVICE + Conv3dWgradOutputGradientTileAccessIteratorAnalytic( + Params const ¶ms, + Conv3dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)) { + + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + // initialize filter_k for every contiguous iteration + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + filter_k_[c] = threadblock_offset.row() + thread_coord.contiguous() + + c * ThreadMap::Delta::kContiguous; + } + + // initialize n, p, q offset for every strided iteration + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_nzpq_[s] = threadblock_offset.column() + thread_coord.strided() + + s * ThreadMap::Delta::kStrided; + + } + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv3dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, layout); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + // moves to the next GEMM-K offset (offset_nzpq_) in GEMM-A by a CTA-K tile + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_nzpq_[s] += Shape::kColumn * problem_size_.split_k_slices; + } + } + + /// Returns the coordinate in the output gradient tensor Dy that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + int nzpq = offset_nzpq_[iteration_strided_]; + + int n = nzpq / (problem_size_.Z * problem_size_.P * problem_size_.Q); + int residual = nzpq % (problem_size_.Z * problem_size_.P * problem_size_.Q); + + int z = residual / (problem_size_.P * problem_size_.Q); + residual = residual % (problem_size_.P * problem_size_.Q); + + int p = residual / problem_size_.Q; + int q = residual % problem_size_.Q; + + return TensorCoord(n, z, p, q, filter_k_[iteration_contiguous_]); + } + + + /// Returns true if the current coordinate is within the output gradient tensor Dy + MCTLASS_HOST_DEVICE + bool valid() const { + TensorCoord coord = at(); + + return coord.n() < problem_size_.N && + coord.d() < problem_size_.Z && + coord.h() < problem_size_.P && + coord.w() < problem_size_.Q && + coord.c() < problem_size_.K; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + return reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv3dWgradOutputGradientTileAccessIteratorAnalytic &operator++() { + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv3dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.K % (128/sizeof_bits::value)) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } + +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + + diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_optimized.h new file mode 100644 index 0000000..3dcb897 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/conv3d_wgrad_output_gradient_tile_access_iterator_optimized.h @@ -0,0 +1,310 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (output gradient tile) + matrix from memory. + + This iterator assumes TensorNDHWC layout of tensors in Global Memory. + + The iterator is specialized for each of the three convolution operators: forward propagation (Fprop), + backward data gradient (Dgrad), and backward weight gradient (Wgrad). +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/conv/threadblock/conv3d_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, + typename Element_, + typename ThreadMap_ +> +class Conv3dWgradOutputGradientTileAccessIteratorOptimized { +public: + + // + // Types + // + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorNDHWC; + using ThreadMap = ThreadMap_; + using AccessType = AlignedArray; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 3; + using ConvProblemSize = typename conv::Conv3dProblemSize; + static int const kAccessesPerVector = 1; + static_assert(sizeof_bits::value >= 8, + "WGRAD requires elements of size 8b or greater."); + + // + // Parameters structure + // + + struct Params : Conv3dWgradOutputGradientIteratorOptimizedParams { + // + // Methods + // + MCTLASS_HOST_DEVICE + Params() {} + + MCTLASS_HOST_DEVICE + Params(Conv3dWgradOutputGradientIteratorOptimizedParams const &base) + : Conv3dWgradOutputGradientIteratorOptimizedParams(base) {} + + MCTLASS_HOST_DEVICE + Params(Conv3dProblemSize const &problem_size, Layout const &layout) + : Conv3dWgradOutputGradientIteratorOptimizedParams( + problem_size, + layout, + sizeof_bits::value, + {Shape::kRow, Shape::kColumn}, + ThreadMap::kThreads, + ThreadMap::kElementsPerAccess, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided}) {} + }; + +private: + + Params const ¶ms_; + Conv3dProblemSize const &problem_size_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + char const *pointer_; + + uint32_t predicates_; + int filter_k_; + int offset_nzpq_; + +public: + + MCTLASS_HOST_DEVICE + Conv3dWgradOutputGradientTileAccessIteratorOptimized( + Params const ¶ms, + Conv3dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + predicates_(0), + filter_k_(0), + offset_nzpq_(0) { + + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_k_ = threadblock_offset.row() + thread_coord.contiguous(); + offset_nzpq_ = threadblock_offset.column() + thread_coord.strided(); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int filter_k = filter_k_ + c * ThreadMap::Delta::kContiguous; + int offset_nzpq = offset_nzpq_ + s * ThreadMap::Delta::kStrided; + + bool predicate = valid_(at_(offset_nzpq, filter_k)); + + uint32_t pred = (predicate ? 1u : 0); + + int pred_idx = c + s * ThreadMap::Iterations::kContiguous; + + predicates_ |= (pred << pred_idx); + } + } + + // Offset pointer to (iteration_strided_, iteration_contiguous_) = (0, 0) + pointer_ += ( + offset_nzpq_ * params.layout.stride()[0] + filter_k_ + ) * sizeof_bits::value / 8; + + set_iteration_index(0); + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv3dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, layout); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + // moves to the next GEMM-K offset (offset_npq_) in GEMM-A by a CTA-K tile + offset_nzpq_ += Shape::kColumn * problem_size_.split_k_slices; + + // Clear predicates if needed + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + if (offset_nzpq_ + s * ThreadMap::Delta::kStrided >= params_.NZPQ) { + uint32_t kClearMask = ((1u << ThreadMap::Iterations::kContiguous) - 1) << (s * ThreadMap::Iterations::kContiguous); + predicates_ = (predicates_ & (~kClearMask)); + } + } + pointer_ += params_.inc_next_nzpq; + } + +private: + /// Returns the coordinate in the output gradient tensor Dy that is (offset_nzpq, k) pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at_(int offset_nzpq, int k) const { + + // The subseqnet fast_divmod() operations are equivalent to the following logical computation: + // + // + // int nzpq = offset_nzpq_; + // int n = nzpq / (problem_size_.Z * problem_size_.P * problem_size_.Q); + // int residual = nzpq % (problem_size_.Z * problem_size_.P * problem_size_.Q); + // + // int z = residual / (problem_size_.P * problem_size_.Q); + // residual = residual % (problem_size_.P * problem_size_.Q); + // + // int p = residual / problem_size_.Q; + // int q = residual % problem_size_.Q; + + int residual, n, z, p, q; + fast_divmod(n, residual, offset_nzpq, params_.ZPQ, params_.zpq_mul, params_.zpq_shr); + fast_divmod(z, residual, residual, params_.PQ, params_.pq_mul, params_.pq_shr); + fast_divmod(p, q, residual, problem_size_.Q, params_.q_mul, params_.q_shr); + + return TensorCoord(n, z, p, q, k); + } + + /// Returns true if the coord is within the output gradient tensor Dy + MCTLASS_HOST_DEVICE + bool valid_(TensorCoord coord) const { + + return coord.n() < problem_size_.N && + coord.c() < problem_size_.K; + } + +public: + + /// Returns true if the current coordinate is within the output gradient tensor Dy + MCTLASS_HOST_DEVICE + bool valid() const { + + LongIndex pred_idx = iteration_contiguous_ + iteration_strided_ * ThreadMap::Iterations::kContiguous; + return (predicates_ & (1u << pred_idx)); + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + + return reinterpret_cast( + pointer_ + + iteration_strided_ * params_.offset_next_strided + + iteration_contiguous_ * params_.offset_next_contiguous + ); + + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + Conv3dWgradOutputGradientTileAccessIteratorOptimized &operator++() { + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv3dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.K % (128/sizeof_bits::value)) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } + +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + + diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_direct_conv_params.h b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_direct_conv_params.h new file mode 100644 index 0000000..a176086 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_direct_conv_params.h @@ -0,0 +1,230 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Extracts the host-params objects into non-template code. +*/ + +#pragma once + +#define TRACE_CONV_PARAMS_INITIALIZERS_ENABLED 0 + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" + +#if TRACE_CONV_PARAMS_INITIALIZERS_ENABLED +#include +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Parameters structure used for DepthwiseFpropActivationDirect2dConvTileAccessIteratorOptimized +template +struct Depthwise2dFpropDirectConvParams; + +/// Parameters structure used for DepthwiseFpropActivationDirect2dConvTileAccessIteratorFixedStrideDilation +template +struct Depthwise2dFpropDirectConvActivationIteratorFixedStrideDilationParams; + +/// Parameters structure used for DepthwiseFpropFilterDirectConvTileAccessIteratorOptimized +template +struct Depthwise2dFpropDirectConvFilterIteratorParams; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Parameters structure used for DepthwiseFpropActivationDirect2dConvTileAccessIteratorOptimized +template<> +struct Depthwise2dFpropDirectConvParams { + + using Layout = layout::TensorNHWC; + + Layout layout; + + int32_t activation_tile_h; + int32_t activation_tile_w; + int32_t activation_tile_hw; + FastDivmod activation_tile_w_divmod; + + int filter[2]; + int stride[2]; + int dilation[2]; + int inc_next[2]; + FastDivmod pq_divmod; + FastDivmod q_divmod; + + int activation_load_count; + int activation_storage_elements; + int activation_size; + // + // Methods + // + + MCTLASS_HOST_DEVICE + Depthwise2dFpropDirectConvParams() { } + + MCTLASS_HOST_DEVICE + Depthwise2dFpropDirectConvParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, ///< layout object + MatrixCoord threadblock_shape, ///< CTA threadblock Shape + Layout::TensorCoord threadblock_output_shape, ///< Output tile Shape per threadblock + const int element_size_bits, ///< bits of activation element + const int thread_count, ///< threads per threadblock + const int thread_count_contiguous, ///< number of threads for continuous dimension + const int element_per_load) ///< element per each load + : layout(layout) { + + filter[0] = problem_size.S; + filter[1] = problem_size.R; + + stride[0] = problem_size.stride_w; + stride[1] = problem_size.stride_h; + + dilation[0] = problem_size.dilation_w; + dilation[1] = problem_size.dilation_h; + + // Compute activation_tile size per threadblock because stride and dilation are runtime params. + activation_tile_h = (threadblock_output_shape.h() - 1) * problem_size.stride_h + + (problem_size.R - 1) * problem_size.dilation_h + 1; + activation_tile_w = (threadblock_output_shape.w() - 1) * problem_size.stride_w + + (problem_size.S - 1) * problem_size.dilation_w + 1; + activation_tile_hw = activation_tile_h * activation_tile_w; + + activation_tile_w_divmod = FastDivmod(activation_tile_w); + + /// Below two values could not be templatized because the stride and dilation are runtime params + activation_load_count = (thread_count_contiguous * activation_tile_hw + (thread_count - 1)) / thread_count; + activation_storage_elements = activation_load_count * element_per_load * thread_count; + activation_size = activation_storage_elements * element_size_bits / 8; + + // Fastdivmod for output P, Q + int tiles_p = + (problem_size.P + (threadblock_output_shape.h() - 1)) / (threadblock_output_shape.h()); + int tiles_q = (problem_size.Q + (threadblock_output_shape.w() - 1)) / + (threadblock_output_shape.w()); + + pq_divmod = FastDivmod(tiles_p * tiles_q); + q_divmod = FastDivmod(tiles_q); + + // next S + inc_next[0] = problem_size.dilation_w; + // next R + inc_next[1] = (activation_tile_w * problem_size.dilation_h - (problem_size.S - 1) * problem_size.dilation_w); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Parameters structure used for DepthwiseFpropActivationDirect2dConvTileAccessIteratorFixedStrideDilation +template <> +struct Depthwise2dFpropDirectConvActivationIteratorFixedStrideDilationParams { + using Layout = layout::TensorNHWC; + + Layout layout; + + FastDivmod pq_divmod; + FastDivmod q_divmod; + + int activation_size; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Depthwise2dFpropDirectConvActivationIteratorFixedStrideDilationParams() {} + + MCTLASS_HOST_DEVICE + Depthwise2dFpropDirectConvActivationIteratorFixedStrideDilationParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, ///< Layout object + MatrixCoord threadblock_shape, ///< Threadblock Shape + Layout::TensorCoord threadblock_output_shape, ///< Output tile Shape per threadblock + const int activation_size_ ///< Activation size loaded by iterator + ) + : layout(layout), + activation_size(activation_size_) { + // Fastdivmod for output P, Q + int tiles_p = + (problem_size.P + (threadblock_output_shape.h() - 1)) / (threadblock_output_shape.h()); + int tiles_q = + (problem_size.Q + (threadblock_output_shape.w() - 1)) / (threadblock_output_shape.w()); + + pq_divmod = FastDivmod(tiles_p * tiles_q); + q_divmod = FastDivmod(tiles_q); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Parameters structure used for DepthwiseFpropFilterDirectConvTileAccessIteratorOptimized +template <> +struct Depthwise2dFpropDirectConvFilterIteratorParams { + using Layout = layout::TensorNHWC; + + Layout layout; + + int filter_size; + + bool is_convolution; + // + // Methods + // + + MCTLASS_HOST_DEVICE + Depthwise2dFpropDirectConvFilterIteratorParams() {} + + MCTLASS_HOST_DEVICE + Depthwise2dFpropDirectConvFilterIteratorParams( + Conv2dProblemSize const &problem_size, + Layout const &layout, ///< Layout object + MatrixCoord threadblock_shape, ///< Threadblock Shape + const int filter_size_) ///< Filter size loaded by iterator + : layout(layout), + filter_size(filter_size_), + is_convolution(problem_size.mode == Mode::kConvolution){} +}; + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_activation_tile_access_iterator_direct_conv_fixed_stride_dilation.h b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_activation_tile_access_iterator_direct_conv_fixed_stride_dilation.h new file mode 100644 index 0000000..4922a84 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_activation_tile_access_iterator_direct_conv_fixed_stride_dilation.h @@ -0,0 +1,314 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (activation tile) + matrix from memory. + + This iterator assumes TensorNHWC layout of tensors in Global Memory. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/threadblock/depthwise_direct_conv_params.h" +#include "mctlass/coord.h" +#include "mctlass/mctlass.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template > +class DepthwiseFpropActivationDirect2dConvTileAccessIteratorFixedStrideDilation { + public: + // + // Types + // + + using Shape = Shape_; + using OutputTileShape = OutputTileShape_; + using Element = Element_; + using Layout = Layout_; + using TensorCoord = typename Layout::TensorCoord; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + // Compilation value of stride , dialtion and activation shape + using StrideShape = StrideShape_; + using DilationShape = DilationShape_; + using ActivationShape = ActivationShape_; + + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + static int const kActivationSize = ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess * ThreadMap::kThreads * + sizeof_bits::value / 8; + + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, "Require Iterations::kContiguous == 1"); + + static_assert(OutputTileShape::kN == 1, "Require OutputTileShape::kN == 1"); + static_assert(OutputTileShape::kC == Shape::kColumn, "Require OutputTile shape == channels per threadblock"); + + // + // Parameters structure + // + + using Params = Depthwise2dFpropDirectConvActivationIteratorFixedStrideDilationParams; + + private: + Conv2dProblemSize const &problem_size_; + Params const ¶ms_; + char const *pointer_; + + // Base channels for current threadblock + int base_c_; + // Base activation index for current threadblock + int offset_intial_npq_; + // Base activation coord for current threadblock + TensorCoord activatioin_base_; + // Intial thread positioin + int offset_initial_hwc_; + // Overall load instruction per thread. + int iterator_load_; + // thread loading position. + int iterator_hwc_; + // activation N is inside the Tensor or not + bool valid_n_; + + public: + + + MCTLASS_HOST_DEVICE + DepthwiseFpropActivationDirect2dConvTileAccessIteratorFixedStrideDilation( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = + MatrixCoord() + ) + : params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + offset_intial_npq_(threadblock_offset.row()), + offset_initial_hwc_(thread_idx), + iterator_load_(0) { + + base_c_ = threadblock_offset.column(); + + set_iteration_index(0); + + set_activation_coord(offset_intial_npq_); + + } + + MCTLASS_HOST_DEVICE + void set_activation_coord(int offset_npq) { + int offset_inital_n, offset_inital_p, offset_inital_q; + int residual; + + params_.pq_divmod(offset_inital_n, residual, offset_npq); + params_.q_divmod(offset_inital_p, offset_inital_q, residual); + + int base_n = offset_inital_n; + + int base_h = + offset_inital_p * OutputTileShape::kH * StrideShape::kRow - problem_size_.pad_h; + + int base_w = + offset_inital_q * OutputTileShape::kW * StrideShape::kColumn - problem_size_.pad_w; + + activatioin_base_ = TensorCoord(base_n, base_h, base_w, base_c_); + + valid_n_ = activatioin_base_.n() < problem_size_.N; + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) { + return Params( + problem_size, + layout, + {Shape::kRow, Shape::kColumn}, + {OutputTileShape::kN, OutputTileShape::kH, OutputTileShape::kW, OutputTileShape::kC}, + kActivationSize); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iterator_hwc_ = offset_initial_hwc_ + index * ThreadMap::kThreads; + iterator_load_ = index; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + // Go to next threadblock + offset_intial_npq_ += problem_size_.split_k_slices; + + set_iteration_index(0); + + set_activation_coord(offset_intial_npq_); + } + + /// Returns the coordinate in the activations tensor X that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + int c = iterator_hwc_ % ThreadMap::Detail::ShapeVec::kContiguous ; + int next = iterator_hwc_ / ThreadMap::Detail::ShapeVec::kContiguous ; + int h = next / ActivationShape::kW; + int w = next % ActivationShape::kW; + + c = c * AccessType::kElements; + + return activatioin_base_ + TensorCoord(0, h, w, c); + } + + /// Returns true if the current coordinate is within the activations tensor X + MCTLASS_HOST_DEVICE + bool valid() const { + TensorCoord coord = at(); + bool valid_c = coord.c() < problem_size_.C; + bool valid_h = coord.h() >= 0 && coord.h() < problem_size_.H; + bool valid_w = coord.w() >= 0 && coord.w() < problem_size_.W; + return valid_n_ ? valid_c & valid_h & valid_w : 0; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + AccessType const *ptr = + reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + + return ptr; + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + DepthwiseFpropActivationDirect2dConvTileAccessIteratorFixedStrideDilation &operator++() { + + ++iterator_load_; + iterator_hwc_ += ThreadMap::kThreads; + + if (iterator_load_ < ThreadMap::Iterations::kCount) { + return *this; + } + + iterator_load_ = 0; + iterator_hwc_ = offset_initial_hwc_; + + return *this; + } + + /// Determines the activation size loaded by iterator + MCTLASS_HOST_DEVICE + int get_load_size() { + return kActivationSize; + } + + /// Determines the iterations needed + MCTLASS_HOST_DEVICE + int get_iteration_num() { + return ThreadMap::Iterations::kCount; + } + + /// Determines whether the Depthwise fprop can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check stride and dilation constraint + if (problem_size.stride_h != StrideShape::kRow || problem_size.stride_w != StrideShape::kColumn) { + return Status::kErrorInvalidProblem; + } + + if (problem_size.dilation_h != DilationShape::kRow || problem_size.dilation_w != DilationShape::kColumn) { + return Status::kErrorInvalidProblem; + } + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_activation_tile_access_iterator_direct_conv_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_activation_tile_access_iterator_direct_conv_optimized.h new file mode 100644 index 0000000..9ed3543 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_activation_tile_access_iterator_direct_conv_optimized.h @@ -0,0 +1,291 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM A (activation tile) + matrix from memory. + + This iterator assumes TensorNHWC layout of tensors in Global Memory. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/threadblock/depthwise_direct_conv_params.h" +#include "mctlass/coord.h" +#include "mctlass/mctlass.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template > +class DepthwiseFpropActivationDirect2dConvTileAccessIteratorOptimized { + public: + // + // Types + // + + using Shape = Shape_; + using OutputTileShape = OutputTileShape_; + using Element = Element_; + using Layout = Layout_; + using TensorCoord = typename Layout::TensorCoord; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, "Require Iterations::kContiguous == 1"); + + static_assert(OutputTileShape::kN == 1, "Require OutputTileShape::kN == 1"); + static_assert(OutputTileShape::kC == Shape::kColumn, "Require OutputTile shape == channels per threadblock"); + + // + // Parameters structure + // + + using Params = Depthwise2dFpropDirectConvParams; + + private: + Conv2dProblemSize const &problem_size_; + Params const ¶ms_; + char const *pointer_; + + // Base channels for current threadblock + int base_c_; + // Base activation index for current threadblock + int offset_intial_npq_; + // Base activation coord for current threadblock + TensorCoord activatioin_base_; + // Intial thread positioin + int offset_initial_hwc_; + // Overall load instruction per thread. + int iterator_load_; + // thread loading position. + int iterator_hwc_; + // Number of loads for activations tensor X. + const int number_of_loads_; + + public: + + + MCTLASS_HOST_DEVICE + DepthwiseFpropActivationDirect2dConvTileAccessIteratorOptimized( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = + MatrixCoord() + ) + : params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + offset_intial_npq_(threadblock_offset.row()), + offset_initial_hwc_(thread_idx), + iterator_load_(0), + number_of_loads_(params.activation_load_count) { + + base_c_ = threadblock_offset.column(); + + set_activation_coord(offset_intial_npq_); + + set_iteration_index(0); + } + + MCTLASS_HOST_DEVICE + void set_activation_coord(int offset_npq) { + int offset_inital_n, offset_inital_p, offset_inital_q; + int residual; + + params_.pq_divmod(offset_inital_n, residual, offset_npq); + params_.q_divmod(offset_inital_p, offset_inital_q, residual); + + int base_n = offset_inital_n; + + int base_h = + offset_inital_p * OutputTileShape::kH * problem_size_.stride_h - problem_size_.pad_h; + + int base_w = + offset_inital_q * OutputTileShape::kW * problem_size_.stride_w - problem_size_.pad_w; + + activatioin_base_ = TensorCoord(base_n, base_h, base_w, base_c_); + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) { + return Params( + problem_size, + layout, + {Shape::kRow, Shape::kColumn}, + {OutputTileShape::kN, OutputTileShape::kH, OutputTileShape::kW, OutputTileShape::kC}, + sizeof_bits::value, + ThreadMap::kThreads, + ThreadMap::Detail::ShapeVec::kContiguous, + ThreadMap::kElementsPerAccess); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iterator_hwc_ = offset_initial_hwc_ + index * ThreadMap::kThreads; + iterator_load_ = index; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_HOST_DEVICE + void advance() { + // Go to next threadblock + offset_intial_npq_ += problem_size_.split_k_slices; + + set_activation_coord(offset_intial_npq_); + } + + /// Returns the coordinate in the activations tensor X that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + int c = iterator_hwc_ % ThreadMap::Detail::ShapeVec::kContiguous ; + int next = iterator_hwc_ / ThreadMap::Detail::ShapeVec::kContiguous ; + int h, w; + params_.activation_tile_w_divmod(h, w, next) ; + + c = c * AccessType::kElements; + + return activatioin_base_ + TensorCoord(0, h, w, c); + } + + /// Returns true if the current coordinate is within the activations tensor X + MCTLASS_HOST_DEVICE + bool valid() const { + TensorCoord coord = at(); + + return coord.n() < problem_size_.N && coord.h() >= 0 && coord.h() < problem_size_.H && + coord.w() >= 0 && coord.w() < problem_size_.W && coord.c() < problem_size_.C; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + TensorCoord coord = at(); + LongIndex offset = params_.layout(coord); + + AccessType const *ptr = + reinterpret_cast(pointer_ + offset * sizeof_bits::value / 8); + + return ptr; + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + DepthwiseFpropActivationDirect2dConvTileAccessIteratorOptimized &operator++() { + + ++iterator_load_; + iterator_hwc_ += ThreadMap::kThreads; + + if (iterator_load_ < number_of_loads_) { + return *this; + } + + iterator_load_ = 0; + iterator_hwc_ = offset_initial_hwc_; + + return *this; + } + + /// Determines the activation size loaded by iterator + MCTLASS_HOST_DEVICE + int get_load_size() { + return params_.activation_size; + } + + /// Determines the iterations needed + MCTLASS_HOST_DEVICE + int get_iteration_num() { + return number_of_loads_; + } + + /// Determines whether the Depthwise fprop can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + // check alignment constraint on iterator's contiguous dimension + if (problem_size.C % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_direct_conv_multistage.h b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_direct_conv_multistage.h new file mode 100644 index 0000000..f96bc21 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_direct_conv_multistage.h @@ -0,0 +1,551 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a multistage threadblock-scoped Implicit GEMM Convolution kernel. +*/ + +#pragma once + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/cache_operation.h" +#include "mctlass/conv/threadblock/depthwise_mma_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Epilogue stores the data into global memory + typename Epilogue_, + /// iterator implementation variants + conv::IteratorAlgorithm IteratorAlgorithm_ = conv::IteratorAlgorithm::kOptimized, + /// Used for partial specialization + typename Enable = bool> +class DepthwiseFpropDirectConvMultipleStage : + public DepthwiseDirectConvMmaBase { +public: + ///< Base class + using Base = DepthwiseDirectConvMmaBase; + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Policy describing tuning details + using Policy = Policy_; + + using Epilogue = Epilogue_; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + static conv::IteratorAlgorithm const kItertorAlgorithm = IteratorAlgorithm_; + + // + // Dependent types + // + + /// Fragment of accumulator tile + + using ElementC = typename Policy::Operator::ElementC; + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Internal structure exposed for introspection. + struct Detail { + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + }; + + private: + + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + DepthwiseFpropDirectConvMultipleStage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + MCTLASS_DEVICE + void copy_tiles_and_advance(IteratorA &iterator_A, + IteratorB &iterator_B, + int group_start_A = 0, + int group_start_B = 0) { + if (kItertorAlgorithm == conv::IteratorAlgorithm::kFixedStrideDilation) { + // Number of iterators is a static value. + iterator_A.set_iteration_index(group_start_A * IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast(this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + } + ++this->smem_iterator_A_; + } + } else { + // Number of iterators is a runtime value. + iterator_A.set_iteration_index(group_start_A * IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < iterator_A.get_iteration_num(); ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast(this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + } + ++this->smem_iterator_A_; + } + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA &iterator_A, + ///< Params of global memory iterator + typename IteratorA::Params const &iterator_a_params, + ///< iterator over B operand in global memory + IteratorB &iterator_B, + ///< Params of global memory iterator + typename IteratorB::Params const &iterator_b_params, + ///< initial value of accumulator + FragmentC const &src_accum, + /// Epilogue + Epilogue &epilogue, + ///< Output operator + typename Epilogue::OutputOp const &output_op, + ///< Tile iterator for destination + typename Epilogue::OutputTileIterator &destination_iterator, + ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + typename Epilogue::OutputTileIterator &source_iterator, + + int split_k_slices = 1 + ) { + + // + // Prologue + // + + // Issue several complete stages + MCTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; ++stage, --gemm_k_iterations) { + + if (stage == 0) { + iterator_B.set_iteration_index(0); + this->smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast(this->smem_iterator_B_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + } + + ++this->smem_iterator_B_; + } + } + + if(kItertorAlgorithm == conv::IteratorAlgorithm::kFixedStrideDilation){ + // Number of iterators is compilation static. + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast(this->smem_iterator_A_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + + } else { + // Number of iterators is a runtime value. + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_num(iterator_A.get_iteration_num()); + this->smem_iterator_A_.set_iteration_index(0); + + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < iterator_A.get_iteration_num(); ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast(this->smem_iterator_A_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + } + + // Move to the next stage + iterator_A.advance(); + + this->smem_iterator_A_.add_tile_offset({1, 0}); + + // Inserts a fence to group cp.async instructions into stages. + mctlass::arch::cp_async_fence(); + } + + ///////////////////////////////////////////////////////////////////////////// + // Waits until kStages-2 stages have committed. + mctlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[2]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpTransformedFragmentA warp_transformed_frag_A[2]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.setup_initial_status(iterator_a_params); + + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0], + warp_loaded_frag_A[0], warp_loaded_frag_B[0]); + + // + // Mainloop + // + + unsigned int iterations = 0; + constexpr int inner_loop_iterations = round_up(Base::kWarpGemmIterations, 2); + + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { // Each iteration is a cta tile. + + accum.clear(); + + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < inner_loop_iterations; ++warp_mma_k) { + if (Base::kWarpGemmIterations % 2 == 0 || warp_mma_k + 1 != Base::kWarpGemmIterations) { + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Shape::kK); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Shape::kK); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + } + + if (warp_mma_k > 0) + warp_mma.transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % 2], + warp_loaded_frag_B[warp_mma_k % 2]); + + // Issue global->shared copies for the next stage + int group_start_iteration_A, group_start_iteration_B; + + if (warp_mma_k == 0) { + group_start_iteration_A = 0; + group_start_iteration_B = 0; + copy_tiles_and_advance( + iterator_A, iterator_B, group_start_iteration_A, group_start_iteration_B); + } + + if (warp_mma_k < Base::kWarpGemmIterations) { + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum + ); + } + + if (warp_mma_k + 1 == inner_loop_iterations) + warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + if (warp_mma_k + 2 == inner_loop_iterations) { + // Inserts a fence to group cp.async instructions into stages. + mctlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages of cp.async have committed + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next cta + iterator_A.advance(); + + this->smem_iterator_A_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({-Base::kStages, 0}); + + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.advance(- (Base::kStages-1) * iterator_A.get_load_size()); + smem_read_stage_idx = 0; + } else { + this->warp_tile_iterator_A_.advance(iterator_A.get_load_size()); + ++smem_read_stage_idx; + } + + if (kItertorAlgorithm == conv::IteratorAlgorithm::kFixedStrideDilation) { + this->warp_tile_iterator_A_.setup_initial_status(iterator_a_params); + } + + // goback to start position. B has no multiple stage + this->warp_tile_iterator_B_.add_tile_offset({-Policy::kPartitionsK * Shape::kK, 0}); + + --gemm_k_iterations; + } + } + + // + // Epilogue + // + int32_t smem_base_offset = iterator_B.get_load_size() + (iterations % Base::kStages) * iterator_A.get_load_size(); + + destination_iterator.set_tile_index(iterations * split_k_slices); + + source_iterator.set_tile_index(iterations * split_k_slices); + + epilogue(output_op, destination_iterator, accum, source_iterator, smem_base_offset); + + ++iterations; + } + + // Insert fence and wait for all outstanding cp.async operations to commit. + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_filter_tile_access_iterator_direct_conv_optimized.h b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_filter_tile_access_iterator_direct_conv_optimized.h new file mode 100644 index 0000000..0db00a3 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_filter_tile_access_iterator_direct_conv_optimized.h @@ -0,0 +1,261 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of convolution tiles mapped to GEMM B (filter tile) + matrix from memory. + + This iterator assumes TensorNHWC layout of tensors in Global Memory. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/threadblock/conv2d_params.h" +#include "mctlass/conv/threadblock/conv2d_fprop_filter_tile_access_iterator_analytic.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +template > +class DepthwiseFpropFilterDirectConvTileAccessIteratorOptimized { +public: + // + // Types + // + + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + using TensorRef = mctlass::TensorRef; + using TensorCoord = typename Layout::TensorCoord; + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + static IteratorAlgorithm const kIteratorAlgorithm = conv::IteratorAlgorithm::kOptimized; + static StrideSupport const kStrideSupport = conv::StrideSupport::kStrided; + static int const kConvDim = 2; + using ConvProblemSize = typename conv::Conv2dProblemSize; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static int const kFilterSize = ThreadMap::Iterations::kCount * ThreadMap::kElementsPerAccess * ThreadMap::kThreads * + sizeof_bits::value / 8; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + // + // Simplifying assertions + // + static_assert(ThreadMap::Iterations::kContiguous == 1, + "Require Iterations::kContiguous == 1"); + + // + // Parameters structure + // + using Params = Depthwise2dFpropDirectConvFilterIteratorParams; + + protected: + + Conv2dProblemSize const &problem_size_; + Params const ¶ms_; + LongIndex iteration_contiguous_; + LongIndex iteration_strided_; + LongIndex iteration_vector_; + char const *pointer_; + + int filter_k_; + int offset_trs_[ThreadMap::Iterations::kStrided]; + +public: + + + + MCTLASS_HOST_DEVICE + DepthwiseFpropFilterDirectConvTileAccessIteratorOptimized( + Params const ¶ms, + Conv2dProblemSize const &problem_size, + Element const *ptr, + int thread_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + params_(params), + problem_size_(problem_size), + pointer_(reinterpret_cast(ptr)), + filter_k_(0) { + + layout::PitchLinearCoord thread_coord = ThreadMap::initial_offset(thread_idx); + + filter_k_ = threadblock_offset.column() + thread_coord.contiguous(); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + offset_trs_[s] = threadblock_offset.row() + thread_coord.strided() + s * ThreadMap::Delta::kStrided; + } + + set_iteration_index(0); + } + + MCTLASS_HOST_DEVICE + static Params getParams(Conv2dProblemSize const &problem_size, Layout const &layout) { + return Params(problem_size, layout, {Shape::kRow, Shape::kColumn}, kFilterSize); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(Index index) { + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset * 8 / sizeof_bits::value; + } + + MCTLASS_HOST_DEVICE + void advance() { + // Do nothing because the filter is persistent in the SMEM + } + + /// Returns the coordinate in the filter tensor W that is currently pointed to + /// by the iterator. + MCTLASS_HOST_DEVICE + TensorCoord at() const { + + int k = filter_k_ + iteration_vector_ * AccessType::kElements; + int trs = offset_trs_[iteration_strided_]; + + return TensorCoord(k, trs, 0 , 0); // As a 2D-matrix + } + + /// Returns true if the current coordinate is within the activations tensor W + MCTLASS_HOST_DEVICE + bool valid() const { + + TensorCoord coord = at(); + + return coord.n() < problem_size_.K && + coord.h() < Shape::kColumn; + } + + /// Returns a pointer to the vector starting at the current coordinate + MCTLASS_HOST_DEVICE + AccessType const *get() const { + TensorCoord coord = at(); + int64_t offset = coord.n(); + if (params_.is_convolution) { + offset += (Shape::kColumn - coord.h() - 1)* problem_size_.K; + } else { + offset += coord.h() * problem_size_.K; + } + + return reinterpret_cast(pointer_ + + offset * sizeof_bits::value / 8); + } + + /// Increments to the next memory access + MCTLASS_HOST_DEVICE + DepthwiseFpropFilterDirectConvTileAccessIteratorOptimized &operator++() { + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + iteration_vector_ = 0; + + ++iteration_contiguous_; + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + iteration_contiguous_ = 0; + + ++iteration_strided_; + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + iteration_strided_ = 0; + + return *this; + } + + /// Determines the filter size loaded by iterator + MCTLASS_HOST_DEVICE + int get_load_size() { + return kFilterSize; + } + + /// Determines whether the Implicit GEMM can execute the given problem. + MCTLASS_HOST_DEVICE + static Status can_implement(Conv2dProblemSize const &problem_size) { + + // check alignment constraint on iterator's contiguous dimension + if (problem_size.K % AccessType::kElements) { + return Status::kErrorInvalidProblem; + } + + // check whether runtime filter size is same as templated filter size. + if ((problem_size.R * problem_size.S) != Shape::kColumn) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_pipelined.h b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_pipelined.h new file mode 100644 index 0000000..0f17c11 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_fprop_pipelined.h @@ -0,0 +1,336 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/aligned_buffer.h" +#include "mctlass/numeric_conversion.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/mma_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Transformation applied to A operand + typename TransformA_ = NumericArrayConverter< + typename SmemIteratorA_::Element, + typename IteratorA_::Element, + IteratorA_::Fragment::kElements>, + /// + /// Transformation applied to A operand + typename TransformB_ = NumericArrayConverter< + typename SmemIteratorB_::Element, + typename IteratorB_::Element, + IteratorB_::Fragment::kElements>, + /// Used for partial specialization + typename Enable = bool +> +class DepthwiseFpropPipelined : public gemm::threadblock::MmaBase { +public: + + ///< Base class + using Base = gemm::threadblock::MmaBase; + + using Shape = Shape_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using IteratorA = IteratorA_; ///< Iterates over tiles of A operand in global memory + using IteratorB = IteratorB_; ///< Iterates over tiles of B operand in global memory + using ElementC = ElementC_; ///< Data type of accumulator matrix + using LayoutC = LayoutC_; ///< Layout of accumulator matrix + using Policy = Policy_; ///< Policy describing tuning details + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + using TransformA = TransformA_; + using TransformB = TransformB_; + + // + // Dependent types + // + + /// Fragment of operand A loaded from global memory + using FragmentA = typename IteratorA::Fragment; + + /// Fragment of operand B loaded from global memory + using FragmentB = typename IteratorB::Fragment; + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Obtain the arch tag from the warp-level operator + using ArchTag = typename Policy::Operator::ArchTag; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + // staticaly assert kStages for MmaPipelined is two (Double-buffered pipeline) + static_assert((Base::kStages==2), "MmaPipelined requires kStages set to value 2"); + +private: + + using WarpFragmentA = typename Operator::FragmentA; + using WarpFragmentB = typename Operator::FragmentB; + +protected: + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + DepthwiseFpropPipelined( + typename Base::SharedStorage &shared_storage, ///< Shared storage needed for internal use by threadblock-scoped GEMM + int thread_idx, ///< ID within the threadblock + int warp_idx, ///< ID of warp + int lane_idx ///< ID of each thread within a warp + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) { + + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + int gemm_k_iterations, ///< number of iterations of the mainloop + FragmentC &accum, ///< destination accumulator tile + IteratorA iterator_A, ///< iterator over A operand in global memory + IteratorB iterator_B, ///< iterator over B operand in global memory + FragmentC const &src_accum, ///< source accumulator tile + int gemm_k_iterations_per_channel = 0, ///< number of iterations per channel + TransformA transform_A = TransformA(), ///< transformation applied to A fragment + TransformB transform_B = TransformB()) { ///< transformation applied to B fragment + + // + // Prologue + // + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + FragmentA tb_frag_A; + FragmentB tb_frag_B; + + tb_frag_A.clear(); + tb_frag_B.clear(); + + // The last kblock is loaded in the prolog + iterator_A.load(tb_frag_A); + iterator_B.load(tb_frag_B); + + ++iterator_A; + ++iterator_B; + + this->smem_iterator_A_.store(transform_A(tb_frag_A)); + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math instructions + WarpFragmentA warp_frag_A[2]; + WarpFragmentB warp_frag_B[2]; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + Operator warp_mma; + + int smem_write_stage_idx = 1; + // Depthwise specific + int channel_start_index = 0; + int rs_plane_idx = 0; + + // Issue loads during the first warp-level matrix multiply-add *AFTER* issuing + // shared memory loads (which have the tightest latency requirement). + + // + // Mainloop + // + + // Note: The main loop does not support Base::kWarpGemmIterations == 2. + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > 0; --gemm_k_iterations) { + // + // Loop over GEMM K dimension + // + + if(rs_plane_idx == gemm_k_iterations_per_channel - 1){ + // Reset interation index. + iterator_B.set_iteration_index(0); + } + + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if this is the last group + // as the case may be. + + if (warp_mma_k == Base::kWarpGemmIterations - 1) { + + // Write fragments to shared memory + this->smem_iterator_A_.store(transform_A(tb_frag_A)); + + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + __syncthreads(); + + if(rs_plane_idx == gemm_k_iterations_per_channel - 1){ + // Move to next set of filter groups. + channel_start_index += Base::kWarpGemmIterations; + } + + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + // Add negative offsets to return iterators to the 'start' of the circular buffer in shared memory + if (smem_write_stage_idx == 1) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + } + else { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations, + 0}); + } + + smem_write_stage_idx ^= 1; + } + + this->warp_tile_iterator_A_.set_kgroup_index(channel_start_index + (warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index(channel_start_index + (warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k == 0) { + + iterator_A.load(tb_frag_A); + iterator_B.load(tb_frag_B); + + ++iterator_A; + ++iterator_B; + } + + warp_mma(accum, warp_frag_A[warp_mma_k % 2], + warp_frag_B[warp_mma_k % 2], accum); + } + + rs_plane_idx = (rs_plane_idx == gemm_k_iterations_per_channel - 1) ? 0: (rs_plane_idx + 1); + + } + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_mma_base.h b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_mma_base.h new file mode 100644 index 0000000..043fc49 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_mma_base.h @@ -0,0 +1,229 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a directconv threadblock-scoped Depthwise kernel. +*/ + +#pragma once + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Policy object describing MmaTensorOp +template < + /// Warp-level GEMM operator (concept: gemm::warp::Mma) + typename Operator_, + /// Padding used for A operand in shared memory (concept: MatrixShape) + typename SmemPaddingA_, + /// Padding used for B operand in shared memory (concept: MatrixShape) + typename SmemPaddingB_, + /// + typename ThreadMapA_, + /// + typename ThreadMapB_, + /// Number of partitions of K dimension of GEMM + int PartitionsK = 1> +struct DepthwiseDirectConvMmaPolicy { + /// Warp-level GEMM operator (concept: gemm::warp::MmaTensorOp or gemm::warp::MmaSimt) + using Operator = Operator_; + + /// Padding used for A operand in shared memory + using SmemPaddingA = SmemPaddingA_; + + /// Padding used for B operand in shared memory + using SmemPaddingB = SmemPaddingB_; + + using ThreadMapA = ThreadMapA_; + using ThreadMapB = ThreadMapB_; + + /// Number of partitions of K dimension + static int const kPartitionsK = PartitionsK; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class DepthwiseDirectConvMmaBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Policy describing tuning details + using Policy = Policy_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = mctlass::gemm:: + GemmShape; + + /// Number of warp-level GEMM oeprations + /// kWarpGemmIterations could be even and odd. + static int const kWarpGemmIterations = (WarpGemm::kK / Operator::Policy::MmaShape::kK); + + /// Number of stages + static int const kStages = Stages; + + /// Tensor reference to the A operand + using TensorRefA = TensorRef; + + /// Tensor reference to the B operand + using TensorRefB = TensorRef; + + static_assert(kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + class SharedStorage { + public: + // + // Type definitions + // + + /// Shape of the A matrix operand in shared memory + using ShapeA = MatrixShape<1, // Not determined at compile-time :( + Shape::kN + Policy::SmemPaddingA::kRow>; + + /// Shape of the B matrix operand in shared memory + using ShapeB = MatrixShape; // Tile N = 64? + + public: + // + // Data members + // + + // Let persistent B matrix in front of dynamic matrix A + /// Buffer for B operand + AlignedBuffer operand_B; + + /// Buffer for A operand + /// Not be determined at compile-time -- Just to get a Smem start address. + AlignedBuffer operand_A; + public: + // + // Methods + // + + /// Returns a layout object for the A matrix + MCTLASS_DEVICE + static typename Operator::LayoutA LayoutA() { + return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn}); + } + + /// Returns a layout object for the B matrix + MCTLASS_HOST_DEVICE + static typename Operator::LayoutB LayoutB() { + return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); + } + + /// Returns a TensorRef to the A operand + MCTLASS_HOST_DEVICE + TensorRefA operand_A_ref() { return TensorRefA{operand_A.data(), LayoutA()}; } + + /// Returns a TensorRef to the B operand + MCTLASS_HOST_DEVICE + TensorRefB operand_B_ref() { return TensorRefB{operand_B.data(), LayoutB()}; } + }; + + protected: + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A operand from shared memory + typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + + public: + /// Construct from tensor references + MCTLASS_DEVICE + DepthwiseDirectConvMmaBase( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx), + warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) {} +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_mma_core_with_lane_access_size.h b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_mma_core_with_lane_access_size.h new file mode 100644 index 0000000..8517235 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/depthwise_mma_core_with_lane_access_size.h @@ -0,0 +1,952 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic properties needed by CTA-level GEMMs assuming expectations about data + layout of the global memory fragments, data types, and internal tile sizes. + + Partial specializations for threadblock::Mma operations targeting depthwise related simt instructions. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/warp/mma_depthwise_simt.h" + +#include "mctlass/gemm/threadblock/mma_pipelined.h" +#include "mctlass/gemm/threadblock/mma_singlestage.h" + +#include "mctlass/gemm/threadblock/mma_base.h" +#include "mctlass/conv/threadblock/depthwise_mma_base.h" + +#include "mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear_direct_conv.h" + +#include "mctlass/arch/cache_operation.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + + +namespace detail { +// +// Convert a WarpShapeM which is the whole tile of elements into the number of elements (2D) held by +// each partitions within warp. +// The goal is for each thread's tile of elements to be as square as +// possible for performance (4x4 will be faster than 2x8). +template // The number of partitions within the warp +struct SimtWarpShape { + // kP * kQ * WarpNumThreadsM = WarpShapeM + // If needed, enable more specializations. +}; +template <> +struct SimtWarpShape<4, 4> { + static constexpr int kP = 1; + static constexpr int kQ = 1; +}; + +template <> +struct SimtWarpShape<4, 2> { + static constexpr int kP = 2; + static constexpr int kQ = 1; +}; + +template <> +struct SimtWarpShape<4, 1> { + static constexpr int kP = 2; + static constexpr int kQ = 2; +}; + +template <> +struct SimtWarpShape<8, 1> { + static constexpr int kP = 2; + static constexpr int kQ = 4; +}; +template <> +struct SimtWarpShape<8, 2> { + static constexpr int kP = 2; + static constexpr int kQ = 2; +}; +template <> +struct SimtWarpShape<8, 4> { + static constexpr int kP = 1; + static constexpr int kQ = 2; +}; + +template <> +struct SimtWarpShape<16, 1> { + static constexpr int kP = 4; + static constexpr int kQ = 4; +}; +template <> +struct SimtWarpShape<16, 2> { + static constexpr int kP = 2; + static constexpr int kQ = 4; +}; +template <> +struct SimtWarpShape<16, 4> { + static constexpr int kP = 2; + static constexpr int kQ = 2; +}; + +template +struct SimtWarpShape<25, WarpNumThreadsM> { + static_assert(WarpNumThreadsM == 1, "WarpShapeM could not be evenly splited by threads"); + static constexpr int kP = 5; + static constexpr int kQ = 5; +}; + +template <> +struct SimtWarpShape<32, 1> { + static constexpr int kP = 4; + static constexpr int kQ = 8; +}; + +template <> +struct SimtWarpShape<32, 2> { + static constexpr int kP = 4; + static constexpr int kQ = 4; +}; + +template <> +struct SimtWarpShape<32, 4> { + static constexpr int kP = 2; + static constexpr int kQ = 4; +}; + +} // namespace detail +template < + /// Shape of threadblock-scoped matrix multiply operator + typename Shape, + /// Shape of warp-level matrix multiply operator + typename WarpShape, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape, + /// Element data type of A operand + typename ElementA, + /// Layout of operand A + typename LayoutA, + /// Element data type of B operand + typename ElementB, + /// Layout of operand B + typename LayoutB, + /// Data type of accumulator + typename ElementC, + /// Layout of accumulator + typename LayoutC, + /// Indicates type of math operator (arch::OpClassSimt or arch::OpClassTensorOp) + typename OperatorClass, + /// Size of a warp-scoped per thread access + int kLaneAccessSizeA_ = 0, + /// Size of a warp-scoped per thread access + int kLaneAccessSizeB_ = 0, + /// Number of stages + int Stages = 2, + /// Operation performed by MMA + typename Operator = typename platform::conditional< + (platform::is_same::value) && + (platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value), + mctlass::arch::OpMultiplyAddSaturate, + mctlass::arch::OpMultiplyAdd>::type, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA = + mctlass::arch::CacheOperation::Global, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB = + mctlass::arch::CacheOperation::Global, + /// per-element transformation for elements of A + ComplexTransform TransformA = ComplexTransform::kNone, + /// per-element transformation for elements of B + ComplexTransform TransformB = ComplexTransform::kNone, + bool IsComplex = false // (is_complex::value || is_complex::value) +> +struct DepthwiseMmaCoreWithLaneAccessSize; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Shape of threadblock-scoped matrix multiply operator + typename Shape, + /// Shape of threadblock-scoped output tile + typename ThreadBlockOutputShape, + /// Shape of filter shape per threadblock + typename FilterShape, + /// Shape of warp-level matrix multiply operator + typename WarpShape, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape, + /// Element data type of A operand + typename ElementA, + /// Layout of operand A + typename LayoutA, + /// Element data type of B operand + typename ElementB, + /// Layout of operand B + typename LayoutB, + /// Data type of accumulator + typename ElementC, + /// Layout of accumulator + typename LayoutC, + /// Indicates type of math operator (arch::OpClassSimt or arch::OpClassTensorOp) + typename OperatorClass, + /// Size of a warp-scoped per thread access + int kLaneAccessSizeA_ = 0, + /// Size of a warp-scoped per thread access + int kLaneAccessSizeB_ = 0, + /// Number of stages + int Stages = 2, + /// Operation performed by MMA + typename Operator = typename platform::conditional< + (platform::is_same::value) && + (platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value), + mctlass::arch::OpMultiplyAddSaturate, + mctlass::arch::OpMultiplyAdd>::type, + /// Iterator algo type + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kAnalytic, + /// Stride ( MatrixShape ) + typename StrideShape = mctlass::MatrixShape<-1, -1>, + /// Dilation ( MatrixShape ) + typename DilationShape = mctlass::MatrixShape<-1, -1>, + /// Activation Shape loaded by threadblock + typename ActivationShape = mctlass::conv::TensorNHWCShape<-1,-1,-1,-1>, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA = + mctlass::arch::CacheOperation::Global, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB = + mctlass::arch::CacheOperation::Global, + /// per-element transformation for elements of A + ComplexTransform TransformA = ComplexTransform::kNone, + /// per-element transformation for elements of B + ComplexTransform TransformB = ComplexTransform::kNone, + bool IsComplex = false // (is_complex::value || is_complex::value) +> +struct DepthwiseDirectConvMmaCoreWithLaneAccessSize; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Shape of threadblock-scoped matrix multiply operator + typename Shape, + /// Shape of warp-level matrix multiply operator + typename WarpShape, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape, + /// Element data type of A operand + typename ElementA, + /// Layout of operand A + typename LayoutA, + /// Element data type of B operand + typename ElementB, + /// Layout of operand B + typename LayoutB, + /// Data type of accumulator + typename ElementC, + /// Layout of accumulator + typename LayoutC, + /// Indicates type of math operator (arch::OpClassSimt or arch::OpClassTensorOp) + typename OperatorClass, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// per-element transformation for elements of A + ComplexTransform TransformA, + /// per-element transformation for elements of B + ComplexTransform TransformB, + bool IsComplex +> +struct DepthwiseMmaCoreWithLaneAccessSize< + Shape, WarpShape, InstructionShape, + ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, + OperatorClass, -1, -1, Stages, Operator, AccumulatorsInRowMajor, + CacheOpA, CacheOpB, TransformA, TransformB, IsComplex +> : mctlass::gemm::threadblock::DefaultMmaCore< + Shape, WarpShape, InstructionShape, + ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, + OperatorClass, Stages, Operator, AccumulatorsInRowMajor, + CacheOpA, CacheOpB, TransformA, TransformB, IsComplex +> {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: column-major +/// Operator: simt class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Size of a warp-scoped per thread access (a value of -1 indicates the default) + int kLaneAccessSizeA_, + /// Size of a warp-scoped per thread access (a value of -1 indicates the default) + int kLaneAccessSizeB_, + /// Operation performed by GEMM + typename Operator_> +struct DepthwiseMmaCoreWithLaneAccessSize, + ElementA_, + layout::RowMajor, + ElementB_, + layout::ColumnMajor, + ElementC_, + LayoutC_, + arch::OpClassSimt, + kLaneAccessSizeA_, + kLaneAccessSizeB_, + 2, + Operator_> : public mctlass::gemm::threadblock::DefaultMmaCore, + ElementA_, + layout::RowMajor, + ElementB_, + layout::ColumnMajor, + ElementC_, + LayoutC_, + arch::OpClassSimt, + 2, + Operator_> { + using Base = mctlass::gemm::threadblock::DefaultMmaCore, + ElementA_, + layout::RowMajor, + ElementB_, + layout::ColumnMajor, + ElementC_, + LayoutC_, + arch::OpClassSimt, + 2, + Operator_>; + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = mctlass::gemm::GemmShape<1, 1, 1>; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + + static int const kLaneAccessSizeA = kLaneAccessSizeA_; + static int const kLaneAccessSizeB = kLaneAccessSizeB_; + + // Divisility requirements + static_assert( kLaneAccessSizeA > 0 && kLaneAccessSizeB > 0, + "Size of a warp-scoped per thread access should be larger then ZERO" ); + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = typename Base::WarpCount; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = mctlass::gemm::warp::WarpSize::value; + + static int const kElementsPerAccess = 1; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory are same as base class + // + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = mctlass::gemm::threadblock::detail::simt_get_warp_threads_m(); + static const int WarpNumThreadsN = kWarpSize / WarpNumThreadsM; + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = kLaneAccessSizeA / sizeof_bits::value; + static const int numElementsB = kLaneAccessSizeB / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + + static int const kPaddingM = mctlass::gemm::threadblock::detail::simt_transpose_padding(kWarpSize, Shape::kK, sizeof_bits::value); + static int const kPaddingN = mctlass::gemm::threadblock::detail::simt_transpose_padding(kWarpSize, Shape::kK, sizeof_bits::value); + + static_assert(!(kPaddingM % LaneM) && !(kPaddingN % LaneN), + "Padding must be divisible by Lane"); + + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::conv::warp::MmaDepthwiseSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy /// Policy describing warp-level MmaSimtOp (concept: MmaSimtOp policy) + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = mctlass::gemm::threadblock::MmaPolicy< + MmaWarpSimt, + MatrixShape, // skew for A matrix to avoid SMEM bank conflicts + MatrixShape<0, kPaddingN>, // skew for B matrix to avoid SMEM bank conflicts + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: row-major +/// Operator: simt class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of threadblock-scoped output tile (concept: TensorNHWCShape) + typename ThreadBlockOutputShape_, + /// Shape of filter shape per threadblock + typename FilterShape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Size of a warp-scoped per thread access + int kLaneAccessSizeA_, + /// Number of stages + int Stages_, + /// Operation performed by GEMM + typename Operator_> +struct DepthwiseDirectConvMmaCoreWithLaneAccessSize, + ElementA_, + layout::RowMajor, + ElementB_, + layout::ColumnMajor, + ElementC_, + LayoutC_, + arch::OpClassSimt, + kLaneAccessSizeA_, + 128, + Stages_, + Operator_> { + using Shape = Shape_; + using FilterShape = FilterShape_; + using WarpShape = WarpShape_; + using InstructionShape = mctlass::gemm::GemmShape<1, 1, 1>; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + + static int const kLaneAccessSizeB = 128; + + // Divisility requirements + static_assert( kLaneAccessSizeB > 0, + "Size of a warp-scoped per thread access should be larger then ZERO" ); + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = mctlass::gemm::GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + 1 + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = mctlass::gemm::warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + // For Gmem load + static int const kElementsPerAccessA = 128 / sizeof_bits::value; + static int const kElementsPerAccessB = 128 / sizeof_bits::value; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajor; + using SmemLayoutB = layout::RowMajor; + + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, // Set kStrided = 1 because activation shape is runtime value. + kThreads, + kElementsPerAccessA + >; + + /// ThreadMap of iterator A + using SmemThreadMapA = IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIteratorDirectConv< + MatrixShape<1, Shape::kN>, // set kRow is 1 because it is a runtime value + ElementA, + SmemLayoutA, + 0, + SmemThreadMapA, // was IteratorThreadMapA + true // Dynamic iterations. + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccessB + >; + + /// Transpose the ThreadMap of iterator B + using SmemThreadMapB = IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIteratorDirectConv< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + SmemThreadMapB, // was IteratorThreadMapB + false // static iterations. + >; + + // + // Warp-level matrix multiply operator + // + // Groups per threads + // Fp32: 2 groups + // Fp16: 2 groups + static const int GroupsPerThread = sizeof(ElementB) > 1 ? 2 : 4; + // Define the warp-level op + static const int WarpNumThreadsN = mctlass::const_min(WarpShape::kN / GroupsPerThread, kWarpSize); + static const int WarpNumThreadsM = kWarpSize / WarpNumThreadsN; + + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + + // Get output P, Q per thread + static const int TileP = mctlass::conv::threadblock::detail::SimtWarpShape::kP; + static const int TileQ = mctlass::conv::threadblock::detail::SimtWarpShape::kQ; + + static const int LaneLayout = 1; + static const int numElementsB = kLaneAccessSizeB / sizeof_bits::value; + static const int LaneN = mctlass::const_min(numElementsB, WarpShape::kN / WarpNumThreadsN); + + // Define the output tile computed by each thread + using ThreadOutputShape = mctlass::conv::TensorNHWCShape<1, TileP, TileQ, LaneN>; + + // Fetch the channel with same access size + static const int LaneM = LaneN; + + // No paddings + static int const kPaddingM = 0; + static int const kPaddingN = 0; + + static_assert(!(kPaddingM % LaneM) && !(kPaddingN % LaneN), + "Padding must be divisible by Lane"); + + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::conv::warp::MmaDepthwiseDirectConvSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> + FilterShape, /// Shape of filter shape per threadblock - concept: gemm::GemmShape + ThreadOutputShape, /// Size of the output tile computed by thread - concept: conv::TensorNHWCShape<> + ThreadBlockOutputShape_, /// Size of the output tile computed by threadblock - concept: conv::TensorNHWCShape<> + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy /// Policy describing warp-level MmaSimtOp (concept: MmaSimtOp policy) + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = mctlass::conv::threadblock::DepthwiseDirectConvMmaPolicy< + MmaWarpSimt, + MatrixShape, // skew for A matrix to avoid SMEM bank conflicts + MatrixShape<0, kPaddingN>, // skew for B matrix to avoid SMEM bank conflicts + IteratorThreadMapA, + IteratorThreadMapB, + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: row-major +/// Operator: simt class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of threadblock-scoped output tile (concept: TensorNHWCShape) + typename ThreadBlockOutputShape_, + /// Shape of filter shape per threadblock + typename FilterShape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Size of a warp-scoped per thread access + int kLaneAccessSizeA_, + /// Number of stages + int Stages_, + /// Operation performed by GEMM + typename Operator_, + /// Stride ( MatrixShape ) + typename StrideShape_, + /// Dilation ( MatrixShape ) + typename DilationShape_, + /// Activation Shape loaded by threadblock + typename ActivationShape_> +struct DepthwiseDirectConvMmaCoreWithLaneAccessSize, + ElementA_, + layout::RowMajor, + ElementB_, + layout::ColumnMajor, + ElementC_, + LayoutC_, + arch::OpClassSimt, + kLaneAccessSizeA_, + 128, + Stages_, + Operator_, + IteratorAlgorithm::kFixedStrideDilation, + StrideShape_, + DilationShape_, + ActivationShape_> { + using Shape = Shape_; + using FilterShape = FilterShape_; + using WarpShape = WarpShape_; + using InstructionShape = mctlass::gemm::GemmShape<1, 1, 1>; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + using StrideShape = StrideShape_; + using DilationShape = DilationShape_; + using ThreadBlockOutputShape = ThreadBlockOutputShape_; + using ActivationShape = ActivationShape_; + + static int const kLaneAccessSizeB = 128; + + // Divisility requirements + static_assert( kLaneAccessSizeB > 0, + "Size of a warp-scoped per thread access should be larger then ZERO" ); + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = mctlass::gemm::GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + 1 + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = mctlass::gemm::warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + // For Gmem load + static int const kElementsPerAccessA = 128 / sizeof_bits::value; + static int const kElementsPerAccessB = 128 / sizeof_bits::value; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajor; + using SmemLayoutB = layout::RowMajor; + + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccessA + >; + + /// ThreadMap of iterator A + using SmemThreadMapA = IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIteratorDirectConv< + MatrixShape, + ElementA, + SmemLayoutA, + 0, + SmemThreadMapA, // was IteratorThreadMapA + false // static iterations. + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccessB + >; + + /// Transpose the ThreadMap of iterator B + using SmemThreadMapB = IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIteratorDirectConv< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + SmemThreadMapB, // was IteratorThreadMapB + false // static iterations. + >; + + // + // Warp-level matrix multiply operator + // + // Groups per threads + // Fp32: 2 groups + // Fp16: 2 groups + static const int GroupsPerThread = sizeof(ElementB) > 1 ? 2 : 4; + // Define the warp-level op + static const int WarpNumThreadsN = mctlass::const_min(WarpShape::kN / GroupsPerThread, kWarpSize); + static const int WarpNumThreadsM = kWarpSize / WarpNumThreadsN; + + static const int TileP = mctlass::conv::threadblock::detail::SimtWarpShape::kP; + static const int TileQ = mctlass::conv::threadblock::detail::SimtWarpShape::kQ; + + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + + static const int LaneLayout = 1; + static const int numElementsB = kLaneAccessSizeB / sizeof_bits::value; + static const int LaneN = mctlass::const_min(numElementsB, WarpShape::kN / WarpNumThreadsN); + + // Define the output tile computed by each thread + using ThreadOutputShape = mctlass::conv::TensorNHWCShape<1, TileP, TileQ, LaneN>; + + // Fetch the channel with same access size + static const int LaneM = LaneN; + + // No paddings + static int const kPaddingM = 0; + static int const kPaddingN = 0; + + static_assert(!(kPaddingM % LaneM) && !(kPaddingN % LaneN), + "Padding must be divisible by Lane"); + + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::conv::warp::MmaDepthwiseDirectConvSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> + FilterShape, /// Shape of filter shape per threadblock - concept: gemm::GemmShape + ThreadOutputShape, /// Size of the output tile computed by thread - concept: conv::TensorNHWCShape<> + ThreadBlockOutputShape, /// Size of the output tile computed by threadblock - concept: conv::TensorNHWCShape<> + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy, /// Policy describing warp-level MmaSimtOp (concept: MmaSimtOp policy) + IteratorAlgorithm::kFixedStrideDilation, /// Iterator algo type + StrideShape, /// Stride ( MatrixShape ) + DilationShape, /// Dilation ( MatrixShape ) + ActivationShape /// Activation Shape loaded by threadblock + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = mctlass::conv::threadblock::DepthwiseDirectConvMmaPolicy< + MmaWarpSimt, + MatrixShape, // skew for A matrix to avoid SMEM bank conflicts + MatrixShape<0, kPaddingN>, // skew for B matrix to avoid SMEM bank conflicts + IteratorThreadMapA, + IteratorThreadMapB, + WarpCount::kK + >; +}; +} // namespace threadblock +} // namespace conv +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_fprop_fusion_multistage.h b/csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_fprop_fusion_multistage.h new file mode 100644 index 0000000..ff8836c --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_fprop_fusion_multistage.h @@ -0,0 +1,802 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a multistage threadblock-scoped fused activation's + scale+bias+relu and Implicit GEMM Convolution kernel. + + The original implicit gemm will store out-of-bound data as zeroes in the + shared memory because zeros into the tensor core, zeroes out of the tensor + cores. The result is remained the same. When fusing scale+bias+relu + into the mainloop, it is no longer true because + + 0 x scale + bias = bias + + which is no longer always 0. So, instead of storing zeroes, this fused + kernel stores the out-of-bound data as a special NaN (0x7eff), when applying + scale+bias+relu, the code is like + + if (data == 0x7eff) + data = 0; + else + data = scale+bias+relu(data, scale, bias); + + See include/mctlass/conv/warp/scale_bias_relu_transformation.h for the + elementwise computation. See include/mctlass/arch/memory_sm80.h for nan fill. +*/ + +#pragma once + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/cache_operation.h" +#include "mctlass/gemm/gemm.h" + +#include "mctlass/gemm/warp/scale_bias_tile_iterator.h" +#include "mctlass/conv/warp/scale_bias_relu_transform.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Element type of scale and bias vectors + typename ElementScaleBias_, + /// Layout of scale and bias vectors + typename LayoutScaleBias_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// WarpIterator to load Scale or Bias vector from the shared memory + typename WarpIteratorScaleBias_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class MmaFpropFusionBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Element type of scale and bias vectors + using ElementScaleBias = ElementScaleBias_; + + /// Layout of scale and bias vectors + using LayoutScaleBias = LayoutScaleBias_; + + ///< Policy describing tuning details + using Policy = Policy_; + + ///< WarpIterator to load Scale or Bias vector from the shared memory + using WarpIteratorScaleBias = WarpIteratorScaleBias_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = mctlass::gemm::GemmShape; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = + (WarpGemm::kK / Operator::Policy::MmaShape::kK); + + /// Number of stages + static int const kStages = Stages; + + /// Tensor reference to the A operand + using TensorRefA = TensorRef; + + /// Tensor reference to the scale and bias vectors + using TensorRefScaleBias = TensorRef; + + /// Tensor reference to the B operand + using TensorRefB = TensorRef; + + static_assert(kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + static_assert((kWarpGemmIterations % 2) == 0, + "Inner loop iteration must be an even number."); + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + class SharedStorage { + public: + // + // Type definitions + // + + /// Shape of the A matrix operand in shared memory + using ShapeA = MatrixShape; + + /// Shape of the A scale and bias vectors in shared memory + using ShapeScaleBias = + MatrixShape<1 + Policy::SmemPaddingA::kRow, + 2 * Shape::kK * kStages + Policy::SmemPaddingA::kColumn>; + + /// Shape of the B matrix operand in shared memory + using ShapeB = + MatrixShape; + + public: + // + // Data members + // + + /// Buffer for A operand + AlignedBuffer operand_A; + + /// Buffer for B operand + AlignedBuffer operand_B; + + /// Buffer for A operand Scale and Bias + AlignedBuffer operand_A_scale_bias; + + public: + + // + // Methods + // + + /// Returns a layout object for the A matrix + MCTLASS_DEVICE + static typename Operator::LayoutA LayoutA() { + return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn}); + } + + /// Returns a layout object for the B matrix + MCTLASS_HOST_DEVICE + static typename Operator::LayoutB LayoutB() { + return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); + } + + /// Returns a layout object for the A scale and bias vectors + MCTLASS_DEVICE + static LayoutScaleBias LayoutScaleBias() { + return LayoutScaleBias::packed( + {ShapeScaleBias::kRow, ShapeScaleBias::kColumn}); + } + + /// Returns a TensorRef to the A operand + MCTLASS_HOST_DEVICE + TensorRefA operand_A_ref() { + return TensorRefA{operand_A.data(), LayoutA()}; + } + + /// Returns a TensorRef to the B operand + MCTLASS_HOST_DEVICE + TensorRefB operand_B_ref() { + return TensorRefB{operand_B.data(), LayoutB()}; + } + + /// Returns a TensorRef to the A operand Scale vector + MCTLASS_HOST_DEVICE + TensorRefScaleBias operand_A_scale_bias_ref() { + return TensorRefScaleBias{operand_A_scale_bias.data(), LayoutScaleBias()}; + } + }; + + protected: + + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A operand from shared memory + typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of A operand scale and bias vector + /// from shared memory + WarpIteratorScaleBias warp_tile_iterator_A_scale_bias_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaFpropFusionBase( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx), + warp_tile_iterator_A_scale_bias_( + shared_storage.operand_A_scale_bias_ref(), lane_idx), + warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) {} +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// Iterates over vectors of scale and bias vector in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorScaleBias_, + /// Iterates over vectors of scale and bias vector in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorScaleBias_, + /// Cache operation for scale/bias operand + mctlass::arch::CacheOperation::Kind CacheOpScaleBias, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// WarpIterator to load Scale or Bias vector from the shared memory + typename WarpIteratorScaleBias_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class ImplicitGemmFpropFusionMultistage + : public MmaFpropFusionBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Iterates over tiles of the scale and bias vectors in global memory + using IteratorScaleBias = IteratorScaleBias_; + ///< WarpIterator to load Scale or Bias vector from the shared memory + using WarpIteratorScaleBias = WarpIteratorScaleBias_; + ///< Policy describing tuning details + using Policy = Policy_; + ///< Base class + using Base = MmaFpropFusionBase; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + using SmemIteratorScaleBias = SmemIteratorScaleBias_; + + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + static mctlass::arch::CacheOperation::Kind const kCacheOpScaleBias = + CacheOpScaleBias; + + // + // Dependent types + // + + /// Fragment of accumulator tile + + using ElementC = typename Policy::Operator::ElementC; + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Internal structure exposed for introspection. + struct Detail { + + static_assert(Base::kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + }; + + private: + + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpLoadedFragmentScaleBias = + typename WarpIteratorScaleBias::Fragment; + + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of A operand scale vector to shared memory + SmemIteratorScaleBias smem_iterator_A_scale_bias_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + ImplicitGemmFpropFusionMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_A_scale_bias_(shared_storage.operand_A_scale_bias_ref(), + thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_A_scale_bias_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + MCTLASS_DEVICE + void copy_tiles_and_advance(IteratorA &iterator_A, + IteratorScaleBias &iterator_A_scale_bias, + IteratorB &iterator_B, int group_start_A = 0, + int group_start_B = 0) { + iterator_A.set_iteration_index(group_start_A); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / 8; + + // Uses nan fill for out of bound data + mctlass::arch::cp_async_nan( + dst_ptr, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + + ++this->smem_iterator_A_; + } + } + + // Async Copy for operand A scale and bias vector. Scale and bias vectors + // are small. One iteration is enough. + if (group_start_A == 0) { + typename IteratorScaleBias::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_scale_bias_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorScaleBias::kElementsPerAccess / 8; + + mctlass::arch::cp_async( + dst_ptr, iterator_A_scale_bias.get(), iterator_A_scale_bias.valid()); + } + + iterator_B.set_iteration_index(group_start_B); + + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + ++this->smem_iterator_B_; + } + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< iterator over scale and bias vectors in global memory + IteratorScaleBias iterator_A_scale_bias, + ///< initial value of accumulator + FragmentC const &src_accum, + ///< number of iterations per channel + int gemm_k_iterations_per_channel = 0, + ///< Imaginary strides used for planar-complex only - ignored here + int64_t imag_stride_A = 0, + int64_t imag_stride_B = 0) { + + // + // Prologue + // + + // Issue several complete stages + MCTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; + ++stage, --gemm_k_iterations) { + + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / 8; + + // Uses Nan fill for out of bound data + mctlass::arch::cp_async_nan( + dst_ptr, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + ++this->smem_iterator_A_; + } + + // Async Copy for operand A scale and bias vectors. Scale and bias + // vectors are small. One iteration is enough. + { + typename IteratorScaleBias::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_scale_bias_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorScaleBias::kElementsPerAccess / 8; + + mctlass::arch::cp_async( + dst_ptr, iterator_A_scale_bias.get(), iterator_A_scale_bias.valid()); + } + + iterator_B.set_iteration_index(0); + this->smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + ++this->smem_iterator_B_; + } + + // Move to the next stage + iterator_A.advance(); + iterator_A_scale_bias.advance(); + iterator_B.advance(); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_A_scale_bias_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Inserts a fence to group cp.async instructions into stages. + mctlass::arch::cp_async_fence(); + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // Waits until kStages-2 stages have committed. + mctlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[2]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpLoadedFragmentScaleBias warp_loaded_frag_A_scale_bias[2]; + WarpTransformedFragmentA warp_transformed_frag_A[2]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + mctlass::conv::warp::FpropScaleBiasReluTransform + elementwise_transform; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_A_scale_bias_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_A_scale_bias_.load( + warp_loaded_frag_A_scale_bias[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_A_scale_bias_; + ++this->warp_tile_iterator_B_; + + // Start issuing the first group of the next stage outside of the mainloop + copy_tiles_and_advance(iterator_A, iterator_A_scale_bias, iterator_B); + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0], + warp_loaded_frag_A[0], warp_loaded_frag_B[0]); + + elementwise_transform(warp_transformed_frag_A[0], + warp_loaded_frag_A_scale_bias[0]); + + // + // Mainloop + // + + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_A_scale_bias_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_A_scale_bias_.load( + warp_loaded_frag_A_scale_bias[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_A_scale_bias_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) { + warp_mma.transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % 2], + warp_loaded_frag_B[warp_mma_k % 2]); + + elementwise_transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_loaded_frag_A_scale_bias[warp_mma_k % 2]); + } + + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum + ); + + // Issue global->shared copies for the next stage + int group_start_iteration_A, group_start_iteration_B; + + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + group_start_iteration_A = 0; + group_start_iteration_B = 0; + } else { + group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + } + + copy_tiles_and_advance(iterator_A, iterator_A_scale_bias, iterator_B, + group_start_iteration_A, + group_start_iteration_B); + + + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + elementwise_transform( + warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_A_scale_bias[(warp_mma_k + 1) % 2]); + } + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + // Inserts a fence to group cp.async instructions into stages. + mctlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages of cp.async have committed + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.advance(); + iterator_A_scale_bias.advance(); + iterator_B.advance(); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_A_scale_bias_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_A_scale_bias_.add_tile_offset( + {0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_A_scale_bias_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + } + } + + } + + // Insert fence and wait for all outstanding cp.async operations to commit. + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_multistage.h b/csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_multistage.h new file mode 100644 index 0000000..cc2780b --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_multistage.h @@ -0,0 +1,539 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a multistage threadblock-scoped Implicit GEMM Convolution kernel. +*/ + +#pragma once + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/cache_operation.h" +#include "mctlass/gemm/threadblock/mma_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class ImplicitGemmMultistage : + public gemm::threadblock::MmaBase { +public: + ///< Base class + using Base = gemm::threadblock::MmaBase; + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Policy describing tuning details + using Policy = Policy_; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + + using ElementC = typename Policy::Operator::ElementC; + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Internal structure exposed for introspection. + struct Detail { + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + // Optional staged-accumulation (e.g., tf32x3 kernels) for improved numerical + // accuracy, where each mainloop iteration first accumulates into a temporary + // set of freshly-cleared accumulators, which are subsequently added to the + // final accumulator set. + static bool const kStagedAccumulation = arch::UseStagedAccumulation::value; + }; + + private: + + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + ImplicitGemmMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + MCTLASS_DEVICE + void copy_tiles_and_advance( + IteratorA &iterator_A, IteratorB &iterator_B, + int group_start_A = 0, int group_start_B = 0) { + + iterator_A.set_iteration_index(group_start_A * + IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + } + + iterator_B.set_iteration_index(group_start_B * + IteratorB::kAccessesPerVector); + + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + } + ++this->smem_iterator_B_; + } + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< initial value of accumulator + FragmentC const &src_accum, + ///< number of iterations per channel + int gemm_k_iterations_per_channel = 0, + ///< Imaginary strides used for planar-complex only - ignored here + int64_t imag_stride_A = 0, + int64_t imag_stride_B = 0) { + + // + // Prologue + // + + // Issue several complete stages + MCTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; + ++stage, --gemm_k_iterations) { + + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + + iterator_B.set_iteration_index(0); + this->smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + } + + ++this->smem_iterator_B_; + } + + // Move to the next stage + iterator_A.advance(); + iterator_B.advance(); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Inserts a fence to group cp.async instructions into stages. + mctlass::arch::cp_async_fence(); + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // Waits until kStages-2 stages have committed. + mctlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[2]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpTransformedFragmentA warp_transformed_frag_A[2]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + // Start issuing the first group of the next stage outside of the mainloop + copy_tiles_and_advance(iterator_A, iterator_B); + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0], + warp_loaded_frag_A[0], warp_loaded_frag_B[0]); + + // tf32x3 kernels use staging accumulation. warp_mma uses a temporary + // accumulator and this temporary accumulator is added to the final + // accumulator once in every mainloop iteration. + plus plus_accum; + + FragmentC tmp_accum; + + if (Detail::kStagedAccumulation) { + tmp_accum.clear(); + } + + // + // Mainloop + // + + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) + warp_mma.transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % 2], + warp_loaded_frag_B[warp_mma_k % 2]); + + // Issue global->shared copies for the next stage + int group_start_iteration_A, group_start_iteration_B; + + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + group_start_iteration_A = 0; + group_start_iteration_B = 0; + } else { + group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + } + + copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, + group_start_iteration_B); + + if (Detail::kStagedAccumulation) { + warp_mma( + tmp_accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + tmp_accum + ); + + if (warp_mma_k == 0) { + accum = plus_accum(accum, tmp_accum); + tmp_accum.clear(); + } + } else { + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum + ); + } + + if (warp_mma_k + 1 == Base::kWarpGemmIterations) + warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + // Inserts a fence to group cp.async instructions into stages. + mctlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages of cp.async have committed + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.advance(); + iterator_B.advance(); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + } + } + + } + + if (Detail::kStagedAccumulation) { + accum = plus_accum(accum, tmp_accum); + } + + // Insert fence and wait for all outstanding cp.async operations to commit. + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_pipelined.h b/csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_pipelined.h new file mode 100644 index 0000000..83e163e --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_pipelined.h @@ -0,0 +1,320 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/aligned_buffer.h" +#include "mctlass/numeric_conversion.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/mma_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Transformation applied to A operand + typename TransformA_ = NumericArrayConverter< + typename SmemIteratorA_::Element, + typename IteratorA_::Element, + IteratorA_::Fragment::kElements>, + /// + /// Transformation applied to A operand + typename TransformB_ = NumericArrayConverter< + typename SmemIteratorB_::Element, + typename IteratorB_::Element, + IteratorB_::Fragment::kElements>, + /// Used for partial specialization + typename Enable = bool +> +class ImplicitGemmPipelined : public gemm::threadblock::MmaBase { +public: + + ///< Base class + using Base = gemm::threadblock::MmaBase; + + using Shape = Shape_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using IteratorA = IteratorA_; ///< Iterates over tiles of A operand in global memory + using IteratorB = IteratorB_; ///< Iterates over tiles of B operand in global memory + using ElementC = ElementC_; ///< Data type of accumulator matrix + using LayoutC = LayoutC_; ///< Layout of accumulator matrix + using Policy = Policy_; ///< Policy describing tuning details + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + using TransformA = TransformA_; + using TransformB = TransformB_; + + // + // Dependent types + // + + /// Fragment of operand A loaded from global memory + using FragmentA = typename IteratorA::Fragment; + + /// Fragment of operand B loaded from global memory + using FragmentB = typename IteratorB::Fragment; + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Obtain the arch tag from the warp-level operator + using ArchTag = typename Policy::Operator::ArchTag; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + // staticaly assert kStages for MmaPipelined is two (Double-buffered pipeline) + static_assert((Base::kStages==2), "MmaPipelined requires kStages set to value 2"); + +private: + + using WarpFragmentA = typename Operator::FragmentA; + using WarpFragmentB = typename Operator::FragmentB; + +protected: + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + ImplicitGemmPipelined( + typename Base::SharedStorage &shared_storage, ///< Shared storage needed for internal use by threadblock-scoped GEMM + int thread_idx, ///< ID within the threadblock + int warp_idx, ///< ID of warp + int lane_idx ///< ID of each thread within a warp + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) { + + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + int gemm_k_iterations, ///< number of iterations of the mainloop + FragmentC &accum, ///< destination accumulator tile + IteratorA iterator_A, ///< iterator over A operand in global memory + IteratorB iterator_B, ///< iterator over B operand in global memory + FragmentC const &src_accum, ///< source accumulator tile + int gemm_k_iterations_per_channel = 0, ///< number of iterations per channel + TransformA transform_A = TransformA(), ///< transformation applied to A fragment + TransformB transform_B = TransformB()) { ///< transformation applied to B fragment + + // + // Prologue + // + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + FragmentA tb_frag_A; + FragmentB tb_frag_B; + + tb_frag_A.clear(); + tb_frag_B.clear(); + + // The last kblock is loaded in the prolog + iterator_A.load(tb_frag_A); + iterator_B.load(tb_frag_B); + + ++iterator_A; + ++iterator_B; + + this->smem_iterator_A_.store(transform_A(tb_frag_A)); + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math instructions + WarpFragmentA warp_frag_A[2]; + WarpFragmentB warp_frag_B[2]; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + Operator warp_mma; + + int smem_write_stage_idx = 1; + + // Issue loads during the first warp-level matrix multiply-add *AFTER* issuing + // shared memory loads (which have the tightest latency requirement). + + // + // Mainloop + // + + // Note: The main loop does not support Base::kWarpGemmIterations == 2. + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > 0; --gemm_k_iterations) { + // + // Loop over GEMM K dimension + // + + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if this is the last group + // as the case may be. + + if (warp_mma_k == Base::kWarpGemmIterations - 1) { + + // Write fragments to shared memory + this->smem_iterator_A_.store(transform_A(tb_frag_A)); + + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + __syncthreads(); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + // Add negative offsets to return iterators to the 'start' of the circular buffer in shared memory + if (smem_write_stage_idx == 1) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + } + else { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations, + 0}); + } + + smem_write_stage_idx ^= 1; + } + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k == 0) { + + iterator_A.load(tb_frag_A); + iterator_B.load(tb_frag_B); + + ++iterator_A; + ++iterator_B; + } + + warp_mma(accum, warp_frag_A[warp_mma_k % 2], + warp_frag_B[warp_mma_k % 2], accum); + } + } + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_wgrad_fusion_multistage.h b/csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_wgrad_fusion_multistage.h new file mode 100644 index 0000000..839b450 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/implicit_gemm_wgrad_fusion_multistage.h @@ -0,0 +1,729 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a multistage threadblock-scoped fused activation's scale+bias+relu and + Implicit GEMM Convolution kernel. + + The original implicit gemm will store out-of-bound data as zeroes in the + shared memory because zeros into the tensor core, zeroes out of the tensor + cores. The result is remained the same. When fusing scale+bias+relu + into the mainloop, it is no longer true because + + 0 x scale + bias = bias + + which is no longer always 0. So, instead of storing zeroes, this fused + kernel stores the out-of-bound data as a special NaN (0x7eff), when applying + scale+bias+relu, the code is like + + if (data == 0x7eff) + data = 0; + else + data = scale+bias+relu(data, scale, bias); + + The biggest difference compared with the fused Fprop and scale+bias+relu is + that scale and bias are loop invariant in Wgrad so that they only needs to + be loaded once before the mainloop. + + See include/mctlass/conv/warp/scale_bias_relu_transformation.h for the + elementwise computation. See include/mctlass/arch/memory_sm80.h for nan fill. + + +*/ + +#pragma once + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/cache_operation.h" +#include "mctlass/gemm/gemm.h" + +#include "mctlass/gemm/warp/scale_bias_tile_iterator.h" +#include "mctlass/conv/warp/scale_bias_relu_transform.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Element type of scale and bias vectors + typename ElementScaleBias_, + /// Layout of scale and bias vectors + typename LayoutScaleBias_, + /// Element type of scale and bias vectors + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class MmaWgradFusionBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Element type of scale and bias vectors + using ElementScaleBias = ElementScaleBias_; + + /// Layout of scale and bias vectors + using LayoutScaleBias = LayoutScaleBias_; + + ///< Policy describing tuning details + using Policy = Policy_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = mctlass::gemm::GemmShape; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = + (WarpGemm::kK / Operator::Policy::MmaShape::kK); + + /// Number of stages + static int const kStages = Stages; + + /// Tensor reference to the A operand + using TensorRefA = TensorRef; + + /// Tensor reference to the B operand + using TensorRefB = TensorRef; + + static_assert(kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + static_assert((kWarpGemmIterations % 2) == 0, + "Inner loop iteration must be an even number."); + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + class SharedStorage { + public: + // + // Type definitions + // + + /// Shape of the A matrix operand in shared memory + using ShapeA = MatrixShape; + + /// Shape of the B matrix operand in shared memory + using ShapeB = + MatrixShape; + + public: + // + // Data members + // + + /// Buffer for A operand + AlignedBuffer operand_A; + + /// Buffer for B operand + AlignedBuffer operand_B; + + public: + + // + // Methods + // + + /// Returns a layout object for the A matrix + MCTLASS_DEVICE + static typename Operator::LayoutA LayoutA() { + return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn}); + } + + /// Returns a layout object for the B matrix + MCTLASS_HOST_DEVICE + static typename Operator::LayoutB LayoutB() { + return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); + } + + /// Returns a TensorRef to the A operand + MCTLASS_HOST_DEVICE + TensorRefA operand_A_ref() { + return TensorRefA{operand_A.data(), LayoutA()}; + } + + /// Returns a TensorRef to the B operand + MCTLASS_HOST_DEVICE + TensorRefB operand_B_ref() { + return TensorRefB{operand_B.data(), LayoutB()}; + } + }; + + protected: + + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A operand from shared memory + typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaWgradFusionBase( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx), + warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) {} +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// Iterates over vectors of scale and bias vector in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorScaleBias_, + /// Iterates over vectors of scale and bias vector i + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class ImplicitGemmWgradFusionMultistage + : public MmaWgradFusionBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Iterates over tiles of the scale and bias vectors in global memory + using IteratorScaleBias = IteratorScaleBias_; + ///< Policy describing tuning details + using Policy = Policy_; + ///< Base class + using Base = MmaWgradFusionBase; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + + using ElementC = typename Policy::Operator::ElementC; + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Internal structure exposed for introspection. + struct Detail { + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + static int const kBBufferSize = + ((sizeof(typename Operator::ElementC) == 4) && + ((platform::is_same::value && + platform::is_same::value)) && + (Operator::Shape::kM >= 64 && Operator::Shape::kN >= 64)) + ? 1 + : 2; + }; + + private: + + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpLoadedFragmentScaleBias = typename IteratorScaleBias::Fragment; + + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + int warp_idx_m_; + + int warp_idx_n_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + ImplicitGemmWgradFusionMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) { + + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + warp_idx_m_ = warp_idx_mn % Base::WarpCount::kM; + warp_idx_n_ = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m_, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n_}); + } + + MCTLASS_DEVICE + void copy_tiles_and_advance(IteratorA &iterator_A, + IteratorB &iterator_B, + int group_start_A = 0, int group_start_B = 0) { + + iterator_A.set_iteration_index(group_start_A); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + + ++this->smem_iterator_A_; + } + } + + iterator_B.set_iteration_index(group_start_B); + + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / 8; + + // Uses nan fill for out of bound data + mctlass::arch::cp_async_nan( + dst_ptr, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + ++this->smem_iterator_B_; + } + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< iterator over scale and bias vectors in global memory + IteratorScaleBias iterator_B_scale_bias, + ///< initial value of accumulator + FragmentC const &src_accum, + ///< number of iterations per channel + int gemm_k_iterations_per_channel = 0, + ///< Imaginary strides used for planar-complex only - ignored here + int64_t imag_stride_A = 0, + int64_t imag_stride_B = 0) { + + // + // Prologue + // + + WarpLoadedFragmentScaleBias warp_loaded_frag_B_scale_bias; + iterator_B_scale_bias.add_tile_offset({0, warp_idx_n_}); + iterator_B_scale_bias.load(warp_loaded_frag_B_scale_bias); + + // Issue several complete stages + MCTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; + ++stage, --gemm_k_iterations) { + + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + ++this->smem_iterator_A_; + } + + iterator_B.set_iteration_index(0); + this->smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / 8; + + // Uses Nan fill for out of bound data + mctlass::arch::cp_async_nan( + dst_ptr, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + ++this->smem_iterator_B_; + } + + // Move to the next stage + iterator_A.advance(); + iterator_B.advance(); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Inserts a fence to group cp.async instructions into stages. + mctlass::arch::cp_async_fence(); + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // Waits until kStages-2 stages have committed. + mctlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[Detail::kBBufferSize]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpTransformedFragmentA warp_transformed_frag_A[Detail::kBBufferSize]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + mctlass::conv::warp::WgradScaleBiasReluTransform + elementwise_transform; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + // Start issuing the first group of the next stage outside of the mainloop + copy_tiles_and_advance(iterator_A, iterator_B); + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0], + warp_loaded_frag_A[0], warp_loaded_frag_B[0]); + + elementwise_transform(warp_transformed_frag_B[0], + warp_loaded_frag_B_scale_bias); + + // + // Mainloop + // + + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + if (Detail::kBBufferSize == 2) { + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % Detail::kBBufferSize]); + ++this->warp_tile_iterator_A_; + } + + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) { + warp_mma.transform(warp_transformed_frag_A[warp_mma_k % Detail::kBBufferSize], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % Detail::kBBufferSize], + warp_loaded_frag_B[warp_mma_k % 2]); + + elementwise_transform(warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_B_scale_bias); + } + + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % Detail::kBBufferSize], + warp_transformed_frag_B[warp_mma_k % 2], + accum + ); + + if (Detail::kBBufferSize == 1) { + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + ++this->warp_tile_iterator_A_; + + } + + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % Detail::kBBufferSize], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % Detail::kBBufferSize], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + elementwise_transform( + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_B_scale_bias); + } + + // Issue global->shared copies for the next stage + int group_start_iteration_A, group_start_iteration_B; + + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + group_start_iteration_A = 0; + group_start_iteration_B = 0; + } else { + group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + } + + copy_tiles_and_advance(iterator_A, iterator_B, + group_start_iteration_A, + group_start_iteration_B); + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + // Inserts a fence to group cp.async instructions into stages. + mctlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages of cp.async have committed + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.advance(); + iterator_B.advance(); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + } + } + + } + + // Insert fence and wait for all outstanding cp.async operations to commit. + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/predicated_scale_bias_vector_access_iterator.h b/csrc/mctlass/include/mctlass/conv/threadblock/predicated_scale_bias_vector_access_iterator.h new file mode 100644 index 0000000..9b36733 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/predicated_scale_bias_vector_access_iterator.h @@ -0,0 +1,471 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates calculating the address and predicates to the load of scale and bias vectors. + + This iterator uses masks to guard out-of-bounds accesses. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/mctlass.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedScaleBiasVectorAccessIterator +/// +template +class PredicatedScaleBiasVectorAccessIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for fprop pitch-linear data. +/// +template +class PredicatedScaleBiasVectorAccessIterator { + public: + + using ThreadblockShape = ThreadblockShape_; + using Element = Element_; + using Layout = layout::PitchLinear; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + static int const kElementsPerAccess = 128 / sizeof_bits::value; + static int const kThreads = ThreadblockShape::kContiguous / kElementsPerAccess; + + using AccessType = AlignedArray; + + using Params = PredicatedScaleBiasVectorAccessIteratorParams; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + /// Parameters object with precomputed internal state + Params const ¶ms_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + + int problem_size_trs; + int problem_size_c; + int filter_trs_; + + TensorCoord thread_offset_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Extent of tensor + Conv2dProblemSize const &problem_size, + /// Pointer to the start of the scale vector + ConstPointer scale_pointer, + /// Pointer to the start of the bias vector + ConstPointer bias_pointer, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : params_(params), + problem_size_trs(problem_size.R * problem_size.S), + problem_size_c(problem_size.C), + filter_trs_(0) { + pointer_ = (thread_id < kThreads) + ? reinterpret_cast( + const_cast(scale_pointer)) + : reinterpret_cast( + const_cast(bias_pointer)); + + // Per-thread offset in logical coordinates of tensor + int thread_base = (thread_id < kThreads) ? 0 : kThreads; + + thread_offset_ = + threadblock_offset + + TensorCoord((thread_id - thread_base) * kElementsPerAccess, 0); + + set_iteration_index(0); + } + + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Extent of tensor + Conv3dProblemSize const &problem_size, + /// Pointer to the start of the scale vector + ConstPointer scale_pointer, + /// Pointer to the start of the bias vector + ConstPointer bias_pointer, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : params_(params), + problem_size_trs(problem_size.T * problem_size.R * problem_size.S), + problem_size_c(problem_size.C), + filter_trs_(0) { + pointer_ = (thread_id < kThreads) + ? reinterpret_cast( + const_cast(scale_pointer)) + : reinterpret_cast( + const_cast(bias_pointer)); + + // Per-thread offset in logical coordinates of tensor + int thread_base = (thread_id < kThreads) ? 0 : kThreads; + + thread_offset_ = + threadblock_offset + + TensorCoord((thread_id - thread_base) * kElementsPerAccess, 0); + + set_iteration_index(0); + } + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Extent of tensor + Conv2dProblemSize const &problem_size, + /// Pointer to start of scale vector + ConstPointer scale_pointer, + /// Pointer to start of scale vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id) + : PredicatedScaleBiasVectorAccessIterator(params, problem_size, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Extent of tensor + Conv3dProblemSize const &problem_size, + /// Pointer to start of scale vector + ConstPointer scale_pointer, + /// Pointer to start of scale vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id) + : PredicatedScaleBiasVectorAccessIterator(params, problem_size, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) {} + + /// Advances an iterator along logical dimensions of matrix in units of whole threadblock tiles + MCTLASS_DEVICE + void add_tile_offset( + TensorCoord const &tile_offset) { + thread_offset_ = + thread_offset_ + + TensorCoord(ThreadblockShape::kContiguous * tile_offset.contiguous(), 0); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + + return reinterpret_cast( + pointer_ + + (thread_offset_.contiguous() * sizeof_bits::value / 8)); + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator &operator++() { + return *this; + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + void advance() { + // moves to the next tile + ++filter_trs_; + if (filter_trs_ == problem_size_trs) { + filter_trs_ = 0; + add_tile_offset(TensorCoord(1, 0)); + } + } + + /// Increment and return an instance to self. + MCTLASS_DEVICE + PredicatedScaleBiasVectorAccessIterator operator++(int) { + PredicatedScaleBiasVectorAccessIterator self(*this); + operator++(); + return self; + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + uint32_t enabled = 0; + +//#if defined(_MSC_VER) || (__CUDACC_VER_MAJOR__ < 11) +#if 1 + enabled = threadIdx.x < kThreads * 2; +#else 0 + asm volatile( + "{\n" + " .reg .u32 tid_reg;\n" + " .reg .pred p;\n" + " mov.u32 tid_reg, %%tid.x;\n" + " setp.lt.u32 p, tid_reg, %1;\n" + " selp.u32 %0, 1, 0, p;\n" + "}\n" : "+r"(enabled) :"n"(kThreads * 2)); +#endif + + return ((thread_offset_.contiguous() < problem_size_c) && enabled); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedScaleBiasVectorAccessIterator { + public: + + using ThreadblockShape = ThreadblockShape_; + using Element = Element_; + using Layout = layout::RowMajor; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedScaleBiasVectorAccessIterator< + layout::PitchLinearShape, + Element, + layout::PitchLinear>; + + using AccessType = typename UnderlyingIterator::AccessType; + static int const kElementsPerAccess = UnderlyingIterator::kElementsPerAccess; + + using Params = PredicatedScaleBiasVectorAccessIteratorParams; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + ///< Precomputed parameters object + Params const ¶ms, + ///< Extent of tensor + Conv2dProblemSize const &problem_size, + ///< Pointer to the start of the scale vector + ConstPointer scale_pointer, + ///< Pointer to the start of the bias vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params, problem_size, scale_pointer, bias_pointer, + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row())) {} + + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + ///< Precomputed parameters object + Params const ¶ms, + ///< Extent of tensor + Conv3dProblemSize const &problem_size, + ///< Pointer to the start of the scale vector + ConstPointer scale_pointer, + ///< Pointer to the start of the bias vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params, problem_size, scale_pointer, bias_pointer, + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Conv2dProblemSize const &problem_size, ///< Extent of tensor + ConstPointer scale_pointer, ///< Pointer to the start of the scale vector + ConstPointer bias_pointer, ///< Pointer to the start of the bias vector + int thread_id ///< ID of each participating thread + ) + : PredicatedScaleBiasVectorAccessIterator(params, problem_size, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Conv3dProblemSize const &problem_size, ///< Extent of tensor + ConstPointer scale_pointer, ///< Pointer to the start of the scale vector + ConstPointer bias_pointer, ///< Pointer to the start of the bias vector + int thread_id ///< ID of each participating thread + ) + : PredicatedScaleBiasVectorAccessIterator(params, problem_size, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// threadblock tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator operator++(int) { + PredicatedScaleBiasVectorAccessIterator self(*this); + operator++(); + return self; + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + void advance() { + iterator_.advance(); + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/predicated_scale_bias_vector_iterator.h b/csrc/mctlass/include/mctlass/conv/threadblock/predicated_scale_bias_vector_iterator.h new file mode 100644 index 0000000..54f7104 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/predicated_scale_bias_vector_iterator.h @@ -0,0 +1,371 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates calculating the address and predicates to the load of scale and bias vectors. + + This iterator uses masks to guard out-of-bounds accesses. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/mctlass.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedScaleBiasVectorIterator +/// +template +class PredicatedScaleBiasVectorIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for wgrad pitch-linear data. +/// +template +class PredicatedScaleBiasVectorIterator { + public: + + using WarpShape = WarpShape_; + using Element = Element_; + using Layout = layout::PitchLinear; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + static int const kElementsPerAccess = 1; + + using AccessType = AlignedArray; + + static int const kIterations = WarpShape::kContiguous / 8; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array<__half2, 2 * kIterations * kElementsPerAccess>; + + /// Parameters object is precomputed state and is host-constructible + using Params = Conv2dWgradActivationIteratorOptimizedParams; + + private: + // + // Data members + // + + /// Parameters object with precomputed internal state + Params const ¶ms_; + + /// Internal pointer to first access of tile + ConstPointer scale_pointer_; + ConstPointer bias_pointer_; + + /// Size of tensor + Conv2dProblemSize problem_size_; + + int32_t thread_offset_; + + // Channel dimension in contiguous dimension stays constant for each gemm_iteration_k + int32_t filter_c_[kIterations]; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Extent of tensor + Conv2dProblemSize const &problem_size, + /// Pointer to the start of the scale vector + ConstPointer scale_pointer, + /// Pointer to the start of the bias vector + ConstPointer bias_pointer, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : params_(params), + problem_size_(problem_size), + scale_pointer_(scale_pointer), + bias_pointer_(bias_pointer) { + + thread_offset_ = threadblock_offset.contiguous() + (thread_id % 32) / 4; + } + + /// Construct a PredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Extent of tensor + Conv2dProblemSize const &problem_size, + /// Pointer to start of scale vector + ConstPointer scale_pointer, + /// Pointer to start of scale vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id) + : PredicatedScaleBiasVectorIterator(params, problem_size, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + /// Advances an iterator along logical dimensions of matrix in units of whole warp tiles + MCTLASS_DEVICE + void add_tile_offset( + TensorCoord const &tile_offset) { + + thread_offset_ += (WarpShape::kContiguous * tile_offset.contiguous()); + + MCTLASS_PRAGMA_UNROLL + for(int c = 0; c < kIterations; ++c) { + int rsc_offset = thread_offset_ + c * 8; + + int residual, tmp; + params_.sc_divmod(tmp, residual, rsc_offset); + params_.c_divmod(tmp, filter_c_[c], residual); + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + frag.fill(__float2half2_rn(0.0f)); + __half2 *frag_ptr = reinterpret_cast<__half2 *>(&frag); + + // load scale + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + + mctlass::arch::global_load< + __half, + sizeof(AccessType) + >( + frag_ptr[c * 2].x, + scale_pointer_ + filter_c_[c], + true + ); + } + + // load bias + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + + mctlass::arch::global_load< + __half, + sizeof(AccessType) + >( + frag_ptr[c * 2 + 1].x, + bias_pointer_ + filter_c_[c], + true + ); + } + + // duplicate scale + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + frag_ptr[c * 2].y = frag_ptr[c * 2].x; + } + + // duplicate bias + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + frag_ptr[c * 2 + 1].y = frag_ptr[c * 2 + 1].x; + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedScaleBiasVectorIterator { + public: + + using WarpShape = WarpShape_; + using Element = Element_; + using Layout = layout::RowMajor; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedScaleBiasVectorIterator< + layout::PitchLinearShape, + Element, + layout::PitchLinear>; + + using AccessType = typename UnderlyingIterator::AccessType; + static int const kElementsPerAccess = UnderlyingIterator::kElementsPerAccess; + using Fragment = typename UnderlyingIterator::Fragment; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedScaleBiasVectorIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default ctor + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Conv2dProblemSize const &problem_size, Layout const &layout) + : params_(problem_size, layout::TensorNHWC(0, 0, 0)){}; + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorIterator( + ///< Precomputed parameters object + Params const ¶ms, + ///< Extent of tensor + Conv2dProblemSize const &problem_size, + ///< Pointer to the start of the scale vector + ConstPointer scale_pointer, + ///< Pointer to the start of the bias vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params.params_, problem_size, scale_pointer, bias_pointer, + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorIterator( + Params const ¶ms, ///< Precomputed parameters object + Conv2dProblemSize const &problem_size, ///< Extent of tensor + ConstPointer scale_pointer, ///< Pointer to the start of the scale vector + ConstPointer bias_pointer, ///< Pointer to the start of the bias vector + int thread_id ///< ID of each participating thread + ) + : PredicatedScaleBiasVectorIterator(params, problem_size, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// threadblock tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + iterator_.load(frag); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace conv +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/conv/threadblock/threadblock_swizzle.h b/csrc/mctlass/include/mctlass/conv/threadblock/threadblock_swizzle.h new file mode 100644 index 0000000..0bb402f --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/threadblock/threadblock_swizzle.h @@ -0,0 +1,193 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Implements several possible threadblock-swizzling functions mapping blockIdx to + Convolution problems. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/platform/platform.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// +MCTLASS_HOST_DEVICE +static int get_strided_dgrad_tile_m( + mctlass::conv::Conv2dProblemSize const &problem_size, + int tile_size_m) { + + // CTAs in M dimension per starting filter position + int tile_m_per_filter = strided_dgrad_tile_m_per_filter(problem_size, tile_size_m); + + // Inflate number of CTAs in M dimension to cover every strating filter position even those that + // may fall out of valid MMA (Dy * w) but are needed to apply epilogue (beta * Dx_source) + // and point-wise fusion + int tile_m = tile_m_per_filter * int(problem_size.stride().product()); + + // There is a possible performance optimization here that leads up to 2x speeds than the current + // MCTLASS strided dgrad performance for stride > filter, i.e., stride={2x2} and filter={1x1}) + // + // * Optimization * + // Only launch CTAs in M dimension which contribute to a row in Dx output + // + // + // * Constraints * + // (A) stride <= filter, for example, stride={2x2} and filter={3x3}: + // - (A.1): There are no constraints for this case and the optimization does + // affect this case functionality or performance. + // (B) stride > filter, for example, stride={2x2} and filter={1x1}: + // - (B.1): Dx output tensor should be zero initialized + // - (B.2): The kernel epilogue cannot apply beta. Thus, beta should be zero + + return tile_m; +} +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Threadblock swizzling function for strided dgrad convolution +struct StridedDgradHorizontalThreadblockSwizzle : + public gemm::threadblock::GemmHorizontalThreadblockSwizzle { + + using Base = gemm::threadblock::GemmHorizontalThreadblockSwizzle; + + MCTLASS_HOST_DEVICE + StridedDgradHorizontalThreadblockSwizzle() { } + + /// Returns the shape of the problem in units of logical tiles + /// For ImplicitGemmConvolution Conv2d problem size: conv_operator(NPQK, NHWC, KRSC) + MCTLASS_HOST_DEVICE + gemm::GemmCoord get_tiled_shape( + mctlass::conv::Operator conv_operator, + mctlass::conv::Conv2dProblemSize const &problem_size, + gemm::GemmCoord tile_size, + int split_k_slices) const { + + gemm::GemmCoord implicit_gemm_problem_size = + mctlass::conv::implicit_gemm_problem_size(conv_operator, problem_size); + + // compute number of tiles in m dimension + int tile_m = get_strided_dgrad_tile_m(problem_size, tile_size.m()); + + // compute number of tiles in n dimension + int tile_n = (implicit_gemm_problem_size.n() + tile_size.n() - 1) / tile_size.n(); + + return gemm::GemmCoord( + tile_m, + tile_n, + split_k_slices); + } + + /// Returns the shape of the problem in units of logical tiles + /// For GEMM problem size (MxNxK) (Do not use base class get_tiled_shape()) + private: + using Base::get_tiled_shape; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Threadblock swizzling function for strided dgrad convolution +template +struct StridedDgradIdentityThreadblockSwizzle : + public gemm::threadblock::GemmIdentityThreadblockSwizzle { + + using Base = gemm::threadblock::GemmIdentityThreadblockSwizzle; + + MCTLASS_HOST_DEVICE + StridedDgradIdentityThreadblockSwizzle() { } + + /// Returns the shape of the problem in units of logical tiles + /// For ImplicitGemmConvolution Conv2d problem size: conv_operator(NPQK, NHWC, KRSC) + MCTLASS_HOST_DEVICE + gemm::GemmCoord get_tiled_shape( + mctlass::conv::Operator conv_operator, + mctlass::conv::Conv2dProblemSize const &problem_size, + gemm::GemmCoord tile_size, + int split_k_slices) const { + + gemm::GemmCoord implicit_gemm_problem_size = + mctlass::conv::implicit_gemm_problem_size(conv_operator, problem_size); + + // compute number of tiles in m dimension + int tile_m = get_strided_dgrad_tile_m(problem_size, tile_size.m()); + + // compute number of tiles in n dimension + int tile_n = (implicit_gemm_problem_size.n() + tile_size.n() - 1) / tile_size.n(); + + return gemm::GemmCoord( + tile_m, + tile_n, + split_k_slices); + } + + /// Returns the shape of the problem in units of logical tiles + /// For GEMM problem size (MxNxK) (Do not use base class get_tiled_shape()) + private: + using Base::get_tiled_shape; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Threadblock swizzling function for GEMMs +template +struct DepthwiseDirect2dConvIdentityThreadblockSwizzle + : public gemm::threadblock::GemmIdentityThreadblockSwizzle { + MCTLASS_HOST_DEVICE + DepthwiseDirect2dConvIdentityThreadblockSwizzle() {} + + /// Returns the shape of the problem in units of logical tiles + MCTLASS_HOST_DEVICE + gemm::GemmCoord get_tiled_shape(mctlass::conv::Operator conv_operator, + mctlass::conv::Conv2dProblemSize const &problem_size, + gemm::GemmCoord tile_size, + int split_k_slices) const { + + gemm::GemmCoord implicit_gemm_problem_size = + mctlass::conv::implicit_gemm_problem_size(conv_operator, problem_size); + + return gemm::GemmCoord(1, + (implicit_gemm_problem_size.n() + tile_size.n() - 1) / tile_size.n(), + split_k_slices); + } +}; + +} // namespace threadblock +} // namespace conv +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/conv/warp/mma_depthwise_simt.h b/csrc/mctlass/include/mctlass/conv/warp/mma_depthwise_simt.h new file mode 100644 index 0000000..8c2ae51 --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/warp/mma_depthwise_simt.h @@ -0,0 +1,380 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/thread/mma.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/thread/depthwise_mma.h" + + +#include "mctlass/gemm/warp/mma_simt_tile_iterator.h" +#include "mctlass/gemm/warp/mma_simt_policy.h" + +#include "mctlass/gemm/warp/mma_simt.h" +#include "mctlass/conv/warp/mma_depthwise_simt_tile_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK = 1, + /// Complex transformation on operand A + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex transformation on operand B + ComplexTransform TransformB = ComplexTransform::kNone, + /// Used for partial specialization + typename Enable = bool> +class MmaDepthwiseSimt + : public mctlass::gemm::warp:: + MmaSimt { + using Base = mctlass::gemm::warp:: + MmaSimt; + +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of multiplicand A + using ElementA = ElementA_; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of multiplicand B + using ElementB = ElementB_; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of accumulator matrix C + using ElementC = ElementC_; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassSimt; + + /// Hard-coded for now + using ArchTag = arch::Sm50; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = TransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = TransformB; + +public: + + /// Iterates over the B operand in memory + using IteratorB = mctlass::conv::warp::DepthwiseMmaSimtTileIterator< + MatrixShape, + mctlass::gemm::Operand::kB, + ElementB, + LayoutB, + Policy, + PartitionsK, + Shape::kK + >; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentB = FragmentB; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaDepthwiseSimt():Base() {} +}; + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Shape of filter shape per threadblock - concept: gemm::GemmShape + typename FilterShape_, + /// Shape of the output tile computed by thread- concept: conv::TensorNHWCShape<> + typename ThreadOutputShape_, + /// Shape of the output tile computed by threadblock - concept: conv::TensorNHWCShape<> + typename ThreadBlockOutputShape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Iterator algo type + conv::IteratorAlgorithm IteratorAlgorithm_ = IteratorAlgorithm::kAnalytic, + /// Stride ( MatrixShape ) + typename StrideShape_ = mctlass::MatrixShape<-1, -1>, + /// Dilation ( MatrixShape ) + typename DilationShape_ = mctlass::MatrixShape<-1, -1>, + /// Activation Shape loaded by threadblock + typename ActivationShape_ = mctlass::conv::TensorNHWCShape<-1,-1,-1,-1>, + /// Number of partitions along K dimension + int PartitionsK = 1, + /// Complex transformation on operand A + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex transformation on operand B + ComplexTransform TransformB = ComplexTransform::kNone, + /// Used for partial specialization + typename Enable = bool> +class MmaDepthwiseDirectConvSimt { + public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Shape of filter shape per threadblock - concept: gemm::GemmShape + using FilterShape = FilterShape_; + + /// Shape of the output tile computed by thread- concept: conv::TensorNHWCShape<> + using ThreadOutputShape = ThreadOutputShape_; + + /// Shape of the output tile computed by threadblock - concept: conv::TensorNHWCShape<> + using ThreadBlockOutputShape = ThreadBlockOutputShape_; + + /// Data type of multiplicand A + using ElementA = ElementA_; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of multiplicand B + using ElementB = ElementB_; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of accumulator matrix C + using ElementC = ElementC_; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Iterator algo type + static conv::IteratorAlgorithm const IteratorAlgorithm = IteratorAlgorithm_; + + /// Stride ( MatrixShape ) + using StrideShape = StrideShape_; + + /// Dilation ( MatrixShape ) + using DilationShape = DilationShape_; + + /// Activation Shape loaded by threadblock + using ActivationShape = ActivationShape_; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassSimt; + + /// Hard-coded for now + using ArchTag = arch::Sm50; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = TransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = TransformB; + + static constexpr bool use_dp4a = (platform::is_same< layout::ColumnMajorInterleaved<4>, LayoutA>::value || + platform::is_same< layout::RowMajorInterleaved<4>, LayoutA >::value) && + platform::is_same< ElementA, int8_t >::value && + platform::is_same< ElementB, int8_t >::value; + + using dp4a_type = typename platform::conditional< use_dp4a , int8_t, bool >::type; + + /// Thread-level matrix multiply accumulate operator + using ThreadMma = mctlass::conv::thread::DepthwiseDirectConvElementwiseInnerProduct< + mctlass::gemm::GemmShape< + Shape::kM / Policy::WarpShape::kRow, // number of output pixels proccessed per thread + Shape::kN / Policy::WarpShape::kColumn, // number of channels proccessed per thread + 1>, + ElementA, + ElementB, + ElementC, + arch::OpMultiplyAdd, + dp4a_type + >; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename ThreadMma::ArchMmaOperator; + + /// Indicates math operator + using MathOperator = typename ArchMmaOperator::Operator; + + /// Shape of the underlying instruction + using InstructionShape = mctlass::gemm::GemmShape<1,1,use_dp4a ? 4 : 1>; + +public: + + /// Iterates over the A operand in memory + using IteratorA = mctlass::conv::warp::DepthwiseDirect2dConvSimtTileIterator< + MatrixShape, // per warp + FilterShape, + ThreadOutputShape, + ThreadBlockOutputShape, + mctlass::gemm::Operand::kA, + ElementA, + Policy, + IteratorAlgorithm, + StrideShape, + DilationShape, + ActivationShape, + PartitionsK, + Shape::kK + >; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentA = FragmentA; + + /// Iterates over the B operand in memory + using IteratorB = mctlass::gemm::warp::MmaSimtTileIterator< + MatrixShape<1, Shape::kN>, + mctlass::gemm::Operand::kB, + ElementB, + LayoutB, + Policy, + PartitionsK, + Shape::kK + >; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentB = FragmentB; + + /// Iterates over the C operand in memory + using IteratorC = mctlass::gemm::warp::MmaSimtTileIterator< + MatrixShape, + mctlass::gemm::Operand::kC, + ElementC, + LayoutC, + Policy + >; + + /// Storage for C tile + using FragmentC = typename ThreadMma::FragmentC; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaDepthwiseDirectConvSimt() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &d, + FragmentA a, + FragmentB b, + FragmentC const &c, int group_idx = 0) const { + + ThreadMma mma; + + mma(d, a, b, c); + } + + /// Transform the mma operands to the required types + MCTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + dst_A = A; + dst_B = B; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace conv +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/conv/warp/mma_depthwise_simt_tile_iterator.h b/csrc/mctlass/include/mctlass/conv/warp/mma_depthwise_simt_tile_iterator.h new file mode 100644 index 0000000..fea37eb --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/warp/mma_depthwise_simt_tile_iterator.h @@ -0,0 +1,862 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Describes the lane policy used by warp-level matrix multiply operators targeting SIMT + instructions +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/conv/convolution.h" + +#include "mctlass/arch/memory_sm75.h" + +#include "mctlass/layout/matrix.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma_simt_policy.h" +#include "mctlass/gemm/warp/mma_simt_tile_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Iterates over operands to warp-level matrix multiply operations targeting SIMT instructions +/// +/// concept: MutableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Operand identity + mctlass::gemm::Operand Operand, + /// Data type of A elements + typename Element_, + /// Layout of operand + typename Layout_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Number of partitions along K dimension - used in sliced-K + int PartitionsK = 1, + /// Group Size along kPartition - used in sliced-K + int PartitionGroupSize = 1 +> +class DepthwiseMmaSimtTileIterator; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization for B operands of row-major layouts +/// +/// Concept: MutableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of A elements + typename Element_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK, + /// Group Size along kPartition - used in sliced-K + int PartitionGroupSize> +class DepthwiseMmaSimtTileIterator + : public mctlass::gemm::warp::MmaSimtTileIterator { + + using Base = mctlass::gemm::warp::MmaSimtTileIterator; + public: + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static mctlass::gemm::Operand const kOperand = mctlass::gemm::Operand::kB; + + /// Element type + using Element = Element_; + + /// Layout of policy + using Layout = layout::RowMajor; + + /// Decomposition of elements among threads + using Policy = Policy_; + + /// TensorRef type for loading element from a tensor + using TensorRef = typename Base::TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Thread-level shape of a fragment + using ThreadShape = typename Base::ThreadShape; + + /// Number of individual loads + using Iterations = typename Base::Iterations; + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + + static_assert(Policy::LaneMmaShape::kN == 1, "Each thread should be 1 element per LDS along the k-dim"); + +private: + + MatrixCoord lane_offset_; + int channel_idx_; + int base_channel_idx_; + int warps_n_; + + public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + DepthwiseMmaSimtTileIterator():Base() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + DepthwiseMmaSimtTileIterator( + TensorRef ref, + int lane_id + ) : Base(ref, lane_id) { + + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + warps_n_ = -1; + channel_idx_ = 0; + base_channel_idx_ = 0; + lane_offset_ = lane_layout.inverse(lane_id) * MatrixCoord(0, Policy::LaneMmaShape::kN); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + DepthwiseMmaSimtTileIterator &add_tile_offset(TensorCoord const &coord) { + + if(warps_n_ == -1){ + warps_n_ = coord.column(); + } + + Base::add_tile_offset(coord); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. (vector loads) + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + Array *dst_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kRow; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Iterations::kColumn; ++n) { + + void const *ptr = this->ref_.data() + + this->ref_.offset({-(channel_idx_ - base_channel_idx_), + n * Policy::WarpShape::kColumn}) + + pointer_offset / Policy::LaneMmaShape::kN; + + // Base_k of a warp + Base_k of current threads. + int thread_k_base_idx = + warps_n_ * Shape::kColumn / Policy::LaneMmaShape::kN + lane_offset_.column(); + + if (channel_idx_ + k == thread_k_base_idx + n * Policy::WarpShape::kColumn) { + // Depthwise kernel would only do computation when channel == k. + // Loads an element when the current computation channel == the k corresponding to this thread. + arch::shared_load(dst_ptr[n + k * Iterations::kColumn], ptr); + } else { + // Reduce SMEM load + dst_ptr[n + k * Iterations::kColumn].fill(Element(0)); + } + } + } + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + if(k_group % PartitionGroupSize == 0 && k_group != 0){ + base_channel_idx_ = k_group; + } + channel_idx_ = k_group; + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Size of filter (concept: gemm::GemmShape) + typename FilterShape_, + /// Size of the matrix to load (concept: MatrixShape) + typename ThreadOutputShape_, + /// Size of the matrix to load (concept: MatrixShape) + typename ThreadBlockOutputShape_, + /// Operand identity + mctlass::gemm::Operand Operand, + /// Data type of A elements + typename Element_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Iterator algo type + conv::IteratorAlgorithm IteratorAlgorithm = IteratorAlgorithm::kAnalytic, + /// Stride ( MatrixShape ) + typename StrideShape = mctlass::MatrixShape<-1, -1>, + /// Dilation ( MatrixShape ) + typename DilationShape = mctlass::MatrixShape<-1, -1>, + /// Activation Shape loaded by threadblock + typename ActivationShape = mctlass::conv::TensorNHWCShape<-1,-1,-1,-1>, + /// Number of partitions along K dimension - used in sliced-K + int PartitionsK = 1, + /// Group Size along kPartition - used in sliced-K + int PartitionGroupSize = 1> +class DepthwiseDirect2dConvSimtTileIterator; + + +/// Specialization for A operands of row-major layouts +/// +/// Concept: MutableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Size of filter (concept: gemm::GemmShape) + typename FilterShape_, + /// Size of the matrix to load (concept: TensorNHWC) + typename ThreadOutputShape_, + /// Size of the matrix to load (concept: TensorNHWC) + typename ThreadBlockOutputShape_, + /// Data type of A elements + typename Element_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Iterator algo type + conv::IteratorAlgorithm IteratorAlgorithm, + /// Stride ( MatrixShape ) + typename StrideShape, + /// Dilation ( MatrixShape ) + typename DilationShape, + /// Activation Shape loaded by threadblock + typename ActivationShape, + /// Number of partitions along K dimension - used in sliced-K + int PartitionsK, + /// Group Size along kPartition - used in sliced-K + int PartitionGroupSize> +class DepthwiseDirect2dConvSimtTileIterator { + public: + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Shape of filter (concept: gemm::GemmShape) + using FilterShape = FilterShape_; + + /// Shape of tile to load (concept: TensorNHWC) + using ThreadOutputShape = ThreadOutputShape_; + + /// Shape of tile to load (concept: TensorNHWC) + using ThreadBlockOutputShape = ThreadBlockOutputShape_; + + /// Operand tag + static mctlass::gemm::Operand const kOperand = mctlass::gemm::Operand::kA; + + /// Element type + using Element = Element_; + + /// Layout of policy + using Layout = layout::RowMajor; + + /// Decomposition of elements among threads + using Policy = Policy_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + // + // Derived quantities + // + + static_assert(!(Shape::kRow % Policy::WarpShape::kRow), + "The warp-level GEMM M size must be divisible by the number of threads arranged along the M dimension."); + + static_assert(Shape::kRow > 0, "Shape::kRow must be greater than zero."); + static_assert(Shape::kColumn > 0, "Shape::kColumn must be greater than zero."); + static_assert(Policy::WarpShape::kRow > 0, "Policy::WarpShape::kRow must be greater than zero."); + static_assert(Shape::kRow / Policy::WarpShape::kRow > 0, "Shape::kRow / Policy::WarpShape::kRow must be greater than zero."); + +// Thread-level shape of a fragment + using ThreadShape = MatrixShape< + ThreadOutputShape::kNHW, // Output tile shape Computed by current threads + ThreadOutputShape::kC + >; + + static_assert(!(ThreadShape::kColumn % Policy::LaneMmaShape::kN), + "Thread-level GEMM must be divisible by Policy::LaneMmaShape."); + + /// Number of individual loads + using Iterations = MatrixShape< + ThreadShape::kRow, + ThreadShape::kColumn / Policy::LaneMmaShape::kN + >; + + using ThreadTileCount = MatrixShape< + ThreadBlockOutputShape::kH / ThreadOutputShape::kH, + ThreadBlockOutputShape::kW / ThreadOutputShape::kW + >; + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + +protected: + + /// Internal reference + mctlass::TensorRef, layout::RowMajor> ref_; + + int activation_offset[ThreadOutputShape::kH][ThreadOutputShape::kW][Iterations::kColumn]; + int iterator_r_; + int iterator_s_; + int iterator_offset_; + + int inc_next_s_ ; + int inc_next_r_ ; + + MatrixCoord lane_offset_; +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + DepthwiseDirect2dConvSimtTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + DepthwiseDirect2dConvSimtTileIterator( + TensorRef ref, + int lane_id + ) { + + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + // Set channel offset + lane_offset_ = lane_layout.inverse(lane_id) * MatrixCoord(0, Policy::LaneMmaShape::kN); + + ref.add_coord_offset(lane_offset_); + + ref_.reset(reinterpret_cast *>(ref.data()), + ref.stride(0) / Policy::LaneMmaShape::kN); + + iterator_r_ = 0; + iterator_s_ = 0; + iterator_offset_ = 0; + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + DepthwiseDirect2dConvSimtTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + template + MCTLASS_HOST_DEVICE + void setup_initial_status(Params const& params) { + + inc_next_s_ = params.inc_next[0]; + inc_next_r_ = params.inc_next[1]; + + // Get base HW offset of current threads + int threadgroup = threadIdx.x / (ThreadBlockOutputShape::kC / ThreadOutputShape::kC); + int base_p_ = + (threadgroup / (ThreadTileCount::kColumn)) * ThreadOutputShape::kH; + int base_q_ = + (threadgroup % (ThreadTileCount::kColumn)) * ThreadOutputShape::kW; + + MCTLASS_PRAGMA_UNROLL + for (int p = 0; p < ThreadOutputShape::kH; ++p) { + MCTLASS_PRAGMA_UNROLL + for (int q = 0; q < ThreadOutputShape::kW; ++q) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < Iterations::kColumn; ++col) { + int base_w = (base_q_ + q) * params.stride[0]; + int base_h = (base_p_ + p) * params.stride[1]; + + int offset = base_h * params.activation_tile_w + base_w; + activation_offset[p][q][col] = offset; + } + } + } + } + + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + DepthwiseDirect2dConvSimtTileIterator &add_tile_offset(TensorCoord const &coord) { + // Set warp row and col start + lane_offset_ = MatrixCoord({lane_offset_.row() + coord.row() * Shape::kRow, lane_offset_.column()}); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + void advance(int32_t pointer_offset) { + ref_.reset(ref_.data() + pointer_offset / sizeof(Element) / Policy::LaneMmaShape::kN); + iterator_s_ = 0; + iterator_r_ = 0; + iterator_offset_ = 0; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + DepthwiseDirect2dConvSimtTileIterator &operator++() { + ++iterator_s_; + if (iterator_s_ < FilterShape::kColumn) { + iterator_offset_ += inc_next_s_; + + return *this; + } + + iterator_s_ = 0; + + ++iterator_r_; + if (iterator_r_ < FilterShape::kRow) { + iterator_offset_ += inc_next_r_; + return *this; + } + + iterator_r_ = 0; + iterator_offset_ = 0; + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + DepthwiseDirect2dConvSimtTileIterator & operator--() { + // Do nothing + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. (vector loads) + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + Array *dst_ptr = + reinterpret_cast *>(&frag); + + + MCTLASS_PRAGMA_UNROLL + for (int p = 0; p < ThreadOutputShape::kH; ++p) { + MCTLASS_PRAGMA_UNROLL + for (int q = 0; q < ThreadOutputShape::kW; ++q) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Iterations::kColumn; ++n) { + void const *ptr = ref_.data() + + ref_.offset({activation_offset[p][q][n] + (iterator_offset_), + n * Policy::WarpShape::kColumn}) + + pointer_offset / Policy::LaneMmaShape::kN; + arch::shared_load(dst_ptr[n + q + p * ThreadOutputShape::kW], ptr); + } + } + } + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) const { + // Do nothing at present. + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store(Fragment const &frag, Index pointer_offset) const { + store_with_pointer_offset(frag, 0); + } + + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation here + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// +/// Specialization for A operands of row-major layouts +/// +/// Concept: MutableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Size of filter (concept: gemm::GemmShape) + typename FilterShape_, + /// Size of the matrix to load (concept: TensorNHWC) + typename ThreadOutputShape_, + /// Size of the matrix to load (concept: TensorNHWC) + typename ThreadBlockOutputShape_, + /// Data type of A elements + typename Element_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Stride ( MatrixShape ) + typename StrideShape_, + /// Dilation ( MatrixShape ) + typename DilationShape_, + /// Activation Shape loaded by threadblock + typename ActivationShape_, + /// Number of partitions along K dimension - used in sliced-K + int PartitionsK, + /// Group Size along kPartition - used in sliced-K + int PartitionGroupSize> +class DepthwiseDirect2dConvSimtTileIterator { + public: + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Shape of filter (concept: gemm::GemmShape) + using FilterShape = FilterShape_; + + /// Shape of tile to load (concept: TensorNHWC) + using ThreadOutputShape = ThreadOutputShape_; + + /// Shape of tile to load (concept: TensorNHWC) + using ThreadBlockOutputShape = ThreadBlockOutputShape_; + + /// Stride ( MatrixShape ) + using StrideShape = StrideShape_; + + /// Dilation ( MatrixShape ) + using DilationShape = DilationShape_; + + /// Activation Shape loaded by threadblock + using ActivationShape = ActivationShape_; + + /// Operand tag + static mctlass::gemm::Operand const kOperand = mctlass::gemm::Operand::kA; + + /// Element type + using Element = Element_; + + /// Layout of policy + using Layout = layout::RowMajor; + + /// Decomposition of elements among threads + using Policy = Policy_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + // + // Derived quantities + // + + static_assert(!(Shape::kRow % Policy::WarpShape::kRow), + "The warp-level GEMM M size must be divisible by the number of threads arranged " + "along the M dimension."); + + static_assert(Shape::kRow > 0, "Shape::kRow must be greater than zero."); + static_assert(Shape::kColumn > 0, "Shape::kColumn must be greater than zero."); + static_assert(Policy::WarpShape::kRow > 0, "Policy::WarpShape::kRow must be greater than zero."); + static_assert(Shape::kRow / Policy::WarpShape::kRow > 0, + "Shape::kRow / Policy::WarpShape::kRow must be greater than zero."); + + // Activations loaded by threadblock + static int const ThreadActivationShapeH = (ThreadOutputShape::kH - 1) * StrideShape::kRow + + (FilterShape::kRow - 1) * DilationShape::kRow + 1; + + static int const ThreadActivationShapeW = (ThreadOutputShape::kW - 1) * StrideShape::kColumn + + (FilterShape::kColumn - 1) * DilationShape::kColumn + 1; + + using ThreadActivationShape = mctlass::conv:: + TensorNHWCShape<1, ThreadActivationShapeH, ThreadActivationShapeW, ThreadOutputShape::kC>; + + // Thread-level shape of a fragment + using ThreadShape = + MatrixShape; + + static_assert(!(ThreadShape::kColumn % Policy::LaneMmaShape::kN), + "Thread-level GEMM must be divisible by Policy::LaneMmaShape."); + + /// Number of individual loads + using Iterations = + MatrixShape; + + using ThreadTileCount = MatrixShape; + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + + protected: + /// Internal reference + mctlass::TensorRef, layout::RowMajor> ref_; + + Array + activation[ThreadActivationShape::kH][ThreadActivationShape::kW][Iterations::kColumn]; + int iterator_r_; + int iterator_s_; + + + MatrixCoord lane_offset_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + DepthwiseDirect2dConvSimtTileIterator() {} + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + DepthwiseDirect2dConvSimtTileIterator(TensorRef ref, int lane_id) { + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + // Set channel offset + lane_offset_ = lane_layout.inverse(lane_id) * MatrixCoord(0, Policy::LaneMmaShape::kN); + + ref.add_coord_offset(lane_offset_); + + ref_.reset(reinterpret_cast *>(ref.data()), + ref.stride(0) / Policy::LaneMmaShape::kN); + + iterator_r_ = 0; + iterator_s_ = 0; + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + DepthwiseDirect2dConvSimtTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + template + MCTLASS_HOST_DEVICE void setup_initial_status( + Params const ¶ms) { + + // Get base HW offset of current threads + int threadgroup = threadIdx.x / (ThreadBlockOutputShape::kC / ThreadOutputShape::kC); + int base_h = + (threadgroup / (ThreadTileCount::kColumn)) * ThreadOutputShape::kH * StrideShape::kRow; + int base_w = + (threadgroup % (ThreadTileCount::kColumn)) * ThreadOutputShape::kW * StrideShape::kColumn; + + MCTLASS_PRAGMA_UNROLL + for (int h = 0; h < ThreadActivationShape::kH; ++h) { + MCTLASS_PRAGMA_UNROLL + for (int w = 0; w < ThreadActivationShape::kW; ++w) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < Iterations::kColumn; ++col) { + int offset = (base_h + h) * ActivationShape::kW + (base_w + w); + + void const *ptr = ref_.data() + ref_.offset({offset, col * Policy::WarpShape::kColumn}); + arch::shared_load(activation[h][w][col], ptr); + } + } + } + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + DepthwiseDirect2dConvSimtTileIterator &add_tile_offset(TensorCoord const &coord) { + // Set warp row and col start + lane_offset_ = + MatrixCoord({lane_offset_.row() + coord.row() * Shape::kRow, lane_offset_.column()}); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + void advance(int32_t pointer_offset) { + ref_.reset(ref_.data() + pointer_offset / sizeof(Element) / Policy::LaneMmaShape::kN); + iterator_s_ = 0; + iterator_r_ = 0; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + DepthwiseDirect2dConvSimtTileIterator &operator++() { + ++iterator_s_; + if (iterator_s_ < FilterShape::kColumn) { + return *this; + } + + iterator_s_ = 0; + + ++iterator_r_; + if (iterator_r_ < FilterShape::kRow) { + return *this; + } + + iterator_r_ = 0; + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + DepthwiseDirect2dConvSimtTileIterator &operator--() { + // Do nothing + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. (vector loads) + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + Array *dst_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int p = 0; p < ThreadOutputShape::kH; ++p) { + MCTLASS_PRAGMA_UNROLL + for (int q = 0; q < ThreadOutputShape::kW; ++q) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Iterations::kColumn; ++n) { + const int h = p * StrideShape::kRow + iterator_r_ * DilationShape::kRow; + const int w = q * StrideShape::kColumn + iterator_s_ * DilationShape::kColumn; + + dst_ptr[n + q + p * ThreadOutputShape::kW] = activation[h][w][n]; + } + } + } + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { load_with_pointer_offset(frag, 0); } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) const { + // Do nothing at present. + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store(Fragment const &frag, Index pointer_offset) const { + store_with_pointer_offset(frag, 0); + } + + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation here + } +}; + +} // namespace warp +} // namespace conv +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/conv/warp/scale_bias_relu_transform.h b/csrc/mctlass/include/mctlass/conv/warp/scale_bias_relu_transform.h new file mode 100644 index 0000000..d6a436a --- /dev/null +++ b/csrc/mctlass/include/mctlass/conv/warp/scale_bias_relu_transform.h @@ -0,0 +1,225 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level per channel scale+bias+relu before + matrix multiply-accumulate operations targeting Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/platform/platform.h" + +#include "mctlass/numeric_conversion.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/arch/mma_sm75.h" +#include "mctlass/arch/mma_sm80.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/warp/mma_tensor_op_policy.h" + +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace conv { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct FpropScaleBiasReluTransform { + + using T = typename FragmentActivations::Element; + + static int const NumActivations = FragmentActivations::kElements; + static int const NumScaleBias = FragmentScaleBias::kElements; + static int const MmaElements = 2; + // One element has one scale and one bias + static int const MmaScaleBiasPair = 2; + // 16816 has 2 columns + static int const MmaCols = 2; + + using MmaOperand = Array; + using ScaleBiasOperand = Array; + + MCTLASS_DEVICE + void transform(MmaOperand &activations, ScaleBiasOperand const &scale_bias) { + +//#if defined(__MACA_ARCH__) +#if 0 + uint32_t *ptr_activations = reinterpret_cast(&activations); + uint32_t const *ptr_scale_bias = reinterpret_cast(&scale_bias); + + // Apply per channel scale+bias+relu if the data is not a special NaN + // (0x7eff). If it is a special NaN (0x7eff), hard code the output to 0. + + // We assumes the pair of FP16 are either both inbound or both out-of-bound. + // It requires C to be an even number. + asm volatile( + "{\n\t" + " .reg .pred %%p;\n\t" + " .reg .b32 t1;\n\t" + " setp.eq.u32 %%p, %2, %4;\n\t" + " fma.rn.f16x2.relu t1, %1, %2, %3;\n" + " selp.u32 %0, 0, t1, %%p;\n\t" + "}\n" + : "=r"(ptr_activations[0]) + : "r"(ptr_scale_bias[0]), "r"(ptr_activations[0]), + "r"(ptr_scale_bias[1]), "n"(mctlass::arch::OOB_NAN_F16x2)); +#else + printf("scale_bias_relu_transform.h L108 this function cannot working correctly now.\n"); + assert(0); +#endif + } + + MCTLASS_DEVICE + void operator()(FragmentActivations &activations, + FragmentScaleBias const &scale_bias) { + MmaOperand *ptr_activations = reinterpret_cast(&activations); + ScaleBiasOperand const *ptr_scale_bias = + reinterpret_cast(&scale_bias); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < (NumActivations / MmaElements); ++i) { + transform(ptr_activations[i], ptr_scale_bias[(i / MmaScaleBiasPair) % MmaCols]); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct WgradScaleBiasReluTransform { + + using T = typename FragmentActivations::Element; + + static int const NumActivations = FragmentActivations::kElements; + static int const NumScaleBias = FragmentScaleBias::kElements; + static int const MmaElements = 2; + // One element has one scale and one bias + static int const MmaScaleBiasPair = 2; + // 16816 has 2 rows + static int const MmaRows = 2; + + using MmaOperand = Array; + using ScaleBiasOperand = Array<__half2, MmaScaleBiasPair>; + + MCTLASS_DEVICE + void transform(MmaOperand &activations, ScaleBiasOperand const &scale_bias) { + +//#if defined(__MACA_ARCH__) +#if 0 + + __half2 *ptr_activations = reinterpret_cast<__half2 *>(&activations); + uint32_t const *ptr_scale_bias = reinterpret_cast(&scale_bias); + +#if 1 + // CUDA + PTX version + + bool h1_oob = (reinterpret_cast(ptr_activations[0].x) == mctlass::arch::OOB_NAN_F16); + bool h2_oob = (reinterpret_cast(ptr_activations[0].y) == mctlass::arch::OOB_NAN_F16); + + // Apply per channel scale+bias+relu if the data is not a special NaN + // (0x7eff). If it is a special NaN (0x7eff), hard code the output to 0. + + // We cannot gurantee that the pair of F16 are both in bound or both + // out-of-bound because C x R x S can be an odd number. + asm volatile( + "{\n\t" + " fma.rn.f16x2.relu %0, %1, %2, %3;\n" + "}" + : "=r"(reinterpret_cast(ptr_activations[0])) + : "r"(ptr_scale_bias[0]), "r"(reinterpret_cast(ptr_activations[0])), + "r"(ptr_scale_bias[1])); + + reinterpret_cast(ptr_activations[0]) = h1_oob ? + (reinterpret_cast(ptr_activations[0]) & 0xffff0000) : + reinterpret_cast(ptr_activations[0]); + + reinterpret_cast(ptr_activations[0]) = h2_oob ? + (reinterpret_cast(ptr_activations[0]) & 0xffff) : + reinterpret_cast(ptr_activations[0]); +#else + // pure PTX version + + // Apply per channel scale+bias+relu if the data is not a special NaN + // (0x7eff). If it is a special NaN (0x7eff), hard code the output to 0. + asm volatile( + "{\n" + " .reg .b16 t1, t2;\n" + " .reg .b32 t3, t4, t5, t6;\n" + " .reg .pred p1, p2;\n" + " mov.b32 {t1, t2}, %2;\n" + " setp.eq.s16 p1, t1, %4;\n" + " setp.eq.s16 p2, t2, %4;\n" + " fma.rn.f16x2.relu t3, %1, %2, %3;\n" + " and.b32 t4, t3, %5;\n" + " selp.b32 t5, t4, t3, p1;\n" + " and.b32 t6, t5, %6;\n" + " selp.b32 %0, t6, t5, p2;\n" + "}\n" + : "=r"(reinterpret_cast(ptr_activations[0])) + : "r"(ptr_scale_bias[0]), "r"(reinterpret_cast(ptr_activations[0])), + "r"(ptr_scale_bias[1]), "n"(mctlass::arch::OOB_NAN_F16), "n"(0xffff0000), "n"(0x0000ffff)); +#endif +#else + printf("scale_bias_relu_transform.h L205 this function cannot working correctly now.\n"); + assert(0); +#endif + } + + MCTLASS_DEVICE + void operator()(FragmentActivations &activations, + FragmentScaleBias const &scale_bias) { + MmaOperand *ptr_activations = reinterpret_cast(&activations); + ScaleBiasOperand const *ptr_scale_bias = + reinterpret_cast(&scale_bias); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < (NumActivations / MmaElements); ++i) { + transform(ptr_activations[i], ptr_scale_bias[(i / MmaRows)]); + } + } +}; +} // namespace warp +} // namespace conv +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/coord.h b/csrc/mctlass/include/mctlass/coord.h new file mode 100644 index 0000000..810cae8 --- /dev/null +++ b/csrc/mctlass/include/mctlass/coord.h @@ -0,0 +1,480 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief A Coord is a coordinate of arbitrary rank into a tensor or matrix +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mctlass/mctlass.h" + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Statically-sized array specifying Coords within a tensor +template < + int Rank_, ///< Logical rank of coordinate + typename Index_ = int, ///< Index type used for each dimension + typename LongIndex_ = int64_t ///< Long index type used for linear offsets +> +struct Coord { + +public: + + // + // Type and constant definitions + // + + /// Number of elements in Coord + static int const kRank = Rank_; + + /// Index type used to store elements + using Index = Index_; + + /// Type used to represent linear offsets + using LongIndex = LongIndex_; + +private: + + // + // Data members + // + + /// Indices + Index idx[kRank]; + +public: + + // + // Methods + // + + /// Default ctor initializes uniformly + MCTLASS_HOST_DEVICE + explicit Coord(Index value = Index(0)) { + for (int i = 0; i < kRank; ++i) { + idx[i] = value; + } + } + + /// Constructs from an array of integers + MCTLASS_HOST_DEVICE + Coord(Index const (&_idx)[kRank]) { + for (int i = 0; i < kRank; ++i) { + idx[i] = _idx[i]; + } + } + + /// Constructs from some other Coord + template + MCTLASS_HOST_DEVICE + Coord(Coord other) { + for (int i = 0; i < kRank; ++i) { + idx[i] = other[i]; + } + } + + /// Returns a slice of the Coord which may be larger or smaller in rank + /// than this. + template + MCTLASS_HOST_DEVICE + Coord slice(int start = 0, Index identity = 0) const { + Coord result; + for (int i = 0; i < Slice; ++i) { + if (i + start < kRank) { + result[i] = idx[i + start]; + } + else { + result[i] = identity; + } + } + return result; + } + + /// Returns the index of the dimension with least value + MCTLASS_HOST_DEVICE + int min_dim_index() const { + int i = 0; + for (int j = 1; j < kRank; ++j) { + if (idx[j] < idx[i]) { + i = j; + } + } + return i; + } + + /// Returns the index of the dimension with greatest value + MCTLASS_HOST_DEVICE + int max_dim_index() const { + int i = 0; + for (int j = 1; j < kRank; ++j) { + if (idx[j] > idx[i]) { + i = j; + } + } + return i; + } + + /// Returns true if Coord is non-zero. + MCTLASS_HOST_DEVICE + explicit operator bool() const { + for (int i = 0; i < kRank; ++i) { + if (idx[i]) { + return true; + } + } + return false; + } + + /// Returns true if Coord is uniformly zero. + MCTLASS_HOST_DEVICE + bool operator!() const { + for (int i = 0; i < kRank; ++i) { + if (idx[i]) { + return false; + } + } + return true; + } + + /// Element-wise addition + MCTLASS_HOST_DEVICE + Coord operator+(Coord const& b) const { + Coord c; + for (int i = 0; i < kRank; ++i) { + c.idx[i] = idx[i] + b.idx[i]; + } + return c; + } + + /// Element-wise subtraction + MCTLASS_HOST_DEVICE + Coord operator-(Coord const& b) const { + Coord c; + for (int i = 0; i < kRank; ++i) { + c.idx[i] = idx[i] - b.idx[i]; + } + return c; + } + + /// Element-wise multiplication + MCTLASS_HOST_DEVICE + Coord operator*(Coord const& b) const { + Coord c; + for (int i = 0; i < kRank; ++i) { + c.idx[i] = idx[i] * b.idx[i]; + } + return c; + } + + /// Element-wise division + MCTLASS_HOST_DEVICE + Coord operator/(Coord const& b) const { + Coord c; + for (int i = 0; i < kRank; ++i) { + c.idx[i] = idx[i] / b.idx[i]; + } + return c; + } + + /// In-place addition + MCTLASS_HOST_DEVICE + Coord& operator+=(Coord const& b) { + for (int i = 0; i < kRank; ++i) { + idx[i] += b.idx[i]; + } + return *this; + } + + /// In-place subtraction + MCTLASS_HOST_DEVICE + Coord& operator-=(Coord const& b) { + for (int i = 0; i < kRank; ++i) { + idx[i] -= b.idx[i]; + } + return *this; + } + + /// In-place multiplication + MCTLASS_HOST_DEVICE + Coord& operator*=(Coord const& b) { + for (int i = 0; i < kRank; ++i) { + idx[i] *= b.idx[i]; + } + return *this; + } + + /// In-place division + MCTLASS_HOST_DEVICE + Coord& operator/=(Coord const& b) { + for (int i = 0; i < kRank; ++i) { + idx[i] /= b.idx[i]; + } + return *this; + } + + /// Member access operator + MCTLASS_HOST_DEVICE Index& operator[](int dim) { return idx[dim]; } + + /// Member access operator + MCTLASS_HOST_DEVICE Index const& operator[](int dim) const { return idx[dim]; } + + /// Computes the dot product with anotherCoord object + MCTLASS_HOST_DEVICE + LongIndex dot(Coord const& b, LongIndex sum = LongIndex(0)) const { + for (int i = 0; i < kRank; ++i) { + sum += idx[i] * b.idx[i]; + } + return sum; + } + + /// Gets the index of a given Coord element + template + MCTLASS_HOST_DEVICE Index& at() { + return idx[Dim]; + } + + /// Access via index; may limit unrolling potential + MCTLASS_HOST_DEVICE + Index& at(int dim) { return idx[dim]; } + + /// Gets the index of a given Coord element + template + MCTLASS_HOST_DEVICE Index const& at() const { + return idx[Dim]; + } + + /// Access via index; may limit unrolling potential + MCTLASS_HOST_DEVICE + Index const& at(int dim) const { return idx[dim]; } + + /// Determines if two Coord<> objects are equal + MCTLASS_HOST_DEVICE + bool operator==(Coord const& b) const { + bool equal = true; + for (int i = 0; equal && i < kRank; ++i) { + equal = (idx[i] == b.idx[i]); + } + return equal; + } + + /// Not equal + MCTLASS_HOST_DEVICE + bool operator!=(Coord const& b) const { return !(*this == b); } + + /// Clamps a coordinate to a range specified by maximum and minimum values + MCTLASS_HOST_DEVICE + Coord& clamp(Coord const& max, Coord const& min = Coord()) { + for (int i = 0; i < kRank; ++i) { + idx[i] = __NV_STD_MAX(__NV_STD_MIN(idx[i], max.idx[i]), min.idx[i]); + } + return *this; + } + + /// Returns the sum of all elements + MCTLASS_HOST_DEVICE + Index sum() const { + Index sum_(idx[0]); + for (int i = 1; i < kRank; ++i) { + sum_ += idx[i]; + } + return sum_; + } + + /// Returns the product of all elements + MCTLASS_HOST_DEVICE + LongIndex product() const { + LongIndex product_(idx[0]); + for (int i = 1; i < kRank; ++i) { + product_ *= idx[i]; + } + return product_; + } + + /// Less than operator + MCTLASS_HOST_DEVICE + bool operator<(Coord const &b) const { + for (int i = 0; i < kRank; ++i) { + if (!(idx[i] < b[i])) { + return false; + } + } + return true; + } + + /// Less than or equals operator + MCTLASS_HOST_DEVICE + bool operator<=(Coord const &b) const { + for (int i = 0; i < kRank; ++i) { + if (!(idx[i] <= b[i])) { + return false; + } + } + return true; + } + + /// Greater than operator + MCTLASS_HOST_DEVICE + bool operator>(Coord const &b) const { + return !(*this <= b); + } + + /// Greater than or equals operator + MCTLASS_HOST_DEVICE + bool operator>=(Coord const &b) const { + return !(*this < b); + } +}; + +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + + +/// Scalar multiplication +template +MCTLASS_HOST_DEVICE +Coord operator*(Index s, Coord coord) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Rank; ++i) { + coord[i] *= s; + } + return coord; +} + +/// Scalar multiplication +template +MCTLASS_HOST_DEVICE +Coord operator*(Coord coord, Index s) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Rank; ++i) { + coord[i] *= s; + } + return coord; +} + +/// Scalar division +template +MCTLASS_HOST_DEVICE +Coord operator/(Index s, Coord coord) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Rank; ++i) { + coord[i] = s / coord[i]; + } + return coord; +} + +/// Scalar division +template +MCTLASS_HOST_DEVICE +Coord operator/(Coord coord, Index s) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Rank; ++i) { + coord[i] /= s; + } + return coord; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Integer-valued make_Coord +// +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Helper to make a 2-element coordinate +template +MCTLASS_HOST_DEVICE +Coord<1, T> make_Coord(T _0) { + T values[1] = {_0}; + return Coord<1, T>(values); +} + +/// Helper to make a 2-element coordinate +template +MCTLASS_HOST_DEVICE +Coord<2, T> make_Coord(T _0, T _1) { + T values[2] = {_0, _1}; + return Coord<2, T>(values); +} + +/// Helper to make a 3-element coordinate +template +MCTLASS_HOST_DEVICE +Coord<3, T> make_Coord(T _0, T _1, T _2) { + T values[3] = {_0, _1, _2}; + return Coord<3, T>(values); +} + +/// Helper to make a 4-element coordinate +template +MCTLASS_HOST_DEVICE +Coord<4, T> make_Coord(T _0, T _1, T _2, T _3) { + T values[4] = {_0, _1, _2, _3}; + return Coord<4, T>(values); +} + +/// Helper to make a 5-element coordinate +template +MCTLASS_HOST_DEVICE +Coord<5, T> make_Coord(T _0, T _1, T _2, T _3, T _4) { + T values[5] = {_0, _1, _2, _3, _4}; + return Coord<5, T>(values); +} + +/// Helper to make a 1-element coordinate +template +MCTLASS_HOST_DEVICE +Coordmake_Coord_with_padding(T _0) { + Coord coord; + + MCTLASS_PRAGMA_UNROLL + for (int i = N - 1; i > 0; --i) { + coord[i] = 0; + } + + coord[0] = _0; + + return coord; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/core_io.h b/csrc/mctlass/include/mctlass/core_io.h new file mode 100644 index 0000000..2bdc0f1 --- /dev/null +++ b/csrc/mctlass/include/mctlass/core_io.h @@ -0,0 +1,287 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Helpers for printing mctlass/core objects +*/ + +#pragma once + +#include +#include + +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix.h" +#include "mctlass/quaternion.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/tensor_view.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/conv3d_problem_size.h" + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Output operator for CUDA built-in dim3 type +inline std::ostream &operator<<(std::ostream &out, dim3 d) { + return out << d.x << ", " << d.y << ", " << d.z; +} + +/// Output operator for CUDA built-in error type +inline std::ostream &operator<<(std::ostream &out, mcError_t error) { + return out << mcGetErrorString(error); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// stream operators for mctlass namespace // +/////////////////////////////////////////////////////////////////////////////////////////////////// + +template +inline +std::ostream& operator<<(std::ostream& out, Array const& v) { + for (int i = 0; i < Rank; ++i) { + out << (i ? ", " : "") << v[i]; + } + return out; +} + +template +inline +std::ostream& operator<<(std::ostream& out, Coord const& coord) { + for (int i = 0; i < Rank; ++i) { + out << (i ? ", " : "") << coord[i]; + } + return out; +} + +inline +std::istream & operator>>(std::istream &stream, half_t &x) { + float tmp; + stream >> tmp; + x = static_cast(tmp); + return stream; +} + +inline +std::ostream & operator<<(std::ostream &out, half_t const &x) { + return out << float(x); +} + +inline +std::ostream & operator<<(std::ostream &out, bfloat16_t const &x) { + return out << float(x); +} + +inline +std::ostream & operator<<(std::ostream &out, tfloat32_t const &x) { + return out << float(x); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Helper to enable formatted printing of MCTLASS scalar types to an ostream +template +struct ScalarIO { + + /// Value to print + T value; + + /// Default ctor + ScalarIO() { } + + /// Constructs from a value + ScalarIO(T value): value(value) {} +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Default printing to ostream +template +inline std::ostream &operator<<(std::ostream &out, ScalarIO const &scalar) { + return out << scalar.value; +} + +/// Printing to ostream of int8_t as integer rather than character +template <> +inline std::ostream &operator<<(std::ostream &out, ScalarIO const &scalar) { + return out << int(scalar.value); +} + +/// Printing to ostream of uint8_t as integer rather than character +template <> +inline std::ostream &operator<<(std::ostream &out, ScalarIO const &scalar) { + return out << unsigned(scalar.value); +} + + +/// Default printing to ostream for MatrixShape +template +inline +std::ostream & operator<<(std::ostream &out, MatrixShape const &matrix_shape) { + out << "mctlass::MatrixShape::(kRow, kColumn) {" + << mctlass::MatrixShape::kRow <<"," + << mctlass::MatrixShape::kColumn <<"}"; + return out; +} + + +/// Prints matrix to ostream +template +std::ostream & operator<<(std::ostream &out, Matrix const &rhs) { + + for (int i = 0; i < Rows; ++i) { + for (int j = 0; j < Columns; ++j) { + ScalarIO element(rhs.at(i, j)); + out << (j ? ", " : "") << element; + } + out << "\\n"; + } + + return out; +} + +template +std::ostream &operator<<(std::ostream &out, Quaternion const &rhs) { + + out << ScalarIO(rhs.w()) << " "; + if (rhs.x() >= 0) { + out << "+"; + } + + out << ScalarIO(rhs.x()) << "*i "; + if (rhs.y() >= 0) { + out << "+"; + } + + out << ScalarIO(rhs.y()) << "*j "; + if (rhs.z() >= 0) { + out << "+"; + } + + out << ScalarIO(rhs.z()) << "*k"; + + return out; +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// stream operators for mctlass::gemm namespace // +/////////////////////////////////////////////////////////////////////////////////////////////////// +namespace gemm { + +/// Default printing to ostream for GemmShape +template +inline +std::ostream & operator<<(std::ostream &out, GemmShape const &gemm_shape) { + out << "mctlass::gemm::GemmShape::(kM, kN, kK) {" + << mctlass::gemm::GemmShape::kM <<"," + << mctlass::gemm::GemmShape::kN <<"," + << mctlass::gemm::GemmShape::kK << "}"; + return out; +} + +/// Default printing to ostream for GemmCoord +inline +std::ostream & operator<<(std::ostream &out, GemmCoord const &gemm_coord) { + out << "mctlass::gemm::GemmCoord {" + << gemm_coord.m() <<"," + << gemm_coord.n() <<"," + << gemm_coord.k() << "}"; + return out; +} + +} //namespace gemm +/////////////////////////////////////////////////////////////////////////////////////////////////// + + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// stream operators for mctlass namespace // +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Default printing to ostream for PitchLinearShape +template < int Contiguous, int Strided> +inline +std::ostream & operator<<(std::ostream &out, PitchLinearShape const &pitch_linear_shape) { + out << "mctlass::PitchLinearShape:(kContiguous, kStrided) {" + << mctlass::layout::PitchLinearShape::kContiguous <<"," + << mctlass::layout::PitchLinearShape::kStrided <<"}"; + return out; +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// + + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// stream operators for mctlass::conv namespace // +/////////////////////////////////////////////////////////////////////////////////////////////////// +namespace conv { +/// Default printing to ostream for Conv2dProblemSize +inline +std::ostream& operator<<(std::ostream& out, Conv2dProblemSize const& problem) { + out << "NHWC: (" << problem.N << ", " << problem.H << ", " << problem.W << ", " << problem.C << ")" << std::endl + << "KRSC: (" << problem.K << ", " << problem.R << ", " << problem.S << ", " << problem.C / problem.groups << ")" << std::endl + << "NPQK: (" << problem.N << ", " << problem.P << ", " << problem.Q << ", " << problem.K << ")" << std::endl + << "groups: (" << problem.groups << ")" << std::endl + << "Pad_h, Pad_w: (" << problem.pad_h << ", " << problem.pad_w << ")" << std::endl + << "Stride_h, Stride_w: (" << problem.stride_h << ", " << problem.stride_w << ")" << std::endl + << "Dilation_h, Dilation_w: (" << problem.dilation_h << ", " << problem.dilation_w << ")" << std::endl + << "split_k_slices: (" << problem.split_k_slices << ")" << std::endl + << "mode: (" << ((problem.mode==conv::Mode::kConvolution) ? "conv" : "xcross") << ")"; + + return out; +} + + +/// Default printing to ostream for Conv3dProblemSize +inline +std::ostream& operator<<(std::ostream& out, Conv3dProblemSize const& problem) { + out << "NDHWC: (" << problem.N << ", " << problem.D << ", " << problem.H << ", " << problem.W << ", " << problem.C << ")" << std::endl + << "KTRSC: (" << problem.K << ", " << problem.T << ", " << problem.R << ", " << problem.S << ", " << problem.C << ")" << std::endl + << "NZPQK: (" << problem.N << ", " << problem.Z << ", " << problem.P << ", " << problem.Q << ", " << problem.K << ")" << std::endl + << "pad_d, pad_h, pad_w: (" << problem.pad_d << ", " << problem.pad_h << ", " << problem.pad_w << ")" << std::endl + << "stride_d, stride_h, stride_w: (" << problem.stride_d << ", " << problem.stride_h << ", " << problem.stride_w << ")" << std::endl + << "dilation_d, dilation_h, dilation_w: (" << problem.dilation_d << ", " << problem.dilation_h << ", " << problem.dilation_w << ")" << std::endl + << "split_k_slices: (" << problem.split_k_slices << ") " << std::endl + << "mode: (" << ((problem.mode==conv::Mode::kConvolution) ? "conv" : "xcross") << ")"; + + return out; +} + +} // namespace conv +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass +/////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/detail/dependent_false.hpp b/csrc/mctlass/include/mctlass/detail/dependent_false.hpp new file mode 100644 index 0000000..dc546bd --- /dev/null +++ b/csrc/mctlass/include/mctlass/detail/dependent_false.hpp @@ -0,0 +1,86 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::detail { + +/// @brief A bool constant that depends on one or more template parameters. +/// +/// For more detailed documentation and use cases, +/// please see `dependent_false` below. +template +inline constexpr bool dependent_bool_value = Value; + +/// @brief An always-false value that depends on one or more template parameters. +/// +/// This exists because `static_assert(false);` always fails, +/// even if it occurs in the `else` branch of an `if constexpr`. +/// The following example shows how to use `dependent_false` in that case. +/// +/// @code +/// template +/// void foo (T t) +/// { +/// if constexpr (std::is_integral_v) { +/// do_integer_stuff(t); +/// } +/// else if constexpr (std::is_floating_point_v) { +/// do_floating_point_stuff(t); +/// } +/// else { +/// static_assert(dependent_false, "T must be " +/// "an integral or floating-point type."); +/// } +/// } +/// @endcode +/// +/// This implements the C++ Standard Library proposal P1830R1. +/// +/// https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2019/p1830r1.pdf +/// +/// That proposal is under review as of 2022/12/05. +/// The following link shows P1830's current review status. +/// +/// https://github.com/cplusplus/papers/issues/572 +/// +/// P2593R0 proposes an alternate solution to this problem, +/// that would change the C++ language itself. +/// +/// https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2022/p2593r0.html +/// +/// For headers in this library, however, we only consider library solutions +/// as work-arounds for future C++ features. +template +inline constexpr bool dependent_false = dependent_bool_value; + +} // end namespace mctlass::detail diff --git a/csrc/mctlass/include/mctlass/device_kernel.h b/csrc/mctlass/include/mctlass/device_kernel.h new file mode 100644 index 0000000..acd95a8 --- /dev/null +++ b/csrc/mctlass/include/mctlass/device_kernel.h @@ -0,0 +1,120 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for generic MCTLASS kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +// __grid_constant__ was introduced in CUDA 11.7. +#if 0 +# define MCTLASS_GRID_CONSTANT_SUPPORTED +#endif + +// __grid_constant__ can be enabled only on SM70+ +#if defined(MCTLASS_GRID_CONSTANT_SUPPORTED) && defined(__MACA_ARCH__) +# define MCTLASS_GRID_CONSTANT_ENABLED +#endif + +#if ! defined(MCTLASS_GRID_CONSTANT) +# if defined(MCTLASS_GRID_CONSTANT_ENABLED) +# define MCTLASS_GRID_CONSTANT __grid_constant__ +# else +# define MCTLASS_GRID_CONSTANT +# endif +#endif + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////// + +/// Generic MCTLASS kernel template. +template +__global__ +void Kernel(typename Operator::Params params) { + // Dynamic shared memory base pointer + extern __shared__ int SharedStorageBase[]; + + // Declare pointer to dynamic shared memory. + typename Operator::SharedStorage *shared_storage = + reinterpret_cast(SharedStorageBase); + + Operator op; + + op(params, *shared_storage); +} + +/// Generic MCTLASS kernel template. +template +__global__ +void Kernel2(typename Operator::Params params) { + // Dynamic shared memory base pointer + extern __shared__ int SharedStorageBase[]; + + // Declare pointer to dynamic shared memory. + typename Operator::SharedStorage *shared_storage = + reinterpret_cast(SharedStorageBase); + + Operator::invoke(params, *shared_storage); + +} + + +//////////////////////////////////////////////////////////////////////////////// +// +// 3.0 specific launch +// +//////////////////////////////////////////////////////////////////////////////// + +/// Generic MCTLASS kernel template. +template +__global__ +#ifdef __MACACC__ +// Enclosing this in __MACACC__ suppresses MSVC warnings. +__launch_bounds__(Operator::MaxThreadsPerBlock, Operator::MinBlocksPerMultiprocessor) +#endif // __MACACC__ +// void device_kernel(MCTLASS_GRID_CONSTANT typename Operator::Params const params) +void device_kernel(MCTLASS_GRID_CONSTANT typename Operator::Params const params) +{ + // Dynamic shared memory base pointer + extern __shared__ char smem[]; + + Operator op; + op(params, smem); +} + +//////////////////////////////////////////////////////////////////////////////// +} /// namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/epilogue/collective/builders/sm90_builder.inl b/csrc/mctlass/include/mctlass/epilogue/collective/builders/sm90_builder.inl new file mode 100644 index 0000000..badfaea --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/collective/builders/sm90_builder.inl @@ -0,0 +1,545 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "cute/atom/mma_traits_sm90.hpp" +#include "cute/atom/mma_traits_sm90_gmma.hpp" +#include "cute/atom/copy_traits_sm90.hpp" + +#include "mctlass/detail/dependent_false.hpp" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/collective/collective_builder.hpp" +#include "mctlass/epilogue/dispatch_policy.hpp" +#include "mctlass/epilogue/collective/collective_epilogue.hpp" +#include "mctlass/epilogue/thread/linear_combination.h" +#include "mctlass/epilogue/thread/linear_combination_generic.h" +#include "mctlass/epilogue/thread/linear_combination_bias_elementwise.h" + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +/////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::epilogue::collective { + +/////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +// Returns the smem layout atom to be used for C or D matrix +template +constexpr auto +sm90_get_epilogue_smem_swizzle_layout_atom() { + using namespace cute; + + // ColMajor C/D (M-major) + if constexpr (size<0>(GmemStrideType{}) == 1) { + return mctlass::gemm::collective::detail::ss_smem_selector< + cute::GMMA::Major::MN, Element, decltype(get<0>(EpilogueTile_MN{})), decltype(get<1>(EpilogueTile_MN{})) + >(); + } + // RowMajor C/D (N-major) + else if constexpr (size<1>(GmemStrideType{}) == 1) { + return mctlass::gemm::collective::detail::ss_smem_selector< + cute::GMMA::Major::K , Element, decltype(get<0>(EpilogueTile_MN{})), decltype(get<1>(EpilogueTile_MN{})) + >(); + } + else { + static_assert(mctlass::detail::dependent_false, "Unsupported gmem layout."); + } +} + +// Attempts to compute a reasonable epilogue tile based on block tile shape or allows the user to provide one. +template +constexpr auto +sm90_compute_tile_shape_or_override() { + if constexpr (cute::is_same_v) { + + if constexpr (detail::sm90_is_cooperative_v) { + return Shape<_128,_16>{}; + } + else if constexpr (detail::sm90_is_warp_specialized_v) { + return Shape<_64,_32>{}; + } + else { + static_assert(mctlass::detail::dependent_false, "Unsupported schedule."); + } + } + else if constexpr (cute::is_tuple::value) { + EpilogueTileType epi_tile; + constexpr int M = size<0>(shape(epi_tile)); + constexpr int N = size<1>(shape(epi_tile)); + + static_assert(!is_layout::value, "EpilogueTile must be a cute::Tile or cute::Shape"); + static_assert(M == 64 && detail::sm90_is_warp_specialized_v || + M == 128 && detail::sm90_is_cooperative_v, "Unsupported tile shape"); + static_assert(N % 8 == 0, "Unsupported tile shape"); + + return epi_tile; + } + else { + static_assert(mctlass::detail::dependent_false, "Invalid type for EpilogueTileType."); + } +} + +// Selects the largest vectorized smem store atom available +template +constexpr auto +sm90_get_smem_store_op_for_accumulator() { + using namespace cute; + + if constexpr (sizeof(ElementD) == 2 && size<0>(GmemStrideTypeD{}) == 1) { + return SM90_U16x8_STSM_T{}; + } + else if constexpr (sizeof(ElementD) == 2 && size<1>(GmemStrideTypeD{}) == 1) { + return SM90_U32x4_STSM_N{}; + } + else { + // auto-vectorizing store + return DefaultCopy{}; + } +} + +// Selects the largest vectorized smem load atom available +template +constexpr auto +sm90_get_smem_load_op_for_source() { + using namespace cute; + + // Reuse the logic from smem store selector + using SmemStoreOp = decltype(sm90_get_smem_store_op_for_accumulator()); + + if constexpr (cute::is_same_v) { + return SM75_U16x8_LDSM_T{}; + } + else if constexpr (cute::is_same_v) { + return SM75_U32x4_LDSM_N{}; + } + else { + // auto-vectorizing load + return DefaultCopy{}; + } +} + +// Helper for building TMA warp-specialized collective epilogues, specialized by +// the thread-level epilogue operation performed and the dispatch policy to use. +template < + class TileShape_MNK, + class ClusterShape_MNK, + class EpilogueTileType, + class ElementAccumulator, + class ElementCompute, + class ElementC_, + class GmemLayoutTagC_, + int AlignmentC, + class ElementD, + class GmemLayoutTagD, + int AlignmentD, + class Schedule, + class ThreadOp, + class DispatchPolicy +> +struct TmaBuilderImpl { + + // Passing void C disables source load + using ElementC = cute::conditional_t,ElementD,ElementC_>; // prevents void ref breakages + using GmemLayoutTagC = cute::conditional_t,GmemLayoutTagD,GmemLayoutTagC_>; + + using GmemStrideTypeC = gemm::TagToStrideC_t; + using GmemStrideTypeD = gemm::TagToStrideC_t; + + using EpilogueTile_MN = decltype(detail::sm90_compute_tile_shape_or_override< + ElementD, EpilogueTileType, Schedule>()); + + using CollectiveOp = mctlass::epilogue::collective::CollectiveEpilogue< + DispatchPolicy, + TileShape_MNK, + EpilogueTile_MN, + ElementC_, // Need to pass void through to expose via GemmUniversal + GmemStrideTypeC, + ElementD, + GmemStrideTypeD, + ThreadOp, + SM90_TMA_LOAD, + decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom()), + decltype(detail::sm90_get_smem_load_op_for_source()), + SM90_TMA_STORE, + decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom()), + decltype(detail::sm90_get_smem_store_op_for_accumulator()) + >; +}; + +} // namespace detail + +/////////////////////////////////////////////////////////////////////////////// + +// No-smem builder +template < + class TileShape_MNK, + class ClusterShape_MNK, + class EpilogueTileType, + class ElementAccumulator, + class ElementCompute, + class ElementC_, + class GmemLayoutTagC_, + int AlignmentC, + class ElementD, + class GmemLayoutTagD, + int AlignmentD, + class Schedule +> +struct CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + TileShape_MNK, + ClusterShape_MNK, + EpilogueTileType, + ElementAccumulator, + ElementCompute, + ElementC_, + GmemLayoutTagC_, + AlignmentC, + ElementD, + GmemLayoutTagD, + AlignmentD, + Schedule, + cute::enable_if_t>> { + + // Passing void C disables source load + using ElementC = cute::conditional_t, + ElementD, ElementC_>; // prevents cute breakages + using GmemLayoutTagC = cute::conditional_t, + GmemLayoutTagD, GmemLayoutTagC_>; + static constexpr thread::ScaleType::Kind ScaleType = cute::is_void_v ? + thread::ScaleType::OnlyAlphaScaling : thread::ScaleType::Default; + + static constexpr int FragmentSize = 1; + using ThreadOp = thread::LinearCombination< + ElementD, FragmentSize, ElementAccumulator, ElementCompute, + ScaleType, FloatRoundStyle::round_to_nearest, ElementC>; + + using CollectiveOp = mctlass::epilogue::collective::detail::Sm90TmaWarpSpecializedAdapter< + mctlass::epilogue::collective::DefaultEpilogue< + mctlass::gemm::TagToStrideC_t, + mctlass::gemm::TagToStrideC_t, + ThreadOp, + mctlass::gemm::EpilogueDefault> + >; +}; + +// Tma warp-specialized builder +template < + class TileShape_MNK, + class ClusterShape_MNK, + class EpilogueTileType, + class ElementAccumulator, + class ElementCompute, + class ElementC_, + class GmemLayoutTagC, + int AlignmentC, + class ElementD, + class GmemLayoutTagD, + int AlignmentD, + class Schedule +> +struct CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + TileShape_MNK, + ClusterShape_MNK, + EpilogueTileType, + ElementAccumulator, + ElementCompute, + ElementC_, + GmemLayoutTagC, + AlignmentC, + ElementD, + GmemLayoutTagD, + AlignmentD, + Schedule, + cute::enable_if_t || + cute::is_same_v >> { +public: + using ElementC = cute::conditional_t,ElementD,ElementC_>; // prevents void ref breakages + static constexpr thread::ScaleType::Kind ScaleType = cute::is_void_v ? + thread::ScaleType::OnlyAlphaScaling : thread::ScaleType::Default; + + static constexpr int FragmentSize = 4; + using ThreadOp = thread::LinearCombination< + ElementD, FragmentSize, ElementAccumulator, ElementCompute, + ScaleType, FloatRoundStyle::round_to_nearest, ElementC>; + +private: + static constexpr int StagesC = 1; + static constexpr int StagesD = 2; + static constexpr bool DisableReuseSmemC = true; + using Impl = detail::TmaBuilderImpl< + TileShape_MNK, ClusterShape_MNK, EpilogueTileType, ElementAccumulator, ElementCompute, + ElementC_, GmemLayoutTagC, AlignmentC, ElementD, GmemLayoutTagD, AlignmentD, + Schedule, ThreadOp, mctlass::epilogue::Sm90TmaWarpSpecialized>; + +public: + using CollectiveOp = typename Impl::CollectiveOp; +}; + +// Auto builder +template < + class TileShape_MNK, + class ClusterShape_MNK, + class EpilogueTileType, + class ElementAccumulator, + class ElementCompute, + class ElementC, + class GmemLayoutTagC, + int AlignmentC, + class ElementD, + class GmemLayoutTagD, + int AlignmentD, + class Schedule +> +struct CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + TileShape_MNK, + ClusterShape_MNK, + EpilogueTileType, + ElementAccumulator, + ElementCompute, + ElementC, + GmemLayoutTagC, + AlignmentC, + ElementD, + GmemLayoutTagD, + AlignmentD, + Schedule, + cute::enable_if_t>> { + +private: + static constexpr bool IsTmaAligned = mctlass::gemm::collective::detail::is_aligned< + ElementC, AlignmentC, ElementD, AlignmentD, mctlass::gemm::collective::detail::tma_alignment_bytes>(); + + // Current TMA epilogues require sixteen-bit data types and epilogue tile M to be of size 64. + // Only dispatch to the TMA builder if these requirements are satisfied. + static constexpr bool IsSixteenBit = sizeof_bits::value == 16 && sizeof_bits::value == 16; + static constexpr bool IsEpiTileM64 = size<0>(shape(TileShape_MNK{})) == 64; + + using _CollectiveBuilder = CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + TileShape_MNK, + ClusterShape_MNK, + EpilogueTileType, + ElementAccumulator, + ElementCompute, + ElementC, + GmemLayoutTagC, + AlignmentC, + ElementD, + GmemLayoutTagD, + AlignmentD, + cute::conditional_t + >; + +public: + using ThreadOp = typename _CollectiveBuilder::ThreadOp; + using CollectiveOp = typename _CollectiveBuilder::CollectiveOp; +}; + +// Tma warp-specialized builder for elementwise fusion +template < + class TileShape_MNK, + class ClusterShape_MNK, + class EpilogueTileType, + class ElementAccumulator, + class ElementCompute, + class ElementC, + class GmemLayoutTagC, + int AlignmentC, + class ElementD, + class GmemLayoutTagD, + int AlignmentD, + class Schedule +> +struct CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + TileShape_MNK, + ClusterShape_MNK, + EpilogueTileType, + ElementAccumulator, + ElementCompute, + ElementC, + GmemLayoutTagC, + AlignmentC, + ElementD, + GmemLayoutTagD, + AlignmentD, + Schedule, + cute::enable_if_t || + cute::is_base_of_v >> { +public: + static constexpr int FragmentSize = 4; + using ThreadOp = thread::LinearCombinationGeneric< + Schedule::ActivationFunctor, + ElementD, FragmentSize, + ElementAccumulator, ElementCompute, Schedule::Scale, + Schedule::Round>; + +private: + static constexpr int StagesC = 1; + static constexpr int StagesD = 2; + static constexpr bool DisableReuseSmemC = true; + using Impl = detail::TmaBuilderImpl< + TileShape_MNK, ClusterShape_MNK, EpilogueTileType, ElementAccumulator, ElementCompute, + ElementC, GmemLayoutTagC, AlignmentC, ElementD, GmemLayoutTagD, AlignmentD, + Schedule, ThreadOp, mctlass::epilogue::Sm90TmaWarpSpecialized>; + +public: + using CollectiveOp = typename Impl::CollectiveOp; +}; + +// Tma warp-specialized builder for bias + elementwise fusion +template < + class TileShape_MNK, + class ClusterShape_MNK, + class EpilogueTileType, + class ElementAccumulator, + class ElementCompute, + class ElementC_, + class GmemLayoutTagC, + int AlignmentC, + class ElementD, + class GmemLayoutTagD, + int AlignmentD, + class Schedule +> +struct CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + TileShape_MNK, + ClusterShape_MNK, + EpilogueTileType, + ElementAccumulator, + ElementCompute, + ElementC_, + GmemLayoutTagC, + AlignmentC, + ElementD, + GmemLayoutTagD, + AlignmentD, + Schedule, + cute::enable_if_t || + cute::is_base_of_v >> { +private: + // Passing void C disables source load + using ElementC = cute::conditional_t, ElementD, ElementC_>; // prevents void ref breakages + +public: + static constexpr int FragmentSize = 4; + using ThreadOp = thread::LinearCombinationBiasElementwise< + ElementC, ElementAccumulator, ElementCompute, ElementD, typename Schedule::ElementT, FragmentSize, + typename Schedule::ActivationFunctor, typename Schedule::BiasOp, + Schedule::StoreT, typename Schedule::ElementBias>; + +private: + static constexpr int StagesC = 1; + static constexpr int StagesD = 2; + using Impl = detail::TmaBuilderImpl< + TileShape_MNK, ClusterShape_MNK, EpilogueTileType, ElementAccumulator, ElementCompute, + ElementC_, GmemLayoutTagC, AlignmentC, ElementD, GmemLayoutTagD, AlignmentD, + Schedule, ThreadOp, mctlass::epilogue::Sm90TmaWarpSpecializedBiasElementwise>; + +public: + using CollectiveOp = typename Impl::CollectiveOp; +}; + +// CollectiveBuilder that transposed epilogue below is used for sm90 gmma RS TT kernels +// since swapping NNN kernels input matrix and transposing its output at the same time then +// we can get TTN kernel. +template < + class TileShape_MNK, + class ClusterShape_MNK, + class EpilogueTileType, + class ElementAccumulator, + class ElementCompute, + class ElementC_, + class GmemLayoutTagC_, + int AlignmentC, + class ElementD, + class GmemLayoutTagD, + int AlignmentD, + class Schedule +> +struct CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + TileShape_MNK, + ClusterShape_MNK, + EpilogueTileType, + ElementAccumulator, + ElementCompute, + ElementC_, + GmemLayoutTagC_, + AlignmentC, + ElementD, + GmemLayoutTagD, + AlignmentD, + Schedule, + cute::enable_if_t>> { + // Passing void C disables source load + using ElementC = cute::conditional_t, + ElementD, ElementC_>; // prevents cute breakages + using GmemLayoutTagC = cute::conditional_t, + GmemLayoutTagD, GmemLayoutTagC_>; + static constexpr thread::ScaleType::Kind ScaleType = cute::is_void_v ? + thread::ScaleType::OnlyAlphaScaling : thread::ScaleType::Default; + + static constexpr int FragmentSize = 1; + using ThreadOp = thread::LinearCombination< + ElementD, FragmentSize, ElementAccumulator, ElementCompute, + ScaleType, FloatRoundStyle::round_to_nearest, ElementC>; + + using CollectiveOp = mctlass::epilogue::collective::detail::Sm90TmaWarpSpecializedAdapter< + mctlass::epilogue::collective::DefaultEpilogue< + mctlass::gemm::TagToStrideC_t, + mctlass::gemm::TagToStrideC_t, + ThreadOp, + mctlass::gemm::EpilogueTransposed> + >; +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::epilogue::collective diff --git a/csrc/mctlass/include/mctlass/epilogue/collective/collective_builder.hpp b/csrc/mctlass/include/mctlass/epilogue/collective/collective_builder.hpp new file mode 100644 index 0000000..7c13a18 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/collective/collective_builder.hpp @@ -0,0 +1,77 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/detail/dependent_false.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::epilogue::collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Used to specify epilogue subtile shape or dispatch to automatic computation of subtile shape +struct EpilogueTileAuto {}; + +// Used to let the builder pick the epilogue schedule automatically. +// Can be overridden with kernel schedule tags in mctlass/gemm/dispatch_policy.hpp +struct EpilogueScheduleAuto {}; + +template < + class ArchTag, + class OpClass, + class TileShape_MNK, + class ClusterShape_MNK, + class EpilogueTileType, + class ElementAccumulator, + class ElementCompute, + class ElementC, + class GmemLayoutTagC, + int AlignmentC, + class ElementD, + class GmemLayoutTagD, + int AlignmentD, + class Schedule, + class Enable = void +> +struct CollectiveBuilder { + static_assert(mctlass::detail::dependent_false, + "Could not build a collective epilogue for given parameters."); +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::epilogue::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// #include "builders/sm90_builder.inl" +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/collective/collective_epilogue.hpp b/csrc/mctlass/include/mctlass/epilogue/collective/collective_epilogue.hpp new file mode 100644 index 0000000..f3c802d --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/collective/collective_epilogue.hpp @@ -0,0 +1,55 @@ +/*************************************************************************************************** + * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::epilogue::collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + class DispatchPolicy, + class... Args +> +class CollectiveEpilogue { + static_assert(mctlass::detail::dependent_false, "Could not find an epilogue specialization."); +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::epilogue::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "detail.hpp" +#include "default_epilogue.hpp" +#include "epilogue_tensor_broadcast.hpp" +#include "sm70_epilogue_vectorized.hpp" +#include "sm90_epilogue_tma_warpspecialized.hpp" +#include "sm90_epilogue_tma_warpspecialized_bias_elementwise.hpp" +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/collective/default_epilogue.hpp b/csrc/mctlass/include/mctlass/epilogue/collective/default_epilogue.hpp new file mode 100644 index 0000000..9b5cb60 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/collective/default_epilogue.hpp @@ -0,0 +1,216 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing elementwise operations used by epilogues. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/dispatch_policy.hpp" +#include "mctlass/epilogue/collective/detail.hpp" + +#include "cute/tensor.hpp" +#include "cute/numeric/int.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies an element wise operation to all elements within the fragment +/// and writes them out to destination storage. +template < + class StrideC_, + class StrideD_, + class ThreadEpilogueOp_, + class EpilogueSchedule_ +> +class DefaultEpilogue { +public: + // + // Type Aliases + // + using EpilogueSchedule = EpilogueSchedule_; + + // derived types of output thread level operator + using ThreadEpilogueOp = ThreadEpilogueOp_; + using ElementOutput = typename ThreadEpilogueOp::ElementOutput; + using ElementAccumulator = typename ThreadEpilogueOp::ElementAccumulator; + using ElementCompute = typename ThreadEpilogueOp::ElementCompute; + using ElementScalar = ElementCompute; + using ElementC = typename ThreadEpilogueOp::ElementC; + using StrideC = StrideC_; + using ElementD = typename ThreadEpilogueOp::ElementD; + using StrideD = StrideD_; + + using GmemTiledCopyC = void; + using GmemTiledCopyD = void; + + static const int kOutputAlignment = ThreadEpilogueOp::kCount; + using AlignmentType = typename cute::uint_bit::value * kOutputAlignment>::type; + + static_assert(rank(StrideC{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + static_assert(rank(StrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + + struct SharedStorage { }; + + // Host side epilgoue arguments + struct Arguments { + typename ThreadEpilogueOp::Params thread{}; + ElementC const* ptr_C = nullptr; + StrideC dC{}; + ElementD* ptr_D = nullptr; + StrideD dD{}; + }; + + // Device side epilogue params + using Params = Arguments; + + // + // Methods + // + + template + static constexpr Params + to_underlying_arguments( + [[maybe_unused]] ProblemShape const& _, + Arguments const& args, + [[maybe_unused]] void* workspace) { + return args; + } + + MCTLASS_HOST_DEVICE + DefaultEpilogue(Params const& params_) + : params(params_), epilogue_op(params_.thread) { } + + MCTLASS_DEVICE + bool + is_source_needed() { + return epilogue_op.is_source_needed(); + } + + template< + class ProblemShapeMNKL, + class BlockShapeMNK, + class BlockCoordMNKL, + class FrgEngine, class FrgLayout, + class TiledMma, + class ResidueMNK + > + MCTLASS_HOST_DEVICE void + operator()( + ProblemShapeMNKL problem_shape_mnkl, + BlockShapeMNK blk_shape_MNK, + BlockCoordMNKL blk_coord_mnkl, + cute::Tensor const& accumulators, + TiledMma tiled_mma, + ResidueMNK residue_mnk, + int thread_idx, + [[maybe_unused]] char* smem_buf) + { + using namespace cute; + using X = Underscore; + + static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4"); + static_assert(is_static::value, "ThreadBlock tile shape must be static"); + static_assert(rank(BlockShapeMNK{}) == 3, "BlockShapeMNK must be rank 3"); + static_assert(rank(BlockCoordMNKL{}) == 4, "BlockCoordMNKL must be rank 3"); + + // Separate out problem shape for convenience + auto M = get<0>(problem_shape_mnkl); + auto N = get<1>(problem_shape_mnkl); + auto L = get<3>(problem_shape_mnkl); + + auto stride_c = detail::get_epilogue_stride(params.dC); + auto stride_d = detail::get_epilogue_stride(params.dD); + + // Represent the full output tensor + Tensor mC_mnl = make_tensor(make_gmem_ptr(params.ptr_C), make_shape(M,N,L), stride_c); // (m,n,l) + Tensor mD_mnl = make_tensor(make_gmem_ptr(params.ptr_D), make_shape(M,N,L), stride_d); // (m,n,l) + Tensor gC_mnl = local_tile(mC_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{}); // (BLK_M,BLK_N,m,n,l) + Tensor gD_mnl = local_tile(mD_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{}); // (BLK_M,BLK_N,m,n,l) + + // Slice to get the tile this CTA is responsible for + auto [m_coord, n_coord, k_coord, l_coord] = blk_coord_mnkl; + Tensor gC = gC_mnl(_,_,m_coord,n_coord,l_coord); // (BLK_M,BLK_N) + Tensor gD = gD_mnl(_,_,m_coord,n_coord,l_coord); // (BLK_M,BLK_N) + + // Partition source and destination tiles to match the accumulator partitioning + auto thr_mma = tiled_mma.get_thread_slice(thread_idx); + Tensor tCgD = thr_mma.partition_C(gD); // (VEC,THR_M,THR_N) + Tensor tCgC = thr_mma.partition_C(gC); // (VEC,THR_M,THR_N) + + static_assert(is_static::value, "Accumulator layout must be static"); + CUTE_STATIC_ASSERT_V(size(tCgC) == size(tCgD), + "Source and destination must have the same number of elements."); + CUTE_STATIC_ASSERT_V(size(tCgD) == size(accumulators), + "Accumulator count must have the same destination element count."); + + // Make an identity coordinate tensor for predicating our output MN tile + auto cD = make_identity_tensor(make_shape(unwrap(shape<0>(gD)), unwrap(shape<1>(gD)))); + Tensor tCcD = thr_mma.partition_C(cD); + + // source is needed + if (epilogue_op.is_source_needed()) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(accumulators); ++i) { + if (elem_less(tCcD(i), make_coord(get<0>(residue_mnk), get<1>(residue_mnk)))) { + tCgD(i) = epilogue_op(accumulators(i), tCgC(i)); + } + } + } + // source is not needed, avoid load + else { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(accumulators); ++i) { + if (elem_less(tCcD(i), make_coord(get<0>(residue_mnk), get<1>(residue_mnk)))) { + tCgD(i) = epilogue_op(accumulators(i)); + } + } + } + } + +private: + Params params; + ThreadEpilogueOp epilogue_op; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace collective +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/collective/detail.hpp b/csrc/mctlass/include/mctlass/epilogue/collective/detail.hpp new file mode 100644 index 0000000..46a872c --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/collective/detail.hpp @@ -0,0 +1,211 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/pipeline/pipeline.hpp" +#include "mctlass/gemm/dispatch_policy.hpp" +#include "mctlass/epilogue/dispatch_policy.hpp" + +#include "cute/tensor.hpp" +#include "cute/numeric/int.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace collective { + +namespace detail { + +template +static constexpr int elements_per_access_v = mctlass::sizeof_bits::value / mctlass::sizeof_bits::value; + +template +static constexpr bool sm90_is_cooperative_v = + std::is_base_of_v; + +template +static constexpr bool sm90_is_warp_specialized_v = + std::is_base_of_v; + +template +struct EmptyStorage { + MCTLASS_HOST_DEVICE + T* data() { return nullptr; } +}; + +template +MCTLASS_HOST_DEVICE +auto get_epilogue_stride(Stride stride){ + if constexpr (cute::is_base_of_v) { + return cute::make_stride(cute::get<1>(stride), cute::get<0>(stride), cute::get<2>(stride)); + } + else { + return stride; + } +} + +template +struct IsThreadEpilogueOpWithBias { + static constexpr bool value = false; + using type = typename ThreadEpilogueOp::ElementCompute; +}; + +template +struct IsThreadEpilogueOpWithBias > { + static constexpr bool value = true; + using type = typename ThreadEpilogueOp::ElementBias; +}; + +// IF_EPILOGUE_USES_TMA::value will be true only if: +// class T has member CopyOpS2G and T::CopyOpS2G is true +template +struct IF_EPILOGUE_USES_TMA { static constexpr bool value = false; }; + +template +struct IF_EPILOGUE_USES_TMA > +{ static constexpr bool value = true; }; + +// Wrapper class to use operator-style epilogues in sm90 TMA warp-specialized kernels +template +class Sm90TmaWarpSpecializedAdapter : public EpilogueOp { +public: + using LoadPipeline = mctlass::PipelineTransactionAsync<0>; // 0 stage to disable smem alloc + using LoadPipelineState = mctlass::PipelineState<0>; + constexpr static uint32_t TmaTransactionBytes = 0; + + using StorePipeline = mctlass::PipelineTmaStore<1>; // tma store pipe has no smem alloc + using StorePipelineState = mctlass::PipelineState<1>; + + using TensorStorage = typename EpilogueOp::SharedStorage; + using PipelineStorage = typename LoadPipeline::SharedStorage; + + template + MCTLASS_HOST_DEVICE + static constexpr int + get_load_pipe_increment([[maybe_unused]] TileShapeMNK) { + return 1; + } + + template + MCTLASS_HOST_DEVICE + static constexpr int + get_store_pipe_increment([[maybe_unused]] TileShapeMNK) { + return 1; + } + + MCTLASS_DEVICE + static void prefetch_tma_descriptors([[maybe_unused]] typename EpilogueOp::Params const&) + { + } + + // ctor inheritance + using EpilogueOp::EpilogueOp; + + template< + class ProblemShapeMNKL, + class TileShapeMNK, + class TileCoordMNKL, + class TiledMma + > + MCTLASS_DEVICE void + load( + [[maybe_unused]] LoadPipeline load_pipeline, + [[maybe_unused]] LoadPipelineState load_pipe_producer_state, + [[maybe_unused]] ProblemShapeMNKL problem_shape_mnkl, + [[maybe_unused]] TileShapeMNK tile_shape_MNK, + [[maybe_unused]] TileCoordMNKL tile_coord_mnkl, + [[maybe_unused]] TiledMma tiled_mma, + [[maybe_unused]] int thread_idx, + [[maybe_unused]] TensorStorage& shared_tensors) + { + // source load is performed in epilogue operator + } + + MCTLASS_DEVICE void + load_tail( + [[maybe_unused]] LoadPipeline load_pipeline, + [[maybe_unused]] LoadPipelineState load_pipe_producer_state) + { + } + + template< + class ProblemShapeMNKL, + class TileShapeMNK, + class TileCoordMNKL, + class AccEngine, class AccLayout, + class TiledMma + > + MCTLASS_DEVICE void + store( + [[maybe_unused]] LoadPipeline load_pipeline, + [[maybe_unused]] LoadPipelineState load_pipe_consumer_state, + [[maybe_unused]] StorePipeline store_pipeline, + [[maybe_unused]] StorePipelineState store_pipe_producer_state, + ProblemShapeMNKL problem_shape_mnkl, + TileShapeMNK tile_shape_MNK, + TileCoordMNKL tile_coord_mnkl, + cute::Tensor accumulators, + TiledMma tiled_mma, + int thread_idx, + TensorStorage& shared_tensors) + { + constexpr int BLK_M_RANK = rank<0>(tile_shape_MNK); + auto m_max_coord = unwrap(cute::transform(make_seq{}, [&](auto i) { + return get<0,i>(problem_shape_mnkl) - get<0,i>(tile_shape_MNK) * get<0,i>(tile_coord_mnkl); + })); + + constexpr int BLK_N_RANK = rank<1>(tile_shape_MNK); + auto n_max_coord = unwrap(cute::transform(make_seq{}, [&](auto i) { + return get<1,i>(problem_shape_mnkl) - get<1,i>(tile_shape_MNK) * get<1,i>(tile_coord_mnkl); + })); + + auto residue_mnk = make_tuple(m_max_coord, n_max_coord, Int<0>{}); + + (*this)( + problem_shape_mnkl, + tile_shape_MNK, + tile_coord_mnkl, + accumulators, + tiled_mma, + residue_mnk, + thread_idx, + reinterpret_cast(&shared_tensors)); + } + +}; + +} // namespace detail +} // namespace collective +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/collective/epilogue_tensor_broadcast.hpp b/csrc/mctlass/include/mctlass/epilogue/collective/epilogue_tensor_broadcast.hpp new file mode 100644 index 0000000..2c3a401 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/collective/epilogue_tensor_broadcast.hpp @@ -0,0 +1,245 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Functor for performing tensor-tensor broadacasts atop existing epilogues. + + Concretely, the opeartion performed is the following: + UnaryOp( + BinaryOp1( + BinaryOp0( + Activation((alpha * A @ B) + bias), + beta * C0 + ), + beta * C1 + ) + ) + + where: + - C0 and C1 have the same extents as the output + - BinaryOp0 and BinaryOp1 perform elementwise binary operations + - UnaryOp is an elementwise operation +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/epilogue/collective/detail.hpp" + +#include "cute/tensor.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace collective { +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Collective epilogue that applies elementwise tensor-tensor operations atop other epilogues +/// +template < + class StrideC_, + class StrideD_, + class ThreadEpilogueOp_, + class EpilogueSchedule_ +> +class EpilogueTensorBroadcast { +public: + // + // Type Aliases + // + using EpilogueSchedule = EpilogueSchedule_; + + // derived types of output thread level operator + using ThreadEpilogueOp = ThreadEpilogueOp_; + using ElementOutput = typename ThreadEpilogueOp::ElementOutput; + using ElementAccumulator = typename ThreadEpilogueOp::ElementAccumulator; + using ElementCompute = typename ThreadEpilogueOp::ElementCompute; + using ElementScalar = ElementCompute; + using ElementBias = typename ThreadEpilogueOp::ElementBias; + using ElementC = typename ThreadEpilogueOp::ElementC; + using StrideC = StrideC_; + using ElementD = typename ThreadEpilogueOp::ElementD; + using StrideD = StrideD_; + using ActivationFunctor = typename ThreadEpilogueOp::ActivationFunctor; + + static_assert(rank(StrideC{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + static_assert(rank(StrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + + static constexpr int kOutputAlignment = ThreadEpilogueOp::kCount; + using AlignmentType = typename cute::uint_bit::value * kOutputAlignment>::type; + + static constexpr bool IsBinaryOp0Enabled = ThreadEpilogueOp::IsBinaryOp0Enabled; + static constexpr bool IsBinaryOp1Enabled = ThreadEpilogueOp::IsBinaryOp1Enabled; + static constexpr bool IsUnaryOpEnabled = ThreadEpilogueOp::IsUnaryOpEnabled; + + struct SharedStorage { }; + + // Host side epilogue arguments + struct Arguments { + typename ThreadEpilogueOp::Params thread{}; + StrideC dC{}; + ElementD* ptr_D = nullptr; + StrideD dD{}; + ElementBias* ptr_Bias = nullptr; + ElementC* ptr_C0 = nullptr; + ElementC* ptr_C1 = nullptr; + }; + + // Device side epilogue params + using Params = Arguments; + + // + // Methods + // + + template + static constexpr Params + to_underlying_arguments( + [[maybe_unused]] ProblemShape const& _, + Arguments const& args, + [[maybe_unused]] void* workspace) { + return args; + } + + MCTLASS_HOST_DEVICE + EpilogueTensorBroadcast(Params const& params_) + : params(params_), epilogue_op(params_.thread) { } + + MCTLASS_DEVICE + bool + is_source_needed() { + return epilogue_op.is_source0_needed() || epilogue_op.is_source1_needed(); + } + + template< + class ProblemShapeMNKL, + class BlockShapeMNK, + class BlockCoordMNKL, + class FrgEngine, class FrgLayout, + class TiledMma, + class ResidueMNK + > + MCTLASS_HOST_DEVICE void + operator()( + ProblemShapeMNKL problem_shape_mnkl, + BlockShapeMNK blk_shape_MNK, + BlockCoordMNKL blk_coord_mnkl, + cute::Tensor const& accumulators, + TiledMma tiled_mma, + ResidueMNK residue_mnk, + int thread_idx, + [[maybe_unused]] char* smem_buf) + { + using namespace cute; + using X = Underscore; + + static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4"); + static_assert(is_static::value, "ThreadBlock tile shape must be static"); + static_assert(rank(BlockShapeMNK{}) == 3, "BlockShapeMNK must be rank 3"); + static_assert(rank(BlockCoordMNKL{}) == 4, "BlockCoordMNKL must be rank 4"); + + // Separate out problem shape for convenience + auto M = get<0>(problem_shape_mnkl); + auto N = get<1>(problem_shape_mnkl); + auto L = get<3>(problem_shape_mnkl); + + auto stride_c = detail::get_epilogue_stride(params.dC); + auto stride_d = detail::get_epilogue_stride(params.dD); + auto stride_bias = detail::get_epilogue_stride(Stride<_1, _0, _0>{}); + + // Represent the full output tensor + Tensor mC0_mnl = make_tensor(make_gmem_ptr(params.ptr_C0), make_shape(M,N,L), stride_c); // (m,n,l) + Tensor mC1_mnl = make_tensor(make_gmem_ptr(params.ptr_C1), make_shape(M,N,L), stride_c); // (m,n,l) + Tensor mD_mnl = make_tensor(make_gmem_ptr(params.ptr_D), make_shape(M,N,L), stride_d); // (m,n,l) + Tensor mBias_mnl = make_tensor(make_gmem_ptr(params.ptr_Bias), make_shape(M,N,L), stride_bias); // (m,n,l) + + Tensor gC0_mnl = local_tile(mC0_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{}); // (BLK_M,BLK_N,m,n,l) + Tensor gC1_mnl = local_tile(mC1_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{}); // (BLK_M,BLK_N,m,n,l) + + Tensor gD_mnl = local_tile(mD_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{}); // (BLK_M,BLK_N,m,n,l) + Tensor gBias_mnl = local_tile(mBias_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{}); // (BLK_M,BLK_N,m,n,l) + + // Slice to get the tile this thread block is responsible for + auto [m_coord, n_coord, k_coord, l_coord] = blk_coord_mnkl; + Tensor gC0 = gC0_mnl(_,_,m_coord,n_coord,l_coord); // (BLK_M,BLK_N) + Tensor gC1 = gC1_mnl(_,_,m_coord,n_coord,l_coord); // (BLK_M,BLK_N) + Tensor gD = gD_mnl(_,_,m_coord,n_coord,l_coord); // (BLK_M,BLK_N) + Tensor gBias = gBias_mnl(_,_,m_coord,n_coord,l_coord); // (BLK_M,BLK_N) + + // Partition source and destination tiles to match the accumulator partitioning + auto thr_mma = tiled_mma.get_thread_slice(thread_idx); + Tensor tCgD = thr_mma.partition_C(gD); // (VEC,THR_M,THR_N) + Tensor tCgC0 = thr_mma.partition_C(gC0); // (VEC,THR_M,THR_N) + Tensor tCgC1 = thr_mma.partition_C(gC1); // (VEC,THR_M,THR_N) + Tensor tCgBias = thr_mma.partition_C(gBias); // (VEC,THR_M,THR_N) + + static_assert(is_static::value, + "Accumulator layout must be static"); + CUTE_STATIC_ASSERT_V(size(tCgC0) == size(tCgD), + "Source and destination must have the same number of elements."); + CUTE_STATIC_ASSERT_V(size(tCgC1) == size(tCgD), + "Source and destination must have the same number of elements."); + CUTE_STATIC_ASSERT_V(size(tCgD) == size(accumulators), + "Accumulator count must have the same destination element count."); + CUTE_STATIC_ASSERT_V(size(tCgBias) == size(accumulators), + "Accumulator count must have the same destination element count."); + + auto cD = make_identity_tensor(make_shape(unwrap(shape<0>(gD)), unwrap(shape<1>(gD)))); + Tensor tCcD = thr_mma.partition_C(cD); + + bool bias_needed = params.ptr_Bias != nullptr; + bool c0_needed = (params.ptr_C0 != nullptr) && epilogue_op.is_source0_needed(); + bool c1_needed = (params.ptr_C1 != nullptr) && epilogue_op.is_source1_needed(); + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(accumulators); ++i) { + if (elem_less(tCcD(i), make_coord(get<0>(residue_mnk), get<1>(residue_mnk)))) { + ElementBias bias = bias_needed ? tCgBias(i) : ElementBias(0); + ElementC c0 = c0_needed ? tCgC0(i) : ElementC(0); + ElementC c1 = c1_needed ? tCgC1(i) : ElementC(0); + + tCgD(i) = epilogue_op(accumulators(i), c0, c1, bias); + } + } + } + +private: + Params params; + ThreadEpilogueOp epilogue_op; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace collective +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/collective/sm70_epilogue_vectorized.hpp b/csrc/mctlass/include/mctlass/epilogue/collective/sm70_epilogue_vectorized.hpp new file mode 100644 index 0000000..7ca4cb5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/collective/sm70_epilogue_vectorized.hpp @@ -0,0 +1,335 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing elementwise operations used by epilogues. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "cute/tensor.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies an element wise operation to all elements within the fragment +/// and writes it out to destination storage. +/// +/// Ways to generalize this: +/// - CTA tile shape +/// - vectorization requirements (GMEM) +/// - vectoriz(able) transform() +/// +template < + class StrideC_, + class StrideD_, + class ThreadEpilogueOp_, + class SmemLayout_, + class CopyAtomR2S_, + class TiledCopyS2R_, + class CopyAtomR2G_ +> +class Epilogue { +public: + // + // Type Aliases + // + // derived types of output thread level operator + using ThreadEpilogueOp = ThreadEpilogueOp_; + using ElementAccumulator = typename ThreadEpilogueOp::ElementAccumulator; + using ElementCompute = typename ThreadEpilogueOp::ElementCompute; + using ElementScalar = ElementCompute; + using ElementOutput = typename ThreadEpilogueOp::ElementOutput; + using ElementC = typename ThreadEpilogueOp::ElementC; + using StrideC = StrideC_; + using ElementD = typename ThreadEpilogueOp::ElementD; + using StrideD = StrideD_; + + using SmemLayout = SmemLayout_; + using CopyAtomR2S = CopyAtomR2S_; + using TiledCopyS2R = TiledCopyS2R_; + using CopyAtomR2G = CopyAtomR2G_; + + static const int kOutputAlignment = ThreadEpilogueOp::kCount; + using AlignmentType = typename cute::uint_bit::value * kOutputAlignment>::type; + + static_assert(rank(StrideC{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + static_assert(rank(StrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + + struct SharedStorage + { + cute::array_aligned> smem_epilogue; + }; + + // Host side epilogue arguments + struct Arguments { + typename ThreadEpilogueOp::Params thread{}; + ElementC const* ptr_C = nullptr; + StrideC dC{}; + ElementD* ptr_D = nullptr; + StrideD dD{}; + }; + + // Device side epilogue params + using Params = Arguments; + + // + // Methods + // + + template + static constexpr Params + to_underlying_arguments( + [[maybe_unused]] ProblemShape const& _, + Arguments const& args, + [[maybe_unused]] void* workspace) { + return args; + } + + MCTLASS_HOST_DEVICE + Epilogue(Params const& params_) + : params(params_), epilogue_op(params_.thread) { } + + MCTLASS_DEVICE + bool + is_source_needed() { + return epilogue_op.is_source_needed(); + } + + template< + class ProblemShapeMNKL, + class BlockShapeMNK, + class BlockCoordMNKL, + class FrgEngine, class FrgLayout, + class TiledMma, + class ResidueMNK + > + MCTLASS_DEVICE void + operator()( + ProblemShapeMNKL problem_shape_mnkl, + BlockShapeMNK blk_shape_MNK, + BlockCoordMNKL blk_coord_mnkl, + cute::Tensor const& accumulators, // (MMA,MMA_M,MMA_N) + TiledMma tiled_mma, + ResidueMNK residue_mnk, + int thread_idx, + char* smem_buf) + { + using namespace cute; + using X = Underscore; + + static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4"); + static_assert(is_static::value, "ThreadBlock tile shape must be static"); + static_assert(rank(BlockShapeMNK{}) == 3, "BlockShapeMNK must be rank 3"); + static_assert(rank(BlockCoordMNKL{}) == 4, "BlockCoordMNKL must be rank 3"); + + // synchronizing function for smem reads/writes +#if CUDA_BARRIER_ENABLED + auto synchronize = [] () { mctlass::arch::NamedBarrier::sync(typename TiledCopyS2R::TiledNumThr{}, 0); }; +#else + auto synchronize = [] () { __syncthreads(); }; +#endif + + // Separate out problem shape for convenience + auto M = get<0>(problem_shape_mnkl); + auto N = get<1>(problem_shape_mnkl); + auto L = get<3>(problem_shape_mnkl); + + // Represent the full output tensor + Tensor mC_mnl = make_tensor(make_gmem_ptr(params.ptr_C), make_shape(M,N,L), params.dC); // (m,n,l) + Tensor mD_mnl = make_tensor(make_gmem_ptr(params.ptr_D), make_shape(M,N,L), params.dD); // (m,n,l) + Tensor gC_mnl = local_tile(mC_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{}); // (BLK_M,BLK_N,m,n,l) + Tensor gD_mnl = local_tile(mD_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{}); // (BLK_M,BLK_N,m,n,l) + + // Slice to get the tile this CTA is responsible for + auto [m_coord, n_coord, k_coord, l_coord] = blk_coord_mnkl; + Tensor gC = gC_mnl(_,_,m_coord,n_coord,l_coord); // (BLK_M,BLK_N) + Tensor gD = gD_mnl(_,_,m_coord,n_coord,l_coord); // (BLK_M,BLK_N) + + // Construct a tensor in SMEM that we can partition for rearranging data + SharedStorage& storage = *reinterpret_cast(smem_buf); + Tensor sC = make_tensor(make_smem_ptr(storage.smem_epilogue.data()), SmemLayout{}); // (SMEM_M,SMEM_N) + + // Partition sC to match the accumulator partitioning + auto tiled_r2s = make_tiled_copy_C(CopyAtomR2S{}, tiled_mma); + auto tC = tiled_r2s.get_thread_slice(thread_idx); + Tensor tCaC = tC.retile_S(accumulators); // ((Atom,AtomNum), MMA_M, MMA_N) + Tensor tCsC = tC.partition_D(sC); // ((Atom,AtomNum),PIPE_M,PIPE_N) + + // Tile gD and gC by the shape of SmemLayout first + auto tile = make_shape(size<0>(sC), size<1>(sC)); + Tensor gCt = local_tile(gC, tile, _); // (SMEM_M,SMEM_N,TILE_M,TILE_N) + Tensor gDt = local_tile(gD, tile, _); // (SMEM_M,SMEM_N,TILE_M,TILE_N) + + // Partition sC, gC, and gD for the output + auto tiled_s2r = TiledCopyS2R{}; + auto tD = tiled_s2r.get_thread_slice(thread_idx); + Tensor tDsC = tD.partition_S(sC); // ((Atom,AtomNum),ATOM_M,ATOM_N) + Tensor tDgC = tD.partition_D(gCt); // ((Atom,AtomNum),ATOM_M,ATOM_N,TILE_M,TILE_N) + Tensor tDgD = tD.partition_D(gDt); // ((Atom,AtomNum),ATOM_M,ATOM_N,TILE_M,TILE_N) + + // Allocate intermediate registers on the dst tensors + Tensor tDrC = make_tensor(take<0,3>(shape(tDgC))); // ((Atom,AtomNum),ATOM_M,ATOM_N) + Tensor tDrD = make_tensor(shape(tDrC)); // ((Atom,AtomNum),ATOM_M,ATOM_N) + + // Repeat the D-partitioning for coordinates and predication + Tensor cD = make_identity_tensor(make_shape(size<0>(gD),size<1>(gD))); // (BLK_M,BLK_N) -> (blk_m,blk_n) + Tensor cDt = local_tile(cD, tile, _); // (SMEM_M,SMEM_N,TILE_M,TILE_N) + Tensor tDcD = tD.partition_D(cDt); // ((Atom,AtomNum),ATOM_M,ATOM_N,TILE_M,TILE_N) + + CUTE_STATIC_ASSERT(size<1>(tCaC) % size<3>(tDgC) == 0); // TILE_M divides MMA_M + CUTE_STATIC_ASSERT(size<2>(tCaC) % size<4>(tDgC) == 0); // TILE_N divides MMA_N + CUTE_STATIC_ASSERT(typename TiledCopyS2R::TiledNumThr{} == size<0>(typename TiledMma::AtomLayoutC_TV{})); + +#if 0 + if (thread_idx == 0 && m_coord == 0 && n_coord == 0) { + print("aC : "); print(accumulators.layout()); print("\n"); + print("gC : "); print(gC.layout()); print("\n"); + print("gD : "); print(gD.layout()); print("\n"); + print("sC : "); print(sC.layout()); print("\n"); + print("\n"); + print("tCsC : "); print(tCsC.layout()); print("\n"); + print("tCaC : "); print(tCaC.layout()); print("\n"); + print("\n"); + print("gDt : "); print(gDt.layout()); print("\n"); + print("tDsC : "); print(tDsC.layout()); print("\n"); + print("tDrC : "); print(tDrC.layout()); print("\n"); + print("\n"); + print("tDrD : "); print(tDrD.layout()); print("\n"); + print("tDgC : "); print(tDgC.layout()); print("\n"); + print("tDgD : "); print(tDgD.layout()); print("\n"); + print("\n"); + } +#endif + + // For each tiling needed for SmemLayout to cover shape(gD) + MCTLASS_PRAGMA_UNROLL + for (int step_m = 0; step_m < size<2>(cDt); ++step_m) + { + MCTLASS_PRAGMA_UNROLL + for (int step_n = 0; step_n < size<3>(cDt); ++step_n) + { + // Step 1. Copy to SMEM + MCTLASS_PRAGMA_UNROLL + for (int pipe_m = 0; pipe_m < size<1>(tCsC); ++pipe_m) { + MCTLASS_PRAGMA_UNROLL + for (int pipe_n = 0; pipe_n < size<2>(tCsC); ++pipe_n) { + int mma_m = step_m * size<1>(tCsC) + pipe_m; + int mma_n = step_n * size<2>(tCsC) + pipe_n; + + copy(tiled_r2s, tCaC(_,mma_m,mma_n), tCsC(_,pipe_m,pipe_n)); + } + } + + // Step 2. Wait for SMEM writes to complete + synchronize(); + + // Step 3. Copy from SMEM into a fragment + copy(tiled_s2r, tDsC, tDrC); + + // Step 4. Wait for SMEM reads to complete + synchronize(); + + Tensor tDgDmn = tDgD(_,_,_,step_m,step_n); + Tensor tDcDmn = tDcD(_,_,_,step_m,step_n); + + if (epilogue_op.is_source_needed()) { + // source is needed + Tensor tDgCmn = tDgC(_,_,_,step_m,step_n); + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < size<1>(tDgDmn); ++m) + { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < size<2>(tDgDmn); ++n) + { + // Predication + if (get<0>(tDcDmn(0,m,n)) < get<0>(residue_mnk) && + get<1>(tDcDmn(0,m,n)) < get<1>(residue_mnk)) + { + // Step 5. Elementwise operation with conversion + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<0>(tDrC); ++i) { + tDrD(i,m,n) = epilogue_op(tDrC(i,m,n), tDgCmn(i,m,n)); + } + // Step 6. Copy to GMEM + copy(CopyAtomR2G{}, tDrD(_,m,n), tDgDmn(_,m,n)); + } + } + } + } + else { + // source is not needed, avoid load and lift compute + + // Step 5. Elementwise operation with conversion + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tDrC); ++i) { + tDrD(i) = epilogue_op(tDrC(i)); + } + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < size<1>(tDgDmn); ++m) + { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < size<2>(tDgDmn); ++n) + { + // Predication + if (get<0>(tDcDmn(0,m,n)) < get<0>(residue_mnk) && + get<1>(tDcDmn(0,m,n)) < get<1>(residue_mnk)) + { + // Step 6. Copy to GMEM + copy(CopyAtomR2G{}, tDrD(_,m,n), tDgDmn(_,m,n)); + } + } + } + } + } + } + } + +private: + Params params; + ThreadEpilogueOp epilogue_op; +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace collective +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp b/csrc/mctlass/include/mctlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp new file mode 100644 index 0000000..2756746 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp @@ -0,0 +1,606 @@ +/*************************************************************************************************** + * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing elementwise operations used by epilogues. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/arch/barrier.h" +#include "mctlass/epilogue/dispatch_policy.hpp" +#include "mctlass/epilogue/collective/detail.hpp" +#include "mctlass/epilogue/thread/scale_type.h" + +#include "cute/tensor.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int StagesC_, + int StagesD_, + bool DisableSmemReuseC_, + class BlockTileShape_, // (BLK_M,BLK_N,BLK_K) + class EpilogueTileShape_, // (EPI_TILE_M,EPI_TILE_N) + class ElementC_, + class StrideC_, + class ElementD_, + class StrideD_, + class ThreadEpilogueOp_, + class CopyOpG2S_, + class SmemLayoutAtomC_, + class CopyOpS2R_, + class CopyOpS2G_, + class SmemLayoutAtomD_, + class CopyOpR2S_ +> +class CollectiveEpilogue< + Sm90TmaWarpSpecialized, + BlockTileShape_, + EpilogueTileShape_, + ElementC_, + StrideC_, + ElementD_, + StrideD_, + ThreadEpilogueOp_, + CopyOpG2S_, + SmemLayoutAtomC_, + CopyOpS2R_, + CopyOpS2G_, + SmemLayoutAtomD_, + CopyOpR2S_ +> { +public: + // + // Type Aliases + // + using DispatchPolicy = Sm90TmaWarpSpecialized; + using BlockTileShape = BlockTileShape_; + using EpilogueTileShape = EpilogueTileShape_; + using ThreadEpilogueOp = ThreadEpilogueOp_; + using ElementAccumulator = typename ThreadEpilogueOp::ElementAccumulator; + using ElementCompute = typename ThreadEpilogueOp::ElementCompute; + using ElementScalar = ElementCompute; + using ElementBias = typename detail::IsThreadEpilogueOpWithBias::type; + using ElementOutput = typename ThreadEpilogueOp::ElementOutput; + using ElementC = ElementC_; + using StrideC = StrideC_; + using ElementD = ElementD_; + using StrideD = StrideD_; + + using CopyOpG2S = CopyOpG2S_; + using SmemLayoutAtomC = SmemLayoutAtomC_; + using CopyOpS2R = CopyOpS2R_; + using CopyOpS2G = CopyOpS2G_; + using SmemLayoutAtomD = SmemLayoutAtomD_; + using CopyOpR2S = CopyOpR2S_; + + using GmemTiledCopyC = SM90_TMA_LOAD; + using GmemTiledCopyD = SM90_TMA_STORE; + + constexpr static int kOutputAlignment = ThreadEpilogueOp::kCount; + constexpr static bool iskThreadEpilogueOpWithBias = detail::IsThreadEpilogueOpWithBias::value; + using AlignmentType = typename uint_bit::value * kOutputAlignment>::type; + + static_assert(!is_layout::value && is_tuple::value, "EpilogueTileShape must be a cute::Shape"); + static_assert(rank(BlockTileShape{}) == 3, "BlockTileShape must be rank-3: [BLK_M,BLK_N,BLK_K]"); + static_assert(rank(EpilogueTileShape{}) == 2, "EpilogueTileShape must be rank-2: [EPI_TILE_M,EPI_TILE_N]"); + static_assert(rank(StrideC{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + static_assert(rank(StrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + +private: + using InternalElementC = cute::conditional_t,ElementD,ElementC>; // prevents void ref breakages + constexpr static int StagesC = StagesC_; + constexpr static int StagesD = StagesD_; + constexpr static bool is_source_supported = ThreadEpilogueOp::kScale == mctlass::epilogue::thread::ScaleType::Default || + ThreadEpilogueOp::kScale == mctlass::epilogue::thread::ScaleType::NoBetaScaling; + static_assert((cute::is_void_v && not is_source_supported) || (not cute::is_void_v && is_source_supported), + "Inconsistent C type and Scale kind"); + + // internal optimization to reuse C shared memory for storing D + using SmemLayoutAtomBitsC = decltype(downcast::value>(SmemLayoutAtomC{})); + using SmemLayoutAtomBitsD = decltype(downcast::value>(SmemLayoutAtomD{})); + constexpr static bool ReuseSmemC = not DispatchPolicy::DisableSmemReuseC && + is_source_supported && + sizeof(InternalElementC) == sizeof(ElementD) && + StrideC{} == StrideD{} && + cute::is_same_v; + +public: + using SmemLayoutC = decltype(tile_to_shape( + SmemLayoutAtomC{}, + make_shape(size<0>(BlockTileShape{}), size<1>(BlockTileShape{}), Int{}), + cute::conditional_t(StrideC{}) == 1, Step<_2,_1,_3>, Step<_1,_2,_3>>{} )); + using SmemLayoutD = decltype(tile_to_shape( + SmemLayoutAtomD{}, + make_shape(size<0>(EpilogueTileShape{}), size<1>(EpilogueTileShape{}), Int{}), + cute::conditional_t(StrideD{}) == 1, Step<_2,_1,_3>, Step<_1,_2,_3>>{} )); + + // TMA pipeline for loading C + using LoadPipeline = mctlass::PipelineTransactionAsync; + using LoadPipelineState = mctlass::PipelineState; + constexpr static uint32_t TmaTransactionBytes = + size(take<0,2>(SmemLayoutC{})) * static_cast(sizeof(InternalElementC)); + + // TMA pipeline for storing D + using StorePipeline = mctlass::PipelineTmaStore; + using StorePipelineState = mctlass::PipelineState; + + struct SharedStorage { + struct TensorStorage : aligned_struct<128> { + cute::conditional_t, + array_aligned> smem_C; + alignas(128) cute::conditional_t, + array_aligned> smem_D; + } tensors; + + using PipelineStorage = typename LoadPipeline::SharedStorage; + PipelineStorage pipeline; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + using PipelineStorage = typename SharedStorage::PipelineStorage; + + // Host side epilogue arguments + struct Arguments { + typename ThreadEpilogueOp::Params thread; + ElementC const* ptr_C; + StrideC dC; + ElementD const* ptr_D; + StrideD dD; + }; + + // Device side epilgoue params + struct Params { + using TMA_C = decltype(make_tma_copy( + CopyOpG2S{}, + make_tensor(static_cast(nullptr), + repeat_like(StrideC{}, int32_t(0)), StrideC{}), + SmemLayoutC{}(_,_,0))); + using TMA_D = decltype(make_tma_copy( + CopyOpS2G{}, + make_tensor(static_cast(nullptr), + repeat_like(StrideD{}, int32_t(0)), StrideD{}), + SmemLayoutD{}(_,_,0))); + + typename ThreadEpilogueOp::Params thread{}; + TMA_C tma_load_c; + TMA_D tma_store_d; + }; + + // + // Methods + // + + template + static constexpr Params + to_underlying_arguments( + ProblemShape const& problem_shape, + Arguments const& args, + [[maybe_unused]] void* workspace) { + // Optionally append _1s until problem shape is rank-4 in case its is only rank-3 (MNK) + auto problem_shape_MNKL = append<4>(problem_shape, Int<1>{}); + auto M = get<0>(problem_shape_MNKL); + auto N = get<1>(problem_shape_MNKL); + auto L = get<3>(problem_shape_MNKL); + + typename Params::TMA_C tma_load_c = [&]() { + if constexpr (not cute::is_void_v) { + Tensor tensor_c = make_tensor(static_cast(args.ptr_C), make_layout(make_shape(M,N,L), args.dC)); + return make_tma_copy( + CopyOpG2S{}, + tensor_c, + SmemLayoutC{}(_,_,0)); + } + else { + return typename Params::TMA_C{}; + } + }(); + + Tensor tensor_d = make_tensor(args.ptr_D, make_layout(make_shape(M,N,L), args.dD)); + typename Params::TMA_D tma_store_d = make_tma_copy( + CopyOpS2G{}, + tensor_d, + SmemLayoutD{}(_,_,0)); + + return { + args.thread, + tma_load_c, + tma_store_d + }; + } + + template + MCTLASS_HOST_DEVICE + static constexpr int + get_load_pipe_increment(TileShapeMNK tile_shape_MNK) { + // Compute number of C subtiles (currently always one) + constexpr int epi_m = size<0>(tile_shape_MNK) / size<0>(SmemLayoutC{}); + constexpr int epi_n = size<1>(tile_shape_MNK) / size<1>(SmemLayoutC{}); + + return epi_m * epi_n; + } + + template + MCTLASS_HOST_DEVICE + static constexpr int + get_store_pipe_increment(TileShapeMNK tile_shape_MNK) { + if constexpr (ReuseSmemC) { + return get_load_pipe_increment(tile_shape_MNK); + } + + // Compute number of D subtiles + constexpr int epi_m = size<0>(tile_shape_MNK) / size<0>(SmemLayoutD{}); + constexpr int epi_n = size<1>(tile_shape_MNK) / size<1>(SmemLayoutD{}); + + return epi_m * epi_n; + } + + MCTLASS_HOST_DEVICE + CollectiveEpilogue(Params const& params_) + : params(params_), epilogue_op(params_.thread) { } + + MCTLASS_DEVICE + bool + is_source_needed() { + return epilogue_op.is_source_needed(); + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + MCTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& epilogue_params) { + cute::prefetch_tma_descriptor(epilogue_params.tma_load_c.get_tma_descriptor()); + cute::prefetch_tma_descriptor(epilogue_params.tma_store_d.get_tma_descriptor()); + } + + template< + class ProblemShapeMNKL, + class TileShapeMNK, + class TileCoordMNKL, + class TiledMma + > + MCTLASS_DEVICE void + load( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_producer_state, + ProblemShapeMNKL problem_shape_mnkl, + TileShapeMNK tile_shape_MNK, + TileCoordMNKL tile_coord_mnkl, + TiledMma tiled_mma, + [[maybe_unused]] int thread_idx, + TensorStorage& shared_tensors) { + using namespace cute; + using X = Underscore; + + int warp_idx = canonical_warp_idx(); + int warp_idx_in_warp_group = warp_idx % 4; + int lane_predicate = cute::elect_one_sync(); + + auto M = get<0>(problem_shape_mnkl); + auto N = get<1>(problem_shape_mnkl); + auto L = get<3>(problem_shape_mnkl); + auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl; + + // Represent the full source tensor + Tensor mC_mnl = params.tma_load_c.get_tma_tensor(make_shape(M,N,L)); // (m,n,l) + Tensor gC_mnl = local_tile(mC_mnl, tile_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{}); // (TILE_M,TILE_N,m,n,l) + // Slice to get the gmem tile of C (gC) this CTA is currently responsible for + Tensor gC = gC_mnl(_,_,m_coord,n_coord,l_coord); // (TILE_M,TILE_N) + // Get the corresponding smem tile of C (sC) + Tensor sC = make_tensor(make_smem_ptr(shared_tensors.smem_C.data()), SmemLayoutC{}); // (TILE_M,TILE_N,PIPE) + + // Prepare the thread(b)lock (G)mem to (S)mem TMA copy (bGS_) + ThrCopy thrblk_g2s = params.tma_load_c.get_slice(Int<0>{}); + Tensor bGS_gC = thrblk_g2s.partition_S(gC); // (TMA,TMA_M,TMA_N) + Tensor bGS_sC = thrblk_g2s.partition_D(sC); // (TMA,TMA_M,TMA_N,PIPE) + + auto* tma_barrier = load_pipeline.producer_get_barrier(load_pipe_producer_state); + uint16_t mcast_mask = 0; + + // Execute the TMA load for C + if (warp_idx_in_warp_group == 0 and lane_predicate) { + load_pipeline.producer_acquire(load_pipe_producer_state); + copy(params.tma_load_c.with(*tma_barrier, mcast_mask), bGS_gC, bGS_sC(_,_,_,load_pipe_producer_state.index())); + load_pipeline.producer_commit(load_pipe_producer_state); + } + } + + MCTLASS_DEVICE void + load_tail( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_producer_state) { + int warp_idx = canonical_warp_idx(); + int warp_idx_in_warp_group = warp_idx % 4; + int lane_predicate = cute::elect_one_sync(); + + if (warp_idx_in_warp_group == 0 and lane_predicate) { + load_pipeline.producer_tail(load_pipe_producer_state); + } + } + + template< + class ProblemShapeMNKL, + class TileShapeMNK, + class TileCoordMNKL, + class AccEngine, class AccLayout, + class TiledMma + > + MCTLASS_DEVICE void + store( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_consumer_state, + StorePipeline store_pipeline, + StorePipelineState store_pipe_producer_state, + ProblemShapeMNKL problem_shape_mnkl, + TileShapeMNK tile_shape_MNK, + TileCoordMNKL tile_coord_mnkl, + cute::Tensor accumulators, + TiledMma tiled_mma, + int thread_idx, + TensorStorage& shared_tensors) { + using namespace cute; + using X = Underscore; + + static_assert(is_rmem::value, "Accumulator must be RF resident."); + static_assert(rank(AccLayout{}) == 3, "Accumulator must be MMA-partitioned: (MMA,MMA_M,MMA_N)"); + static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4"); + static_assert(is_static::value, "TileShapeMNK must be static"); + static_assert(rank(TileShapeMNK{}) == 3, "TileShapeMNK must be rank 3"); + static_assert(rank(TileCoordMNKL{}) == 4, "TileCoordMNKL must be rank 4"); + + // Separate out problem shape for convenience + auto M = get<0>(problem_shape_mnkl); + auto N = get<1>(problem_shape_mnkl); + auto L = get<3>(problem_shape_mnkl); + auto mma_tile_m = size<0>(typename TiledMma::TiledShape_MNK{}); + auto mma_tile_n = size<1>(typename TiledMma::TiledShape_MNK{}); + auto epi_tile_m = size<0>(EpilogueTileShape{}); + auto epi_tile_n = size<1>(EpilogueTileShape{}); + + // Represent the full output tensor + Tensor mD_mnl = params.tma_store_d.get_tma_tensor(make_shape(M,N,L)); // (m,n,l) + Tensor gD_mnl = local_tile(mD_mnl, tile_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{}); // (TILE_M,TILE_N,m,n,l) + + // Slice to get the tile this CTA is responsible for + auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl; + Tensor gD = gD_mnl(_,_,m_coord,n_coord,l_coord); // (TILE_M,TILE_N) + + // Construct the smem tensors for source (sC) and output (sD) + Tensor sC = make_tensor(make_smem_ptr(shared_tensors.smem_C.data()), // (TILE_M,TILE_N) + SmemLayoutC{})(_,_,load_pipe_consumer_state.index()); + Tensor bEsD = make_tensor(make_smem_ptr(shared_tensors.smem_D.data()), // (EPI_TILE_M,EPI_TILE_N,PIPE) + SmemLayoutD{}); + + // Tile thread(b)lock tensors by (E)pilogue output tile shape (bE) + Tensor bEsC = local_tile(sC, EpilogueTileShape{}, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + Tensor bEgD = local_tile(gD, EpilogueTileShape{}, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + + // Partition for register to smem copy (tRS_) + using CopyAtomR2S = cute::conditional_t, + Copy_Atom>,ElementD>, + Copy_Atom>; + TiledCopy tiled_r2s = make_tiled_copy_C_atom(CopyAtomR2S{}, tiled_mma); + ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx); + Tensor tRS_rAcc = thread_r2s.retile_S(accumulators); // ((R2S,R2S_V),MMA_M,MMA_N) + Tensor tRS_sD = conditional_return( + thread_r2s.partition_D(recast(bEsC)), // (R2S,R2S_M,R2S_N,EPI_M,EPI_N) + thread_r2s.partition_D(bEsD) ); // (R2S,R2S_M,R2S_N,PIPE) + + // Allocate register tensors + auto tRS_rD_shape = take<0,3>(shape(thread_r2s.partition_S(bEsD))); // (R2S,R2S_M,R2S_N) + Tensor tRS_rC = make_tensor(tRS_rD_shape); // (R2S,R2S_M,R2S_N) + Tensor tRS_rD = make_tensor(tRS_rD_shape); // (R2S,R2S_M,R2S_N) + + // Vectorized fragment view for thread epilogue op + Tensor tRS_rAcc_frg = recast(tRS_rAcc); + Tensor tRS_rC_frg = recast(tRS_rC); + Tensor tRS_rD_frg = recast(tRS_rD); + + // Partition for smem to register copy (tSR_) + TiledCopy tiled_s2r = make_tiled_copy_S(Copy_Atom{}, tiled_r2s); + ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx); + Tensor tSR_sC = thread_s2r.partition_S(bEsC); // (S2R,S2R_M,S2R_N,EPI_M,EPI_N) + Tensor tSR_rC = thread_s2r.retile_D(tRS_rC); // (S2R,S2R_M,S2R_N) + + // Partition for smem to gmem copy (tSG_) + ThrCopy thrblk_s2g = params.tma_store_d.get_slice(Int<0>{}); + Tensor tSG_sD = conditional_return( + thrblk_s2g.partition_S(recast(bEsC)), // (S2G,S2G_M,S2G_N,EPI_M,EPI_N) + thrblk_s2g.partition_S(bEsD) ); // (S2G,S2G_M,S2G_N,PIPE) + Tensor tSG_gD = thrblk_s2g.partition_D(bEgD); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N) + + CUTE_STATIC_ASSERT(size<0>(tRS_rAcc) % ThreadEpilogueOp::kCount == 0, "ThreadEpilogueOp does not vectorize properly"); + CUTE_STATIC_ASSERT(mma_tile_m == epi_tile_m, "EPI_TILE_M must equal MMA_TILE_M"); + CUTE_STATIC_ASSERT(mma_tile_n % epi_tile_n == 0, "EPI_TILE_N must divide MMA_TILE_N"); + + // Thread synchronizer for previously issued waits or fences + // to ensure visibility of smem reads/writes to threads or TMA unit + auto synchronize = [&] () { mctlass::arch::NamedBarrier::sync(size(TiledMma{}), 0); }; + + // Predication for TMA store (one warp issues TMA store) + bool issue_tma_store = (thread_idx / NumThreadsPerWarp) == 0; + + if (epilogue_op.is_source_needed()) { + // Wait for epilogue load to fill smem buffer with C + load_pipeline.consumer_wait(load_pipe_consumer_state); + } + + // Delay issue of TMA store by 1 iteration to achieve better instruction pipelining + PipelineState store_pipe_producer_state_prev = store_pipe_producer_state; + int epi_m_prev = 0, epi_n_prev = 0; + + // For each output tile + MCTLASS_PRAGMA_UNROLL + for (int epi_n = 0; epi_n < size<3>(bEgD); ++epi_n) { + MCTLASS_PRAGMA_UNROLL + for (int epi_m = 0; epi_m < size<2>(bEgD); ++epi_m) { + // The current tile in accumulator + int mma_m = epi_m; + int mma_n = (epi_n * epi_tile_n) / mma_tile_n; + Tensor tRS_rAcc_frg_mn = tRS_rAcc_frg(_,mma_m,mma_n); + + // Elementwise operation with conversion + int r2s_v = epi_n * size(tRS_rD_frg); + if (epilogue_op.is_source_needed()) { + // Copy source tile to register from smem + if constexpr (cute::is_same_v) { + copy(tSR_sC(_,_,_,epi_m,epi_n), tSR_rC); + } + else { + copy(tiled_s2r, tSR_sC(_,_,_,epi_m,epi_n), tSR_rC); + } + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tRS_rD_frg); ++i) { + tRS_rD_frg(i) = epilogue_op(tRS_rAcc_frg_mn(r2s_v + i), tRS_rC_frg(i)); + } + } + else { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tRS_rD_frg); ++i) { + tRS_rD_frg(i) = epilogue_op(tRS_rAcc_frg_mn(r2s_v + i)); + } + } + + if constexpr (ReuseSmemC) { + // Issue the TMA store of the previous iteration + if (not (epi_m == 0 && epi_n == 0)) { + // Make sure smem writes are visible to TMA + mctlass::arch::fence_view_async_shared(); + synchronize(); // ensure all threads have issued their async fence + + // Write the tile to gmem from smem with TMA + if (issue_tma_store) { + copy(params.tma_store_d, tSG_sD(_,_,_,epi_m_prev,epi_n_prev), tSG_gD(_,_,_,epi_m_prev,epi_n_prev)); + } + } + + // Copy output tile to smem from register + if constexpr (cute::is_same_v) { + copy(tRS_rD, tRS_sD(_,_,_,epi_m,epi_n)); + } + else { + copy(tiled_r2s, tRS_rD, tRS_sD(_,_,_,epi_m,epi_n)); + } + } + else { + // Issue the TMA store of the previous iteration + if (not (epi_m == 0 && epi_n == 0)) { + // Make sure smem writes are visible to TMA + mctlass::arch::fence_view_async_shared(); + synchronize(); // ensure all threads have issued their async fence + + // Write the tile to gmem from smem with TMA + if (issue_tma_store) { + copy(params.tma_store_d, tSG_sD(_,_,_,store_pipe_producer_state_prev.index()), tSG_gD(_,_,_,epi_m_prev,epi_n_prev)); + store_pipeline.producer_commit(store_pipe_producer_state_prev); + } + } + + // Wait for a smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + synchronize(); + + // Copy tile to smem from register + if constexpr (cute::is_same_v) { + copy(tRS_rD, tRS_sD(_,_,_,store_pipe_producer_state.index())); + } + else { + copy(tiled_r2s, tRS_rD, tRS_sD(_,_,_,store_pipe_producer_state.index())); + } + + // Advance pipeline state + store_pipe_producer_state_prev = store_pipe_producer_state; + ++store_pipe_producer_state; + } + + epi_m_prev = epi_m; + epi_n_prev = epi_n; + } + } + + if constexpr (ReuseSmemC) { + // Fence and issue the TMA store of the last iteration + mctlass::arch::fence_view_async_shared(); + synchronize(); // ensure all threads have issued their async fence + if (issue_tma_store) { + copy(params.tma_store_d, tSG_sD(_,_,_,epi_m_prev,epi_n_prev), tSG_gD(_,_,_,epi_m_prev,epi_n_prev)); + } + + // Arrive and advance pipeline state + if (issue_tma_store) { + store_pipeline.producer_commit(store_pipe_producer_state); + } + ++store_pipe_producer_state; + + // Wait for a smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + synchronize(); + + // Let dma warp know smem buffer is consumed and empty + if (epilogue_op.is_source_needed()) { + load_pipeline.consumer_release(store_pipe_producer_state); + } + } + else { + // Fence and issue the TMA store of the last iteration + mctlass::arch::fence_view_async_shared(); + synchronize(); // ensure all threads have issued their async fence + if (issue_tma_store) { + copy(params.tma_store_d, tSG_sD(_,_,_,store_pipe_producer_state_prev.index()), tSG_gD(_,_,_,epi_m_prev,epi_n_prev)); + store_pipeline.producer_commit(store_pipe_producer_state_prev); + } + + // Let dma warp know smem buffer is consumed and empty + if (epilogue_op.is_source_needed()) { + load_pipeline.consumer_release(load_pipe_consumer_state); + } + } + } + +private: + Params const& params; + ThreadEpilogueOp epilogue_op; +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace collective +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/collective/sm90_epilogue_tma_warpspecialized_bias_elementwise.hpp b/csrc/mctlass/include/mctlass/epilogue/collective/sm90_epilogue_tma_warpspecialized_bias_elementwise.hpp new file mode 100644 index 0000000..4037575 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/collective/sm90_epilogue_tma_warpspecialized_bias_elementwise.hpp @@ -0,0 +1,679 @@ +/*************************************************************************************************** + * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing pipelined epilogues with bias add and elementwise activation functions. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/arch/barrier.h" +#include "mctlass/epilogue/dispatch_policy.hpp" +#include "mctlass/epilogue/collective/detail.hpp" +#include "mctlass/epilogue/thread/scale_type.h" + +#include "cute/tensor.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int StagesC_, + int StagesD_, + class BlockTileShape_, // (BLK_M,BLK_N,BLK_K) + class EpilogueTileShape_, // (EPI_TILE_M,EPI_TILE_N) + class ElementC_, + class StrideC_, + class ElementD_, + class StrideD_, + class ThreadEpilogueOp_, + class CopyOpG2S_, + class SmemLayoutAtomC_, + class CopyOpS2R_, + class CopyOpS2G_, + class SmemLayoutAtomD_, + class CopyOpR2S_ +> +class CollectiveEpilogue< + Sm90TmaWarpSpecializedBiasElementwise, + BlockTileShape_, + EpilogueTileShape_, + ElementC_, + StrideC_, + ElementD_, + StrideD_, + ThreadEpilogueOp_, + CopyOpG2S_, + SmemLayoutAtomC_, + CopyOpS2R_, + CopyOpS2G_, + SmemLayoutAtomD_, + CopyOpR2S_ +> { +public: + // + // Type Aliases + // + using DispatchPolicy = Sm90TmaWarpSpecializedBiasElementwise; + using BlockTileShape = BlockTileShape_; + using EpilogueTileShape = EpilogueTileShape_; + using ThreadEpilogueOp = ThreadEpilogueOp_; + using ElementAccumulator = typename ThreadEpilogueOp::ElementAccumulator; + using ElementCompute = typename ThreadEpilogueOp::ElementCompute; + using ElementScalar = ElementCompute; + using ElementBias = typename detail::IsThreadEpilogueOpWithBias::type; + using ElementT = typename ThreadEpilogueOp::ElementT; + using ElementOutput = typename ThreadEpilogueOp::ElementOutput; + using ElementC = ElementC_; + using StrideC = StrideC_; + using ElementD = ElementD_; + using StrideD = StrideD_; + using ActivationFunctor = typename ThreadEpilogueOp::ActivationFunctor; + using BinaryOp = typename ThreadEpilogueOp::BinaryOp; + + using CopyOpG2S = CopyOpG2S_; + using SmemLayoutAtomC = SmemLayoutAtomC_; + using CopyOpS2R = CopyOpS2R_; + using CopyOpS2G = CopyOpS2G_; + using SmemLayoutAtomD = SmemLayoutAtomD_; + using CopyOpR2S = CopyOpR2S_; + + using GmemTiledCopyC = SM90_TMA_LOAD; + using GmemTiledCopyD = SM90_TMA_STORE; + + constexpr static bool StoreT = ThreadEpilogueOp::kStoreT; + constexpr static int kOutputAlignment = ThreadEpilogueOp::kCount; + static_assert(detail::IsThreadEpilogueOpWithBias::value, + "Epilogue dispatch policy Sm90TmaWarpSpecializedBiasElementwise requires the use of a thread-level epiogue that supports bias calculation"); + constexpr static bool iskThreadEpilogueOpWithBias = true; + using AlignmentType = typename uint_bit::value * kOutputAlignment>::type; + + static_assert(!is_layout::value && is_tuple::value, "EpilogueTileShape must be a cute::Shape"); + static_assert(rank(BlockTileShape{}) == 3, "BlockTileShape must be rank-3: [BLK_M,BLK_N,BLK_K]"); + static_assert(rank(EpilogueTileShape{}) == 2, "EpilogueTileShape must be rank-2: [EPI_TILE_M,EPI_TILE_N]"); + static_assert(rank(StrideC{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + static_assert(rank(StrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]"); + +private: + using InternalElementC = cute::conditional_t,ElementD,ElementC>; // prevents void ref breakages + constexpr static int StagesC = StagesC_; + constexpr static int StagesD = StagesD_; + constexpr static bool is_source_supported = not cute::is_void_v; + static_assert((cute::is_void_v && not is_source_supported) || (not cute::is_void_v && is_source_supported), + "Inconsistent C type and Scale kind"); + + // internal optimization to reuse C shared memory for storing D + using SmemLayoutAtomBitsC = decltype(downcast::value>(SmemLayoutAtomC{})); + using SmemLayoutAtomBitsD = decltype(downcast::value>(SmemLayoutAtomD{})); + constexpr static bool ReuseSmemC = is_source_supported && + sizeof(InternalElementC) == sizeof(ElementD) && + StrideC{} == StrideD{} && + cute::is_same_v && + not StoreT; + +public: + using SmemLayoutC = decltype(tile_to_shape( + SmemLayoutAtomC{}, + make_shape(size<0>(BlockTileShape{}), size<1>(BlockTileShape{}), Int{}), + cute::conditional_t(StrideC{}) == 1, Step<_2,_1,_3>, Step<_1,_2,_3>>{} )); + using SmemLayoutD = decltype(tile_to_shape( + SmemLayoutAtomD{}, + make_shape(size<0>(EpilogueTileShape{}), size<1>(EpilogueTileShape{}), Int{}), + cute::conditional_t(StrideD{}) == 1, Step<_2,_1,_3>, Step<_1,_2,_3>>{} )); + + // TMA pipeline for loading C + using LoadPipeline = mctlass::PipelineTransactionAsync; + using LoadPipelineState = mctlass::PipelineState; + constexpr static uint32_t TmaTransactionBytes = + size(take<0,2>(SmemLayoutC{})) * static_cast(sizeof(InternalElementC)); + + // TMA pipeline for storing D and T. ReuseSmemC cannot be set to true if StoreT is enabled. + using StorePipeline = mctlass::PipelineTmaStore; + using StorePipelineState = mctlass::PipelineState; + + struct SharedStorage { + struct TensorStorage : aligned_struct<128> { + cute::conditional_t, + array_aligned> smem_C; + alignas(128) cute::conditional_t, + array_aligned> smem_D; + alignas(128) cute::conditional_t, + array_aligned> smem_T; + } tensors; + + using PipelineStorage = typename LoadPipeline::SharedStorage; + PipelineStorage pipeline; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + using PipelineStorage = typename SharedStorage::PipelineStorage; + + // Host side epilogue arguments + struct Arguments { + typename ThreadEpilogueOp::Params thread; + ElementC const* ptr_C; + StrideC dC; + ElementD const* ptr_D; + StrideD dD; + ElementBias const* ptr_Bias = nullptr; + ElementT const* ptr_T = nullptr; + }; + + // Device side epilgoue params + struct Params { + using TMA_C = decltype(make_tma_copy( + CopyOpG2S{}, + make_tensor(static_cast(nullptr), + repeat_like(StrideC{}, int32_t(0)), StrideC{}), + SmemLayoutC{}(_,_,0))); + using TMA_D = decltype(make_tma_copy( + CopyOpS2G{}, + make_tensor(static_cast(nullptr), + repeat_like(StrideD{}, int32_t(0)), StrideD{}), + SmemLayoutD{}(_,_,0))); + using TMA_T = decltype(make_tma_copy( + CopyOpS2G{}, + make_tensor(static_cast(nullptr), + repeat_like(StrideD{}, int32_t(0)), StrideD{}), + SmemLayoutD{}(_,_,0))); + typename ThreadEpilogueOp::Params thread{}; + TMA_C tma_load_c; + TMA_D tma_store_d; + TMA_T tma_store_t; + ElementBias const* ptr_Bias = nullptr; + }; + + // + // Methods + // + + template + static constexpr Params + to_underlying_arguments( + ProblemShape const& problem_shape, + Arguments const& args, + [[maybe_unused]] void* workspace) { + // Optionally append _1s until problem shape is rank-4 in case its is only rank-3 (MNK) + auto problem_shape_MNKL = append<4>(problem_shape, Int<1>{}); + auto M = get<0>(problem_shape_MNKL); + auto N = get<1>(problem_shape_MNKL); + auto L = get<3>(problem_shape_MNKL); + + typename Params::TMA_C tma_load_c = [&]() { + if constexpr (not cute::is_void_v) { + Tensor tensor_c = make_tensor(static_cast(args.ptr_C), make_layout(make_shape(M,N,L), args.dC)); + return make_tma_copy( + CopyOpG2S{}, + tensor_c, + SmemLayoutC{}(_,_,0)); + } + else { + return typename Params::TMA_C{}; + } + }(); + + Tensor tensor_d = make_tensor(args.ptr_D, make_layout(make_shape(M,N,L), args.dD)); + typename Params::TMA_D tma_store_d = make_tma_copy( + CopyOpS2G{}, + tensor_d, + SmemLayoutD{}(_,_,0)); + + typename Params::TMA_T tma_store_t = [&]() { + if constexpr (StoreT) { + Tensor tensor_t = make_tensor(args.ptr_T, make_layout(make_shape(M,N,L), args.dD)); + return make_tma_copy( + CopyOpS2G{}, + tensor_t, + SmemLayoutD{}(_,_,0)); + } + else { + return typename Params::TMA_T{}; + } + }(); + + return { + args.thread, + tma_load_c, + tma_store_d, + tma_store_t, + args.ptr_Bias + }; + } + + template + MCTLASS_HOST_DEVICE + static constexpr int + get_load_pipe_increment(TileShapeMNK tile_shape_MNK) { + // Compute number of C subtiles (currently always one) + constexpr int epi_m = size<0>(tile_shape_MNK) / size<0>(SmemLayoutC{}); + constexpr int epi_n = size<1>(tile_shape_MNK) / size<1>(SmemLayoutC{}); + + return epi_m * epi_n; + } + + template + MCTLASS_HOST_DEVICE + static constexpr int + get_store_pipe_increment(TileShapeMNK tile_shape_MNK) { + if constexpr (ReuseSmemC) { + return get_load_pipe_increment(tile_shape_MNK); + } + + // Compute number of D subtiles + constexpr int epi_m = size<0>(tile_shape_MNK) / size<0>(SmemLayoutD{}); + constexpr int epi_n = size<1>(tile_shape_MNK) / size<1>(SmemLayoutD{}); + + return epi_m * epi_n; + } + + MCTLASS_HOST_DEVICE + CollectiveEpilogue(Params const& params_) + : params(params_), epilogue_op(params_.thread) { } + + MCTLASS_DEVICE + bool + is_source_needed() { + return is_source_supported && epilogue_op.is_source_needed(); + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + MCTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& epilogue_params) { + cute::prefetch_tma_descriptor(epilogue_params.tma_load_c.get_tma_descriptor()); + cute::prefetch_tma_descriptor(epilogue_params.tma_store_d.get_tma_descriptor()); + if constexpr (StoreT) { + cute::prefetch_tma_descriptor(epilogue_params.tma_store_t.get_tma_descriptor()); + } + } + + template< + class ProblemShapeMNKL, + class TileShapeMNK, + class TileCoordMNKL, + class TiledMma + > + MCTLASS_DEVICE void + load( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_producer_state, + ProblemShapeMNKL problem_shape_mnkl, + TileShapeMNK tile_shape_MNK, + TileCoordMNKL tile_coord_mnkl, + TiledMma tiled_mma, + [[maybe_unused]] int thread_idx, + TensorStorage& shared_tensors) { + using namespace cute; + using X = Underscore; + + int warp_idx = canonical_warp_idx(); + int warp_idx_in_warp_group = warp_idx % 4; + int lane_predicate = cute::elect_one_sync(); + + auto M = get<0>(problem_shape_mnkl); + auto N = get<1>(problem_shape_mnkl); + auto L = get<3>(problem_shape_mnkl); + auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl; + + // Represent the full source tensor + Tensor mC_mnl = params.tma_load_c.get_tma_tensor(make_shape(M,N,L)); // (m,n,l) + Tensor gC_mnl = local_tile(mC_mnl, tile_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{}); // (TILE_M,TILE_N,m,n,l) + // Slice to get the gmem tile of C (gC) this CTA is currently responsible for + Tensor gC = gC_mnl(_,_,m_coord,n_coord,l_coord); // (TILE_M,TILE_N) + // Get the corresponding smem tile of C (sC) + Tensor sC = make_tensor(make_smem_ptr(shared_tensors.smem_C.data()), SmemLayoutC{}); // (TILE_M,TILE_N,PIPE) + + // Prepare the thread(b)lock (G)mem to (S)mem TMA copy (bGS_) + ThrCopy thrblk_g2s = params.tma_load_c.get_slice(Int<0>{}); + Tensor bGS_gC = thrblk_g2s.partition_S(gC); // (TMA,TMA_M,TMA_N) + Tensor bGS_sC = thrblk_g2s.partition_D(sC); // (TMA,TMA_M,TMA_N,PIPE) + + auto* tma_barrier = load_pipeline.producer_get_barrier(load_pipe_producer_state); + uint16_t mcast_mask = 0; + + // Execute the TMA load for C + if (warp_idx_in_warp_group == 0 and lane_predicate) { + load_pipeline.producer_acquire(load_pipe_producer_state); + copy(params.tma_load_c.with(*tma_barrier, mcast_mask), bGS_gC, bGS_sC(_,_,_,load_pipe_producer_state.index())); + load_pipeline.producer_commit(load_pipe_producer_state); + } + } + + MCTLASS_DEVICE void + load_tail( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_producer_state) { + int warp_idx = canonical_warp_idx(); + int warp_idx_in_warp_group = warp_idx % 4; + int lane_predicate = cute::elect_one_sync(); + + if (warp_idx_in_warp_group == 0 and lane_predicate) { + load_pipeline.producer_tail(load_pipe_producer_state); + } + } + + template< + class ProblemShapeMNKL, + class TileShapeMNK, + class TileCoordMNKL, + class AccEngine, class AccLayout, + class TiledMma + > + MCTLASS_DEVICE void + store( + LoadPipeline load_pipeline, + LoadPipelineState load_pipe_consumer_state, + StorePipeline store_pipeline, + StorePipelineState store_pipe_producer_state, + ProblemShapeMNKL problem_shape_mnkl, + TileShapeMNK tile_shape_MNK, + TileCoordMNKL tile_coord_mnkl, + cute::Tensor accumulators, + TiledMma tiled_mma, + int thread_idx, + TensorStorage& shared_tensors) { + using namespace cute; + using X = Underscore; + + static_assert(is_rmem::value, "Accumulator must be RF resident."); + static_assert(rank(AccLayout{}) == 3, "Accumulator must be MMA-partitioned: (MMA,MMA_M,MMA_N)"); + static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4"); + static_assert(is_static::value, "TileShapeMNK must be static"); + static_assert(rank(TileShapeMNK{}) == 3, "TileShapeMNK must be rank 3"); + static_assert(rank(TileCoordMNKL{}) == 4, "TileCoordMNKL must be rank 4"); + + // Separate out problem shape for convenience + auto M = get<0>(problem_shape_mnkl); + auto N = get<1>(problem_shape_mnkl); + auto L = get<3>(problem_shape_mnkl); + auto mma_tile_m = size<0>(typename TiledMma::TiledShape_MNK{}); + auto mma_tile_n = size<1>(typename TiledMma::TiledShape_MNK{}); + auto epi_tile_m = size<0>(EpilogueTileShape{}); + auto epi_tile_n = size<1>(EpilogueTileShape{}); + + // Represent the full output tensor + Tensor mD_mnl = params.tma_store_d.get_tma_tensor(make_shape(M,N,L)); // (m,n,l) + Tensor gD_mnl = local_tile(mD_mnl, tile_shape_MNK, make_coord(_,_,_), Step<_1, _1, X>{}); // (TILE_M,TILE_N,m,n,l) + Tensor mT_mnl = params.tma_store_t.get_tma_tensor(make_shape(M,N,L)); // (m,n,l) + Tensor gT_mnl = local_tile(mT_mnl, tile_shape_MNK, make_coord(_,_,_), Step<_1, _1, X>{}); // (TILE_M,TILE_N,m,n,l) + Tensor mBias_mnl = make_tensor(make_gmem_ptr(params.ptr_Bias), make_shape(M,N,L), Stride<_1, _0, _0>{}); // (m,n,l) + Tensor gBias_mnl = local_tile(mBias_mnl, tile_shape_MNK, make_coord(_,_,_), Step<_1,_1,X>{}); // (TILE_M,TILE_N,m,n,l) + + // Slice to get the tile this CTA is responsible for + auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl; + Tensor gD = gD_mnl(_,_,m_coord,n_coord,l_coord); // (TILE_M,TILE_N) + Tensor gT = gT_mnl(_,_,m_coord,n_coord,l_coord); // (TILE_M,TILE_N) + Tensor gBias = gBias_mnl(_,_,m_coord,n_coord,l_coord); // (TILE_M,TILE_N) + + // Construct the smem tensors for source (sC) and output (sD, sT) + Tensor sC = make_tensor(make_smem_ptr(shared_tensors.smem_C.data()), // (TILE_M,TILE_N) + SmemLayoutC{})(_,_,load_pipe_consumer_state.index()); + Tensor bEsD = make_tensor(make_smem_ptr(shared_tensors.smem_D.data()), // (EPI_TILE_M,EPI_TILE_N,PIPE) + SmemLayoutD{}); + Tensor bEsT = make_tensor(make_smem_ptr(shared_tensors.smem_T.data()), // (EPI_TILE_M,EPI_TILE_N,PIPE) + SmemLayoutD{}); + + // Tile thread(b)lock tensors by (E)pilogue output tile shape (bE) + Tensor bEsC = local_tile(sC, EpilogueTileShape{}, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + Tensor bEgD = local_tile(gD, EpilogueTileShape{}, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + Tensor bEgT = local_tile(gT, EpilogueTileShape{}, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + Tensor bEgBias = local_tile(gBias, EpilogueTileShape{}, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N) + + // Partition for register to smem copy (tRS_) + using CopyAtomR2S = cute::conditional_t, + Copy_Atom>,ElementD>, + Copy_Atom>; + TiledCopy tiled_r2s = make_tiled_copy_C_atom(CopyAtomR2S{}, tiled_mma); + ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx); + Tensor tRS_rAcc = thread_r2s.retile_S(accumulators); // ((R2S,R2S_V),MMA_M,MMA_N) + Tensor tRS_sD = conditional_return( + thread_r2s.partition_D(recast(bEsC)), // (R2S,R2S_M,R2S_N,EPI_M,EPI_N) + thread_r2s.partition_D(bEsD) ); // (R2S,R2S_M,R2S_N,PIPE) + Tensor tRS_sT = thread_r2s.partition_D(bEsT); // (R2S,R2S_M,R2S_N,PIPE) + + // Allocate register tensors + auto tRS_rD_shape = take<0,3>(shape(thread_r2s.partition_S(bEsD))); // (R2S,R2S_M,R2S_N) + Tensor tRS_rC = make_tensor(tRS_rD_shape); // (R2S,R2S_M,R2S_N) + Tensor tRS_rD = make_tensor(tRS_rD_shape); // (R2S,R2S_M,R2S_N) + Tensor tRS_rT = make_tensor(tRS_rD_shape); // (R2S,R2S_M,R2S_N) + + Tensor tRS_gBias = thread_r2s.partition_S(bEgBias); // (R2S,R2S_M,R2S_N,EPI_M,EPI_N) + Tensor tRS_rBias = make_tensor(take<0,3>(shape(tRS_gBias))); // (R2S,R2S_M,R2S_N) + + // Vectorized fragment view for thread epilogue op + Tensor tRS_rAcc_frg = recast(tRS_rAcc); + Tensor tRS_rC_frg = recast(tRS_rC); + Tensor tRS_rD_frg = recast(tRS_rD); + Tensor tRS_rT_frg = recast(tRS_rT); + Tensor tRS_rBias_frg = recast(tRS_rBias); + + // Partition for smem to register copy (tSR_) + TiledCopy tiled_s2r = make_tiled_copy_S(Copy_Atom{}, tiled_r2s); + ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx); + Tensor tSR_sC = thread_s2r.partition_S(bEsC); // (S2R,S2R_M,S2R_N,EPI_M,EPI_N) + Tensor tSR_rC = thread_s2r.retile_D(tRS_rC); // (S2R,S2R_M,S2R_N) + + // Partition for smem to gmem copy (tSG_) + ThrCopy thrblk_s2g = params.tma_store_d.get_slice(Int<0>{}); + Tensor tSG_sD = conditional_return( + thrblk_s2g.partition_S(recast(bEsC)), // (S2G,S2G_M,S2G_N,EPI_M,EPI_N) + thrblk_s2g.partition_S(bEsD) ); // (S2G,S2G_M,S2G_N,PIPE) + Tensor tSG_gD = thrblk_s2g.partition_D(bEgD); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N) + + ThrCopy thrblk_s2g_t = params.tma_store_t.get_slice(Int<0>{}); + Tensor tSG_sT = thrblk_s2g_t.partition_S(bEsT); // (S2G,S2G_M,S2G_N,PIPE) + Tensor tSG_gT = thrblk_s2g_t.partition_D(bEgT); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N) + + CUTE_STATIC_ASSERT(size<0>(tRS_rAcc) % ThreadEpilogueOp::kCount == 0, "ThreadEpilogueOp does not vectorize properly"); + CUTE_STATIC_ASSERT(mma_tile_m == epi_tile_m, "EPI_TILE_M must equal MMA_TILE_M"); + CUTE_STATIC_ASSERT(mma_tile_n % epi_tile_n == 0, "EPI_TILE_N must divide MMA_TILE_N"); + + // Thread synchronizer for previously issued waits or fences + // to ensure visibility of smem reads/writes to threads or TMA unit + auto synchronize = [&] () { mctlass::arch::NamedBarrier::sync(size(TiledMma{}), 0); }; + + // Predication for TMA store (one warp issues TMA store) + bool issue_tma_store = (thread_idx / NumThreadsPerWarp) == 0; + + if (is_source_supported && epilogue_op.is_source_needed()) { + // Wait for epilogue load to fill smem buffer with C + load_pipeline.consumer_wait(load_pipe_consumer_state); + } + + // Delay issue of TMA store by 1 iteration to achieve better instruction pipelining + PipelineState store_pipe_producer_state_prev = store_pipe_producer_state; + int epi_m_prev = 0, epi_n_prev = 0; + + // For each output tile + MCTLASS_PRAGMA_UNROLL + for (int epi_n = 0; epi_n < size<3>(bEgD); ++epi_n) { + MCTLASS_PRAGMA_UNROLL + for (int epi_m = 0; epi_m < size<2>(bEgD); ++epi_m) { + // The current tile in accumulator + int mma_m = epi_m; + int mma_n = (epi_n * epi_tile_n) / mma_tile_n; + Tensor tRS_rAcc_frg_mn = tRS_rAcc_frg(_,mma_m,mma_n); + + // Copy bias to registers from gmem + copy(tRS_gBias(_,_,_,epi_m,epi_n), tRS_rBias); + + // Elementwise operation with conversion + int r2s_v = epi_n * size(tRS_rD_frg); + if (is_source_supported && epilogue_op.is_source_needed()) { + // Copy source tile to registers from smem + if constexpr (cute::is_same_v) { + copy(tSR_sC(_,_,_,epi_m,epi_n), tSR_rC); + } + else { + copy(tiled_s2r, tSR_sC(_,_,_,epi_m,epi_n), tSR_rC); + } + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tRS_rD_frg); ++i) { + epilogue_op(tRS_rD_frg(i), tRS_rT_frg(i), tRS_rAcc_frg_mn(r2s_v + i), tRS_rC_frg(i), tRS_rBias_frg(i)); + } + } + else { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tRS_rD_frg); ++i) { + epilogue_op(tRS_rD_frg(i), tRS_rT_frg(i), tRS_rAcc_frg_mn(r2s_v + i), tRS_rBias_frg(i)); + } + } + + if constexpr (ReuseSmemC) { + // If ReuseSmemC is true, StoreT must be false. Therefore, we do not perform copies for T in this block. + + // Issue the TMA store of the previous iteration + if (not (epi_m == 0 && epi_n == 0)) { + // Make sure smem writes are visible to TMA + mctlass::arch::fence_view_async_shared(); + synchronize(); // ensure all threads have issued their async fence + + // Write the tile to gmem from smem with TMA + if (issue_tma_store) { + copy(params.tma_store_d, tSG_sD(_,_,_,epi_m_prev,epi_n_prev), tSG_gD(_,_,_,epi_m_prev,epi_n_prev)); + } + } + + // Copy output tile to smem from register + if constexpr (cute::is_same_v) { + copy(tRS_rD, tRS_sD(_,_,_,epi_m,epi_n)); + } + else { + copy(tiled_r2s, tRS_rD, tRS_sD(_,_,_,epi_m,epi_n)); + } + } + else { + // Issue the TMA store of the previous iteration + if (not (epi_m == 0 && epi_n == 0)) { + // Make sure smem writes are visible to TMA + mctlass::arch::fence_view_async_shared(); + synchronize(); // ensure all threads have issued their async fence + + // Write the tile to gmem from smem with TMA + if (issue_tma_store) { + copy(params.tma_store_d, tSG_sD(_,_,_,store_pipe_producer_state_prev.index()), tSG_gD(_,_,_,epi_m_prev,epi_n_prev)); + if constexpr (StoreT) { + copy(params.tma_store_t, tSG_sT(_,_,_,store_pipe_producer_state_prev.index()), tSG_gT(_,_,_,epi_m_prev,epi_n_prev)); + } + store_pipeline.producer_commit(store_pipe_producer_state_prev); + } + } + + // Wait for a smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + synchronize(); + + // Copy tile to smem from register + if constexpr (cute::is_same_v) { + copy(tRS_rD, tRS_sD(_,_,_,store_pipe_producer_state.index())); + if constexpr (StoreT) { + copy(tRS_rT, tRS_sT(_,_,_,store_pipe_producer_state.index())); + } + } + else { + copy(tiled_r2s, tRS_rD, tRS_sD(_,_,_,store_pipe_producer_state.index())); + if constexpr (StoreT) { + copy(tiled_r2s, tRS_rT, tRS_sT(_,_,_,store_pipe_producer_state.index())); + } + } + + // Advance pipeline state + store_pipe_producer_state_prev = store_pipe_producer_state; + ++store_pipe_producer_state; + } + + epi_m_prev = epi_m; + epi_n_prev = epi_n; + } + } + + if constexpr (ReuseSmemC) { + // If ReuseSmemC is true, StoreT must be false. Therefore, we do not perform copies for T in this block. + + // Fence and issue the TMA store of the last iteration + mctlass::arch::fence_view_async_shared(); + synchronize(); // ensure all threads have issued their async fence + if (issue_tma_store) { + copy(params.tma_store_d, tSG_sD(_,_,_,epi_m_prev,epi_n_prev), tSG_gD(_,_,_,epi_m_prev,epi_n_prev)); + } + + // Arrive and advance pipeline state + if (issue_tma_store) { + store_pipeline.producer_commit(store_pipe_producer_state); + } + ++store_pipe_producer_state; + + // Wait for a smem buffer to be available + if (issue_tma_store) { + store_pipeline.producer_acquire(store_pipe_producer_state); + } + synchronize(); + + // Let dma warp know smem buffer is consumed and empty + if (is_source_supported && epilogue_op.is_source_needed()) { + load_pipeline.consumer_release(store_pipe_producer_state); + } + } + else { + // Fence and issue the TMA store of the last iteration + mctlass::arch::fence_view_async_shared(); + synchronize(); // ensure all threads have issued their async fence + if (issue_tma_store) { + copy(params.tma_store_d, tSG_sD(_,_,_,store_pipe_producer_state_prev.index()), tSG_gD(_,_,_,epi_m_prev,epi_n_prev)); + if (StoreT) { + copy(params.tma_store_t, tSG_sT(_,_,_,store_pipe_producer_state_prev.index()), tSG_gT(_,_,_,epi_m_prev,epi_n_prev)); + } + store_pipeline.producer_commit(store_pipe_producer_state_prev); + } + + // Let dma warp know smem buffer is consumed and empty + if (epilogue_op.is_source_needed()) { + load_pipeline.consumer_release(load_pipe_consumer_state); + } + } + } + +private: + Params const& params; + ThreadEpilogueOp epilogue_op; +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace collective +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/dispatch_policy.hpp b/csrc/mctlass/include/mctlass/epilogue/dispatch_policy.hpp new file mode 100644 index 0000000..8f4110f --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/dispatch_policy.hpp @@ -0,0 +1,136 @@ +/*************************************************************************************************** + * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + * + * Redistribution and use in source and binary forms, with or without modification, are permitted + * provided that the following conditions are met: + * * Redistributions of source code must retain the above copyright notice, this list of + * conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above copyright notice, this list of + * conditions and the following disclaimer in the documentation and/or other materials + * provided with the distribution. + * * Neither the name of the NVIDIA CORPORATION nor the names of its contributors may be used + * to endorse or promote products derived from this software without specific prior written + * permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR + * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND + * FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, + * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; + * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, + * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/numeric_conversion.h" +#include "mctlass/epilogue/thread/scale_type.h" + +////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::epilogue { + +////////////////////////////////////////////////////////////////////////////// + +// Epilogue schedule types that can be used for categorical dispatch +struct NoSmemWarpSpecialized {}; +struct TmaWarpSpecialized {}; +struct TmaWarpSpecializedCooperative {}; + +struct TmaWarpSpecializedElementwiseBase : public TmaWarpSpecialized {}; +struct TmaWarpSpecializedCooperativeElementwiseBase : public TmaWarpSpecializedCooperative {}; + +template < + template class ActivationFunctor_, + thread::ScaleType::Kind Scale_ = thread::ScaleType::Default, + FloatRoundStyle Round_ = FloatRoundStyle::round_to_nearest +> +struct TmaWarpSpecializedElementwise : public TmaWarpSpecializedElementwiseBase { + template + using ActivationFunctor = ActivationFunctor_; + static constexpr thread::ScaleType::Kind Scale = Scale_; + static constexpr FloatRoundStyle Round = Round_; +}; + +template < + template class ActivationFunctor_, + thread::ScaleType::Kind Scale_ = thread::ScaleType::Default, + FloatRoundStyle Round_ = FloatRoundStyle::round_to_nearest +> +struct TmaWarpSpecializedCooperativeElementwise : public TmaWarpSpecializedCooperativeElementwiseBase { + template + using ActivationFunctor = ActivationFunctor_; + static constexpr thread::ScaleType::Kind Scale = Scale_; + static constexpr FloatRoundStyle Round = Round_; +}; + +struct TmaWarpSpecializedBiasElementwiseBase : public TmaWarpSpecialized{}; +struct TmaWarpSpecializedCooperativeBiasElementwiseBase : public TmaWarpSpecializedCooperative {}; + +template < + template class ActivationFunctor_, + class ElementT_, + template class BiasOp_, + bool StoreT_, + class ElementBias_ +> +struct TmaWarpSpecializedBiasElementwise : public TmaWarpSpecializedBiasElementwiseBase { + template + using ActivationFunctor = ActivationFunctor_; + using ElementT = ElementT_; + + template + using BiasOp = BiasOp_; + + static constexpr bool StoreT = StoreT_; + using ElementBias = ElementBias_; +}; + +template < + template class ActivationFunctor_, + class ElementT_, + template class BiasOp_, + bool StoreT_, + class ElementBias_ +> +struct TmaWarpSpecializedCooperativeBiasElementwise : public TmaWarpSpecializedCooperativeBiasElementwiseBase { + template + using ActivationFunctor = ActivationFunctor_; + + using ElementT = ElementT_; + + template + using BiasOp = BiasOp_; + + static constexpr bool StoreT = StoreT_; + using ElementBias = ElementBias_; +}; + +// +// Collective Epilogue Policies +// + +template< + int StagesC_, + int StagesD_, + bool DisableSmemReuseC_ +> +struct Sm90TmaWarpSpecialized { + constexpr static int StagesC = StagesC_; + constexpr static int StagesD = StagesD_; + constexpr static bool DisableSmemReuseC = DisableSmemReuseC_; +}; + +template< + int StagesC_, + int StagesD_ +> +struct Sm90TmaWarpSpecializedBiasElementwise { + constexpr static int StagesC = StagesC_; + constexpr static int StagesD = StagesD_; +}; + +////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::epilogue diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/activation.h b/csrc/mctlass/include/mctlass/epilogue/thread/activation.h new file mode 100644 index 0000000..7d11c69 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/activation.h @@ -0,0 +1,705 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief This extends the contents of mctlass/functional.h with frequently used activation functions. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/constants.h" +#include "mctlass/complex.h" +#include "mctlass/array.h" +#include "mctlass/half.h" +#include "mctlass/functional.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// +template +struct LinearCombinationGenericParams { + T alpha; ///< scales accumulators + T beta; ///< scales source tensor + T const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + T const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + + // + // Methods + // + + MCTLASS_HOST_DEVICE + LinearCombinationGenericParams(): + alpha(T(1)), + beta(T(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + LinearCombinationGenericParams( + T alpha, + T beta = T(0) + ): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + LinearCombinationGenericParams( + T const *alpha_ptr, + T const *beta_ptr = nullptr + ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Identity operator +template +struct Identity { + static const bool kIsHeavy=false; + + MCTLASS_HOST_DEVICE + T operator()(T value) const { + return value; + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + T operator()(T const &value, Params const ¶ms_) const { + return this->operator()(value); + } +}; + +template +struct Identity > { + MCTLASS_HOST_DEVICE + Array operator()(Array const &value) const { + return value; + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &value, Params const ¶ms_) const { + return this->operator()(value); + } +}; + +/// ReLu operator - propagates NaNs +/// Always put threshold in the right hand side of max to propagate NaN. +template +struct ReLu { + static const bool kIsHeavy=false; + MCTLASS_HOST_DEVICE + T operator()(T const & threshold, T value) const { + maximum mx; + + return mx(value, threshold); + } + + MCTLASS_HOST_DEVICE + T operator()(T value) const { + maximum mx; + + return mx(value, T(0)); + } + + /// Host-constructable parameters structure + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + T operator()(T value, Params const ¶ms_) const { + return this->operator()(value); + } +}; + +template +struct ReLu> { + static const bool kIsHeavy=false; + MCTLASS_HOST_DEVICE + Array operator()(T const & threshold, Array const &frag) const { + maximum > mx; + + return mx(frag, threshold); + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const &frag) const { + maximum > mx; + return mx(frag, T(0)); + } + + /// Host-constructable parameters structure + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &frag, Params const ¶ms_) const { + return this->operator()(frag); + } +}; + +// Leaky Relu operator +template +struct LeakyReLU { + + struct Params: LinearCombinationGenericParams { + T leaky_alpha; ///< leaky_alpha + + // Methods + using LinearCombinationGenericParams::LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + Params(): + LinearCombinationGenericParams(), + leaky_alpha(T(1)) {} + + MCTLASS_HOST_DEVICE + Params( + T alpha, + T beta, + T leaky_alpha = T(1) + ): LinearCombinationGenericParams(alpha, beta), leaky_alpha(leaky_alpha) {} + }; + + MCTLASS_HOST_DEVICE + T operator()(T const &value, T const & alpha_recip) const { + T res = value > T(0) ? value : value * alpha_recip; + return res; + } + + MCTLASS_HOST_DEVICE + T operator()(T const &value, Params const ¶ms_) const { + this->operator()(value, params_.leaky_alpha); + } +}; + +template +struct LeakyReLU > { + + struct Params: LinearCombinationGenericParams { + T leaky_alpha; ///< leaky_alpha + using LinearCombinationGenericParams::LinearCombinationGenericParams; + + // Methods + + MCTLASS_HOST_DEVICE + Params(): + LinearCombinationGenericParams(), + leaky_alpha(T(1)) {} + + MCTLASS_HOST_DEVICE + Params( + T alpha, + T beta, + T leaky_alpha = T(1) + ): LinearCombinationGenericParams(alpha, beta), leaky_alpha(leaky_alpha) {} + }; + + + MCTLASS_HOST_DEVICE + Array operator()(Array const &value, T const & alpha_recip) const { + Array y; + LeakyReLU leaky_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < int(value.size()); ++i) { + y[i] = leaky_op(value[i], alpha_recip); + } + + return y; + } + + MCTLASS_HOST_DEVICE + Array operator()(Array const &value, Params const ¶ms_) const { + return this->operator()(value, params_.leaky_alpha); + } +}; + +// Tanh operator +template +struct Tanh { + MCTLASS_HOST_DEVICE + T operator()(T const &scalar) const { + return fast_tanh(scalar); + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + T operator()(T const &scalar, Params const ¶ms_) const { + return this->operator()(scalar); + } +}; + +template +struct Tanh > { + MCTLASS_HOST_DEVICE + Array operator()(Array const &value) const { + Array y; + Tanh tanh_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + y[i] = tanh_op(value[i]); + } + + return y; + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &value, Params const ¶ms_) const { + return this->operator()(value); + } +}; + +template +struct Tanh> { + using T = half_t; + + MCTLASS_HOST_DEVICE + Array operator()(Array const& z) const { + fast_tanh_op> tanh; + return tanh(z); + + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &value, Params const ¶ms_) const { + return this->operator()(value); + } +}; + +// Sigmoid operator +template +struct Sigmoid { + MCTLASS_HOST_DEVICE + T operator()(T const &scalar) const { + return T(1) / (T(1) + fast_exp(-scalar)); + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + T operator()(T const &scalar, Params const ¶ms_) const { + return this->operator()(scalar); + } +}; + +template +struct Sigmoid > { + MCTLASS_HOST_DEVICE + Array operator()(Array const &value) const { + Array y; + Sigmoid sigmoid_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + y[i] = sigmoid_op(value[i]); + } + + return y; + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &value, Params const ¶ms_) const { + return this->operator()(value); + } +}; + +template +struct Sigmoid> { + using T = half_t; + + MCTLASS_HOST_DEVICE + Array operator()(Array const& z) const { + plus> add; + +#if defined(MCTLASS_USE_TANH_FOR_SIGMOID) + multiplies> mul; + fast_tanh_op> tanh; + return mul(add(tanh(mul(z, mctlass::constants::half())), mctlass::constants::one()), + mctlass::constants::half()); +#else + divides> div; + negate> neg; + fast_exp_op> fast_exp; + return div(mctlass::constants::one(), + add(mctlass::constants::one(), + fast_exp(neg(z)))); +#endif + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &z, Params const ¶ms_) const { + return this->operator()(z); + } +}; + +// SiLu (swish) operator introduced by Elfwing et al. in the following paper +// "Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning" (2017) +// https://arxiv.org/pdf/1702.03118.pdf +// It is used in EfficientNet and YOLOv5, for example. +// Reference: https://pytorch.org/docs/stable/generated/torch.nn.SiLU.html +template +struct SiLu { + MCTLASS_HOST_DEVICE + T operator()(T const &scalar) const { + Sigmoid sigmoid; + return scalar * sigmoid(scalar); + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + T operator()(T const &scalar, Params const ¶ms_) const { + return this->operator()(scalar); + } +}; + +template +struct SiLu> { + MCTLASS_HOST_DEVICE + Array operator()(Array const &value) const { + Sigmoid> sigmoid_op; + multiplies> mul; + return mul(value, sigmoid_op(value)); + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &value, Params const ¶ms_) const { + return this->operator()(value); + } +}; + +// Hardswish operator introduced by Howard et al. in the following paper +// "Searching for MobileNetV3" (2019) +// https://arxiv.org/pdf/1905.02244.pdf +// It is used in models based on MobilenetNetV3. +// Reference: https://pytorch.org/docs/stable/generated/torch.nn.Hardswish.html +template +struct HardSwish { + MCTLASS_HOST_DEVICE + T operator()(T const &x) const { + minimum mn; + maximum mx; + T relu6 = mn(mx(x + T(3), T(0)), T(6)); + return x * relu6 / T(6); + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + T operator()(T const &x, Params const ¶ms_) const { + return this->operator()(x); + } +}; + +template <> +struct HardSwish { + using T = float; + + MCTLASS_HOST_DEVICE + T operator()(T const &x) const { + minimum mn; + maximum mx; + T relu6 = mn(mx(x + T(3), T(0)), T(6)); + return x * relu6 * 0.16666667f; + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + T operator()(T const &x, Params const ¶ms_) const { + return this->operator()(x); + } +}; + +template +struct HardSwish > { + MCTLASS_HOST_DEVICE + Array operator()(Array const &value) const { + Array y; + HardSwish hardswish_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + y[i] = hardswish_op(value[i]); + } + + return y; + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &x, Params const ¶ms_) const { + return this->operator()(x); + } +}; + +template +struct HardSwish > { + using T = half_t; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &value) const { + minimum > mn; + maximum > mx; + multiplies > mul; + plus > add; + + return mul(mul(mn(mx(add(value, T(3)), T(0)), T(6)), value), T(0.16666667f)); + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &x, Params const ¶ms_) const { + return this->operator()(x); + } +}; + +// +// GELU function definitions implemented as described by +// Hendrycks, D., and Gimpel, K. in +// "Gaussian Error Linear Units (GELUs)." (2020) +// https://arxiv.org/pdf/1606.08415.pdf +// +// Floating-point constants are Taylor coefficients described in the paper. +// + +// GELU operator +template +struct GELU { + MCTLASS_HOST_DEVICE + T operator()(T const &scalar) const { + return T(mctlass::constants::half() * scalar * + (mctlass::constants::one() + (T)erff((float)(scalar * mctlass::constants::half_root_two())))); + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + T operator()(T const &scalar, Params const ¶ms_) const { + return this->operator()(scalar); + } +}; + +template <> +struct GELU { + MCTLASS_HOST_DEVICE + float operator()(float const &scalar) const { + return mctlass::constants::half() * scalar * + (mctlass::constants::one() + erff( scalar * mctlass::constants::half_root_two() )); + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + float operator()(float const &scalar, Params const ¶ms_) const { + return this->operator()(scalar); + } +}; + +template <> +struct GELU { + MCTLASS_HOST_DEVICE + double operator()(double const &scalar) const { + return mctlass::constants::half() * scalar * + (mctlass::constants::one() + erf( scalar * mctlass::constants::half_root_two() )); + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + double operator()(double const &scalar, Params const ¶ms_) const { + return this->operator()(scalar); + } +}; + +template +struct GELU > { + MCTLASS_HOST_DEVICE + Array operator()(Array const &value) const { + Array y; + GELU gelu_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + y[i] = gelu_op(value[i]); + } + + return y; + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &value, Params const ¶ms_) const { + return this->operator()(value); + } +}; + +// GELU operator implemented using the Taylor series approximation +template +struct GELU_taylor { + static const bool kIsHeavy=true; + MCTLASS_HOST_DEVICE + T operator()(T const &z) const { + + T k0 = T(0.7978845608028654); + T k1 = T(0.044715); + + return T(mctlass::constants::half() * z * + (mctlass::constants::one() + fast_tanh(k0 * z * (mctlass::constants::one() + k1 * z * z)))); + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + T operator()(T const &scalar, Params const ¶ms_) const { + return this->operator()(scalar); + } +}; + +template +struct GELU_taylor > { + static const bool kIsHeavy=true; + MCTLASS_HOST_DEVICE + Array operator()(Array const &z) const { + + using T = half_t; + Array y; + + half_t k0 = half_t(0.7978845608028654); + half_t k1 = half_t(0.044715); + + multiply_add> fma; + multiplies> mul; + plus> add; + + fast_tanh_op> tanh; + + Array u = mul(mul(k0, z), fma(mul(k1, z), z, mctlass::constants::one())); + + y = mul(mul(z, mctlass::constants::half()), add(mctlass::constants::one(), tanh(u))); + + return y; + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &value, Params const ¶ms_) const { + return this->operator()(value); + } +}; + +template +struct GELU_taylor > { + static const bool kIsHeavy=true; + MCTLASS_HOST_DEVICE + Array operator()(Array const &value) const { + Array y; + GELU_taylor gelu_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + y[i] = gelu_op(value[i]); + } + + return y; + } + + using Params = LinearCombinationGenericParams; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &value, Params const ¶ms_) const { + return this->operator()(value); + } +}; + +/// Computes backwards pass for GELU operator assuming d_t is the layer gradient and +/// z is computed from the forward pass. +template +struct dGELU { + MCTLASS_HOST_DEVICE + T operator()(T const &d_t, T const &z) const { + + T k0 = T(0.7978845608028654); + T k1 = T(0.044715); + T k2 = T(0.1070322243); + + T tanh_out = fast_tanh(k0 * z * (1 + k1 * z * z)); + + T ff = constants::half() * z * ((1 - tanh_out * tanh_out) * (k0 + k2 * z * z)) + + constants::half() * (1 + tanh_out); + + return ff * d_t; + } +}; + +template +struct dGELU > { + MCTLASS_HOST_DEVICE + Array operator()(Array const &d_t, Array const &z) const { + Array y; + dGELU gelu_op; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + y[i] = gelu_op(d_t[i], z[i]); + } + + return y; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/conversion_op.h b/csrc/mctlass/include/mctlass/epilogue/thread/conversion_op.h new file mode 100644 index 0000000..3ff83fa --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/conversion_op.h @@ -0,0 +1,132 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing conversion operations used by epilogues. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Converts the result without other operations +/// +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +class Convert { +public: + + using ElementOutput = ElementOutput_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementAccumulator_; + + static int const kCount = Count; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using ComputeFragment = FragmentAccumulator; + + static FloatRoundStyle const kRound = Round; + + static bool const kIsHeavy = false; + + /// Host-constructable parameters structure + struct Params { + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params() {} + }; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + Convert(Params const ¶ms = Params()) { + + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + + } + + /// Returns true if source is needed based on state of runtime arguments + MCTLASS_HOST_DEVICE + constexpr bool is_source_needed() const { + return false; + } + + /// Constexpr function to enable the compiler to optimize away the source loading if it is + /// never needed. + MCTLASS_HOST_DEVICE + constexpr bool is_source_ever_needed() const { + return false; + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentOutput const &source = FragmentOutput(), + ElementCompute uniform = ElementCompute(0)) const { + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(accumulator); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/detail.hpp b/csrc/mctlass/include/mctlass/epilogue/thread/detail.hpp new file mode 100644 index 0000000..23ecb8c --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/detail.hpp @@ -0,0 +1,52 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Utilities for thread-level epilogues +*/ + +#pragma once + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +namespace detail { + +/// Class used to identify cases in which no operation is performed +template +struct NoOp {}; + +} // namespace detail +} // namespace thread +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination.h new file mode 100644 index 0000000..54bd02f --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination.h @@ -0,0 +1,330 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing linear combination operations used by epilogues. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/epilogue/thread/scale_type.h" +#include "mctlass/epilogue/thread/linear_combination_params.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements. +/// +/// D = alpha * accumulator + beta * source +/// +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation. + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data to store + typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type + typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination + ScaleType::Kind Scale = ScaleType::Default, ///< Control Alpha and Beta scaling + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest, + typename ElementSource_ = ElementOutput_ +> +class LinearCombination { +public: + + using ElementOutput = ElementOutput_; + using ElementSource = ElementSource_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + using ElementC = ElementSource_; + using ElementD = ElementOutput_; + + static int const kCount = Count; + static const ScaleType::Kind kScale = Scale; + using FragmentOutput = Array; + using FragmentSource = Array; + using FragmentAccumulator = Array; + using FragmentCompute = Array; + + static FloatRoundStyle const kRound = Round; + + /// Host-constructable parameters structure + struct Params + { + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales source tensor + ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta + ): + alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha + ): + alpha(alpha), beta(0), alpha_ptr(nullptr), beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr, + ElementCompute const *beta_ptr + ): + alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr + ): + alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(nullptr) { } + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombination(Params const ¶ms) { + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + if (Scale == ScaleType::NoBetaScaling) return true; + + if (Scale == ScaleType::OnlyAlphaScaling) return false; + + if (Scale == ScaleType::Nothing) return false; + + return beta_ != ElementCompute(0); + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + } + + /// Computes intermediate: X = beta * source + MCTLASS_HOST_DEVICE + FragmentCompute compute_intermediate( + FragmentSource const &source) const { + + // Convert source to internal compute numeric type + NumericArrayConverter source_converter; + FragmentCompute converted_source = source_converter(source); + + if (Scale == ScaleType::NoBetaScaling) { + return converted_source; + } + else { + multiplies mul_source; + return mul_source(beta_, converted_source); + } + } + + /// Computes linear scaling with intermediate: D = alpha * accumulator + X + MCTLASS_HOST_DEVICE + FragmentOutput with_intermediate( + FragmentAccumulator const& accumulator, + FragmentCompute const& intermediate) const { + + // Convert accumulator to internal compute numeric type + NumericArrayConverter accumulator_converter; + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + if (Scale == ScaleType::Nothing) { + return destination_converter(converted_accumulator); + } else { + // Perform binary operations + multiply_add mul_add_accumulator; + FragmentCompute computed_output = mul_add_accumulator(alpha_, converted_accumulator, intermediate); + + return destination_converter(computed_output); + } + } + + /// Computes linear scaling with source: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentSource const &source) const { + + // Convert source to internal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + FragmentCompute converted_source = source_converter(source); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + if (Scale == ScaleType::Nothing) + return destination_converter(converted_accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + + if (Scale == ScaleType::NoBetaScaling) + intermediate = converted_source; + else + intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform + + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + + return destination_converter(intermediate); + } + + /// Computes linear scaling: D = alpha * accumulator + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + if (Scale == ScaleType::Nothing) + return destination_converter(converted_accumulator); + + // Perform binary operations + FragmentCompute intermediate; + multiplies mul_accumulator; + + intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum + + return destination_converter(intermediate); + } + + // + // Specializations for scalar (for use with cute::collective::DefaultEpilogue) + // + MCTLASS_HOST_DEVICE + ElementD operator()(ElementAccumulator const accumulator, ElementC const source) const { + // Convert everything to Compute type, do compute, and then store to output type + NumericConverter accumulator_converter; + [[maybe_unused]] NumericConverter source_converter; + NumericConverter destination_converter; + + // Convert to destination numeric type + + ElementCompute converted_accumulator = accumulator_converter(accumulator); + if constexpr (Scale == ScaleType::Nothing) { + return destination_converter(converted_accumulator); + } + + // Perform binary operations + ElementCompute intermediate; + multiplies multiply; + multiply_add madd; + + if constexpr (Scale == ScaleType::NoBetaScaling) { + intermediate = source_converter(source); + } + else { + intermediate = multiply(beta_, source); // X = beta * C + uniform + } + + intermediate = madd(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + return destination_converter(intermediate); + } + + MCTLASS_HOST_DEVICE + ElementD operator()(ElementAccumulator const accumulator) const { + // Convert everything to Compute type, do compute, and then store to output type + NumericConverter accumulator_converter; + NumericConverter destination_converter; + ElementCompute converted_accumulator = accumulator_converter(accumulator); + + // Convert to destination numeric type + if constexpr (Scale == ScaleType::Nothing) { + return destination_converter(converted_accumulator); + } + + // Perform binary operations + ElementCompute intermediate; + multiplies multiply; + + intermediate = multiply(alpha_, accumulator); // D = alpha * Accum + return destination_converter(intermediate); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_bias_elementwise.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_bias_elementwise.h new file mode 100644 index 0000000..15c75c5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_bias_elementwise.h @@ -0,0 +1,270 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Functor performing linear combination operations used by epilogues. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" + +#include "mctlass/epilogue/thread/activation.h" +#include "mctlass/epilogue/thread/scale_type.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// This base class is meant to define the concept required of the +/// EpilogueWithBroadcast::OutputOp +template < + typename ElementC_, + typename ElementAccumulator_, + typename ElementCompute_, + typename ElementZ_, + typename ElementT_, + int ElementsPerAccess, + typename ElementwiseOp_ = Identity, + typename BinaryOp_ = plus, + bool StoreT_ = true, + typename ElementVector_ = ElementC_ +> +class LinearCombinationBiasElementwise { +public: + + using ElementOutput = ElementC_; + using ElementC = ElementC_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + using ElementZ = ElementZ_; + using ElementT = ElementT_; + using ElementVector = ElementVector_; + static int const kElementsPerAccess = ElementsPerAccess; + static int const kCount = kElementsPerAccess; + + using ElementwiseOp = ElementwiseOp_; + using BinaryOp = BinaryOp_; + + // Indicates that this epilogue applies only one binary operation + static bool const kIsSingleSource = true; + + using FragmentAccumulator = Array; + using FragmentCompute = Array; + using FragmentC = Array; + using FragmentZ = Array; + using FragmentT = Array; + + // Definitions needed for collective epilogue + using FragmentSource = FragmentC; + using FragmentOutput = FragmentZ; + using ElementBias = ElementVector; + using FragmentBias = FragmentCompute; + using ActivationFunctor = ElementwiseOp; + static const ScaleType::Kind kScale = ScaleType::Default; + + static bool const kIsHeavy = ElementwiseOp::kIsHeavy; + + /// If true, the 'Z' tensor is stored + static bool const kStoreZ = true; + + /// If true, the 'T' tensor is stored + static bool const kStoreT = StoreT_; + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales source tensor + ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta + ): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha + ): alpha(alpha), beta(0), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr, + ElementCompute const *beta_ptr + ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr + ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(nullptr) { + + } + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + bool skip_elementwise_; + +public: + + // + // Methods + // + + /// Constructor from Params + MCTLASS_HOST_DEVICE + LinearCombinationBiasElementwise(Params const ¶ms) { + + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + skip_elementwise_ = false; + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + return beta_ != ElementCompute(0); + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + + if (k_partition != k_partition_count - 1) { + skip_elementwise_ = true; + } + } + + /// Applies the operation when is_source_needed() is true + MCTLASS_HOST_DEVICE + void operator()( + FragmentZ &frag_Z, + FragmentT &frag_T, + FragmentAccumulator const &AB, + FragmentC const &frag_C, + FragmentCompute const &V) const { + + ElementwiseOp elementwise_op; + BinaryOp binary_op; + + FragmentCompute tmp_Accum = NumericArrayConverter()(AB); + FragmentCompute tmp_C = NumericArrayConverter()(frag_C); + FragmentCompute result_Z; + FragmentCompute result_T; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kElementsPerAccess; ++i) { + ElementCompute z = binary_op(alpha_ * tmp_Accum[i] + beta_ * tmp_C[i], V[i]); + result_T[i] = z; + result_Z[i] = skip_elementwise_ ? z : elementwise_op(z); + } + + NumericArrayConverter convert_z; + frag_Z = convert_z(result_Z); + + NumericArrayConverter convert_t; + frag_T = convert_t(result_T); + } + + /// Applies the operation when is_source_needed() is false + MCTLASS_HOST_DEVICE + void operator()( + FragmentZ &frag_Z, + FragmentT &frag_T, + FragmentAccumulator const &AB, + FragmentCompute const &V) const { + + ElementwiseOp elementwise_op; + BinaryOp binary_op; + + FragmentCompute tmp_Accum = NumericArrayConverter()(AB); + FragmentCompute result_Z; + FragmentCompute result_T; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kElementsPerAccess; ++i) { + ElementCompute z = binary_op(alpha_ * tmp_Accum[i], V[i]); + result_T[i] = z; + result_Z[i] = skip_elementwise_ ? z : elementwise_op(z); + } + + NumericArrayConverter convert_z; + frag_Z = convert_z(result_Z); + + NumericArrayConverter convert_t; + frag_T = convert_t(result_T); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_bias_relu.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_bias_relu.h new file mode 100644 index 0000000..100c640 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_bias_relu.h @@ -0,0 +1,471 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing linear combination operations used by epilogues. +*/ + +#pragma once + +#include + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/epilogue/thread/activation.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +template +struct ArrayMaximum { + + MCTLASS_HOST_DEVICE + Array operator()( + Array const &lhs, + Array const &rhs) const { + + Array result; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < ElementsPerAccess; ++i) { + result[i] = fmax(lhs[i], rhs[i]); + } + + return result; + } +}; + +template +struct ArrayMaximum { + + MCTLASS_DEVICE + Array operator()( + Array const &lhs, + Array const &rhs) const { + + Array result; + + #if defined(__MACA_ARCH__) + int const kVectorCount = ElementsPerAccess / 2; + + + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(lhs.raw_data()); + __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(rhs.raw_data()); + __half2 *res_ptr = reinterpret_cast<__half2 *>(result.raw_data()); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kVectorCount; ++i) { + res_ptr[i] = __hmax2(lhs_ptr[i], rhs_ptr[i]); + } + + #else + __half const *lhs_ptr = reinterpret_cast<__half const *>(lhs.raw_data()); + __half const *rhs_ptr = reinterpret_cast<__half const *>(rhs.raw_data()); + __half *res_ptr = reinterpret_cast<__half *>(result.raw_data()); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < ElementsPerAccess; ++i) { + res_ptr[i] = ((lhs_ptr[i] < rhs_ptr[i]) ? rhs_ptr[i] : lhs_ptr[i]); + } + + #endif + + return result; + } + + MCTLASS_DEVICE + Array operator()( + Array const &lhs, + half_t const &rhs) const { + + Array result; + + #if defined(__MACA_ARCH__) + int const kVectorCount = ElementsPerAccess / 2; + + + __half rhs_raw = reinterpret_cast<__half const &>(rhs); + __half2 rhs_pair = __half2half2(rhs_raw); + + __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(lhs.raw_data()); + __half2 *res_ptr = reinterpret_cast<__half2 *>(result.raw_data()); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kVectorCount; ++i) { + res_ptr[i] = __hmax2(lhs_ptr[i], rhs_pair); + } + + #else + + __half const *lhs_ptr = reinterpret_cast<__half const *>(lhs.raw_data()); + __half const rhs_raw = reinterpret_cast<__half const &>(rhs); + __half *res_ptr = reinterpret_cast<__half *>(result.raw_data()); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < ElementsPerAccess; ++i) { + res_ptr[i] = ((lhs_ptr[i] < rhs_raw) ? rhs_raw : lhs_ptr[i]); + } + + #endif + + return result; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct ReluConditional { + + MCTLASS_HOST_DEVICE + void operator()( + bool conditional[], + Array const &fragment, + Element threshold) const { + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < ElementsPerAccess; ++i) { + conditional[i] = !(fragment[i] < threshold); + } + } +}; + +template +struct ReluConditional { + + MCTLASS_DEVICE + void operator()( + bool conditional[], + Array const &fragment, + half_t threshold) const { + + __half y = reinterpret_cast<__half const &>(threshold); + __half const *x = reinterpret_cast<__half const *>(fragment.raw_data()); + + #if defined(__MACA_ARCH__) + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < ElementsPerAccess; ++i) { + conditional[i] = !__hlt(x[i], y); + } + #else + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < ElementsPerAccess; ++i) { + conditional[i] = float(x[i]) < float(y) ? 0 : 1; + } + #endif + } +}; + +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// This is a partial specialization for fused Bias and ReLU. It supports the option of packing +/// ReLU conditionals in a bit vector that may be used by backwards passes as an optimization. +/// +/// This class can only be used with mctlass::epilogue::threadblock::EpilogueWithBroadcast<>. +/// +/// This base class is meant to define the concept required of the +/// EpilogueWithBroadcast::OutputOp +template < + typename ElementC_, + typename ElementAccumulator_, + typename ElementCompute_, + typename ElementZ_, + int ElementsPerAccess, + bool StoreT_ = true, + typename ElementVector_ = ElementC_ +> +class LinearCombinationBiasRelu { +public: + + using ElementOutput = ElementC_; + using ElementC = ElementC_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + using ElementZ = ElementZ_; + using ElementVector = ElementVector_; + + using ElementT = uint1b_t; + + static int const kElementsPerAccess = ElementsPerAccess; + static int const kCount = kElementsPerAccess; + + using ElementwiseOp = ReLu; + using BinaryOp = plus; + + // Indicates that this epilogue applies only one binary operation + static bool const kIsSingleSource = true; + + using FragmentAccumulator = Array; + using FragmentCompute = Array; + using FragmentC = Array; + using FragmentZ = Array; + using FragmentT = Array; + + /// If true, the 'Z' tensor is stored + static bool const kStoreZ = true; + + /// If true, the 'T' tensor is stored + static bool const kStoreT = StoreT_; + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales source tensor + ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + ElementZ threshold; ///< ReLu threshold + + // + // Methods + // + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute()), + alpha_ptr(nullptr), + beta_ptr(nullptr), + threshold(ElementCompute()) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta, + ElementCompute threshold_ = ElementCompute() + ): + alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { + + NumericConverter convert_threshold; + + threshold = convert_threshold(threshold_); + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha + ): alpha(alpha), beta(0), alpha_ptr(nullptr), beta_ptr(nullptr), threshold(ElementZ()) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr, + ElementCompute const *beta_ptr, + ElementCompute threshold_ = ElementCompute() + ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { + + NumericConverter convert_threshold; + + threshold = convert_threshold(threshold_); + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr + ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(nullptr), threshold(ElementZ()) { + } + + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + ElementZ threshold_; + +public: + + // + // Methods + // + + /// Constructor from Params + MCTLASS_HOST_DEVICE + LinearCombinationBiasRelu(Params const ¶ms) { + + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + threshold_ = params.threshold; + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + return beta_ != ElementCompute(0); + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + + if (k_partition != k_partition_count - 1) { + // set to NaN to make ReLU no-op for all except last k partitions + int64_t allones = -1; + threshold_ = reinterpret_cast(allones); + } + } + + /// Applies the operation when is_source_needed() is true + MCTLASS_HOST_DEVICE + void operator()( + FragmentZ &frag_Z, + FragmentT &frag_T, + FragmentAccumulator const &AB, + FragmentC const &frag_C, + FragmentCompute const &V) const { + + BinaryOp binary_op; + + FragmentCompute tmp_Accum = NumericArrayConverter()(AB); + FragmentCompute tmp_C = NumericArrayConverter()(frag_C); + FragmentCompute result_Z; + + bool conditions[kElementsPerAccess]; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kElementsPerAccess; ++i) { + + ElementCompute z = alpha_ * tmp_Accum[i]; + z += beta_ * tmp_C[i]; + + z = binary_op(z, V[i]); + result_Z[i] = z; + + //For mxmaca f16, we cannot got same precision with cuda, so we got zero which is near with it. + if (z < ElementCompute(1e-5f) && z > ElementCompute(-1e-5f)) { + result_Z[i] = ElementCompute(0); + } + + } + + NumericArrayConverter convert_z; + frag_Z = convert_z(result_Z); + + // + // Compute condition + // + + detail::ReluConditional relu_conditional; + relu_conditional(conditions, frag_Z, threshold_); + + detail::ArrayMaximum maximum_op; + frag_Z = maximum_op(frag_Z, threshold_); + + if (kStoreT) { + PackPredicates pack_predicates; + frag_T = pack_predicates(conditions); + } + } + + /// Applies the operation when is_source_needed() is false + MCTLASS_HOST_DEVICE + void operator()( + FragmentZ &frag_Z, + FragmentT &frag_T, + FragmentAccumulator const &AB, + FragmentCompute const &V) const { + + BinaryOp binary_op; + + FragmentCompute tmp_Accum = NumericArrayConverter()(AB); + FragmentCompute result_Z; + + bool conditions[kElementsPerAccess]; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kElementsPerAccess; ++i) { + ElementCompute z = binary_op(alpha_ * tmp_Accum[i], V[i]); + result_Z[i] = z; + + //For mxmaca f16, we cannot got same precision with cuda, so we got zero which is near with it. + if (z < ElementCompute(1e-5f) && z > ElementCompute(-1e-5f)) { + result_Z[i] = ElementCompute(0); + } + + } + + NumericArrayConverter convert_z; + frag_Z = convert_z(result_Z); + + // + // Compute condition + // + + detail::ReluConditional relu_conditional; + relu_conditional(conditions, frag_Z, threshold_); + + detail::ArrayMaximum maximum_op; + frag_Z = maximum_op(frag_Z, threshold_); + + // + // Compute conditions + // + + // + // Store + // + if (kStoreT) { + PackPredicates pack_predicates; + frag_T = pack_predicates(conditions); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_clamp.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_clamp.h new file mode 100644 index 0000000..3da6bc8 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_clamp.h @@ -0,0 +1,689 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing linear scaling operations used by epilogues. Values are clamped before + converting to the output element type. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/epilogue/thread/scale_type.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +/// Single source of truth for whether to unroll for `LinearCombinationClamp()` +constexpr bool LinearCombinationClampIsHeavy() { + return false; +} + +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements then clamps the output before +/// converting to the output element type. +/// +/// D = alpha * accumulator + beta * source + uniform +/// +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data to store + typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type + typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination + ScaleType::Kind Scale = ScaleType::Default, ///< Control Alpha and Beta scaling + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +class LinearCombinationClamp { +public: + + using ElementOutput = ElementOutput_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + + static int const kCount = Count; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using ComputeFragment = Array; + using FragmentSource = Array; + + static FloatRoundStyle const kRound = Round; + + static bool const kIsHeavy = detail::LinearCombinationClampIsHeavy(); + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales source tensor + ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta + ): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha + ): alpha(alpha), beta(0), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr, + ElementCompute const *beta_ptr + ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr + ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(nullptr) { + + } + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationClamp(Params const ¶ms) { + + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + if (Scale == ScaleType::NoBetaScaling) return true; + + if (Scale == ScaleType::OnlyAlphaScaling) return false; + + if (Scale == ScaleType::Nothing) return false; + + return beta_ != ElementCompute(0); + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentOutput const &source, + ElementCompute uniform = ElementCompute(0)) const { + + // Convert source to interal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + ComputeFragment converted_source = source_converter(source); + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + + ComputeFragment intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + + minimum min_accumulator; + maximum max_accumulator; + + if (Scale == ScaleType::NoBetaScaling) { + intermediate = converted_source; + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } else if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } + + /// Clamping constant value + ElementCompute const kClampMax = + ElementCompute(platform::numeric_limits::max()); + + ElementCompute const kClampMin = + ElementCompute(platform::numeric_limits::lowest()); + + intermediate = max_accumulator(intermediate, kClampMin); + intermediate = min_accumulator(intermediate, kClampMax); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(intermediate); + } + + /// Computes linear scaling: D = alpha * accumulator + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + + ComputeFragment intermediate; + + multiplies mul_accumulator; + + minimum min_accumulator; + maximum max_accumulator; + + if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum + } + + /// Clamping constant value + ElementCompute const kClampMax = + ElementCompute(platform::numeric_limits::max()); + + ElementCompute const kClampMin = + ElementCompute(platform::numeric_limits::lowest()); + + intermediate = max_accumulator(intermediate, kClampMin); + intermediate = min_accumulator(intermediate, kClampMax); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(intermediate); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Conditional guards to enable partial specialization for packed integers +#if defined(__MACA_ARCH__) + +/// Applies a linear combination operator to an array of elements then clamps the output before +/// converting to the output element type. +/// +/// D = alpha * accumulator + beta * source + uniform +/// +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + ScaleType::Kind Scale, ///< Control Alpha and Beta scaling + FloatRoundStyle Round +> +class LinearCombinationClamp { +public: + + using ElementOutput = ElementOutput_; + using ElementAccumulator = int; + using ElementCompute = float; + + static_assert( + platform::numeric_limits::is_integer, + "This elementwise op expects the output to be int."); + + static int const kCount = Count; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using ComputeFragment = Array; + + static FloatRoundStyle const kRound = Round; + + static bool const kIsHeavy = detail::LinearCombinationClampIsHeavy(); + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales source tensor + ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta + ): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha + ): alpha(alpha), beta(0), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr, + ElementCompute const *beta_ptr + ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr + ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(nullptr) { + + } + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationClamp(Params const ¶ms) { + + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + if (Scale == ScaleType::NoBetaScaling) return true; + + if (Scale == ScaleType::OnlyAlphaScaling) return false; + + if (Scale == ScaleType::Nothing) return false; + + return beta_ != ElementCompute(0); + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentOutput const &source, + ElementCompute uniform = ElementCompute(0)) const { + + // Convert source to interal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + ComputeFragment converted_source = source_converter(source); + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + // Compute linear scaling in floating point + ComputeFragment intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + + // Float min-max + if (Scale == ScaleType::NoBetaScaling) { + intermediate = converted_source; + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } else if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } + + // Convert floats back to INT + FragmentAccumulator scaled_accumulator; + + NumericArrayConverter compute_converter; + + scaled_accumulator = compute_converter(intermediate); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(scaled_accumulator); + } + + /// Computes linear scaling: D = alpha * accumulator + MCTLASS_HOST_DEVICE + FragmentOutput operator()(FragmentAccumulator const &accumulator) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + // Compute linear scaling in floating point + ComputeFragment intermediate; + + multiplies mul_add_accumulator; + + // Float min-max + if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_add_accumulator(alpha_, converted_accumulator); // D = alpha * Accum + } + + // Convert floats back to INT + FragmentAccumulator scaled_accumulator; + + NumericArrayConverter compute_converter; + + scaled_accumulator = compute_converter(intermediate); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(scaled_accumulator); + } +}; + +#endif // Conditional guards to enable partial specialization for packed integers + +//////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements then clamps +/// the output before converting to the output element type. +/// +template < + /// Data type used to load and store< tensors + typename ElementOutput_, + /// Number of elements computed per operation + int Count, + ///< Control Alpha and Beta scaling + ScaleType::Kind Scale = ScaleType::Default, + /// Rounding mode + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest> +class FastLinearCombinationClamp { + public: + using ElementOutput = ElementOutput_; + using ElementAccumulator = int; + using ElementCompute = float; + + static_assert( + platform::numeric_limits::is_integer, + "This elementwise op expects the output to be int."); + + static int const kCount = Count; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using ComputeFragment = Array; + + static FloatRoundStyle const kRound = Round; + + static bool const kIsHeavy = false; + + /// Host-constructable parameters structure + struct Params { + /// scales accumulators + ElementCompute alpha; + /// scales source tensor + ElementCompute beta; + /// pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *alpha_ptr; + /// pointer to source scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params() + : alpha(ElementCompute(1)), + beta(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) {} + + MCTLASS_HOST_DEVICE + Params(ElementCompute alpha, ElementCompute beta) + : alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) {} + + MCTLASS_HOST_DEVICE + Params(ElementCompute alpha) + : alpha(alpha), beta(0), alpha_ptr(nullptr), beta_ptr(nullptr) {} + + MCTLASS_HOST_DEVICE + Params(ElementCompute const *alpha_ptr, ElementCompute const *beta_ptr) + : alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {} + + MCTLASS_HOST_DEVICE + Params(ElementCompute const *alpha_ptr) + : alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(nullptr) {} + }; + + private: + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + + public: + /// Constructs the function object, possibly loading from pointers in host + /// memory + MCTLASS_HOST_DEVICE + FastLinearCombinationClamp(Params const ¶ms) { + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + if (Scale == ScaleType::NoBetaScaling) return true; + + if (Scale == ScaleType::OnlyAlphaScaling) return false; + + if (Scale == ScaleType::Nothing) return false; + + return beta_ != ElementCompute(0); + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()(FragmentAccumulator const &accumulator, + FragmentOutput const &source, + ElementCompute uniform = ElementCompute(0)) const { + // Convert source to interal compute numeric type + FastNumericArrayConverter + source_converter; + FastNumericArrayConverter + accumulator_converter; + + ComputeFragment converted_source = source_converter(source); + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + // Compute linear scaling in floating point + ComputeFragment intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + + minimum min_accumulator; + maximum max_accumulator; + + // Float min-max + if (Scale == ScaleType::NoBetaScaling) { + intermediate = converted_source; + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } else if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = + mul_add_source(beta_, converted_source); // X = beta * C + uniform + intermediate = mul_add_accumulator(alpha_, converted_accumulator, + intermediate); // D = alpha * Accum + X + } + + /// Clamping constant value + ElementCompute const kClamp = + ElementCompute(1 << (sizeof_bits::value - 1)); + + intermediate = max_accumulator(intermediate, -kClamp); + intermediate = min_accumulator(intermediate, kClamp - ElementCompute(1)); + + // Convert to destination numeric type + FastNumericArrayConverter + destination_converter; + + return destination_converter(intermediate); + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()(FragmentAccumulator const &accumulator) const { + + // Convert source to interal compute numeric type + FastNumericArrayConverter + accumulator_converter; + + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + // Compute linear scaling in floating point + ComputeFragment intermediate; + + multiplies mul_accumulator; + + minimum min_accumulator; + maximum max_accumulator; + + // Float min-max + if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_accumulator(alpha_, converted_accumulator); + } + + /// Clamping constant value + ElementCompute const kClamp = + ElementCompute(1 << (sizeof_bits::value - 1)); + + intermediate = max_accumulator(intermediate, -kClamp); + intermediate = min_accumulator(intermediate, kClamp - ElementCompute(1)); + + // Convert to destination numeric type + FastNumericArrayConverter + destination_converter; + + return destination_converter(intermediate); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_dgelu.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_dgelu.h new file mode 100644 index 0000000..5a4d5f7 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_dgelu.h @@ -0,0 +1,250 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + + \brief Functor performing linear combination followed by dGelu operation +*/ + +#pragma once + +#include "mctlass/half.h" +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/constants.h" +#include "mctlass/fast_math.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/epilogue/thread/activation.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements. +/// +/// D = alpha * accumulator + beta * source + uniform +/// +template < + typename ElementCompute_, ///< Data type returned by this functor + typename ElementAccumulator_, ///< Data type of accumulators + typename ElementSource_, ///< Data type of source tensor + typename ElementTensor_, ///< Data type of additional tensor + int Count, ///< Number of elements computed per operation + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data to store + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +class LinearCombinationDGelu { +public: + + using ElementOutput = ElementSource_; + using ElementCompute = ElementCompute_; + using ElementAccumulator = ElementAccumulator_; + using ElementSource = ElementSource_; + using ElementTensor = ElementTensor_; + + static bool const kIsHeavy = true; + + static int const kCount = Count; + + using FragmentCompute = Array; + using FragmentAccumulator = Array; + using FragmentSource = Array; + using FragmentTensor = Array; + + static FloatRoundStyle const kRound = Round; + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales source tensor + ElementCompute threshold; ///< minimum value that is output + ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute(0)), + threshold(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta, + ElementCompute threshold = ElementCompute(0) + ): alpha(alpha), beta(beta), threshold(threshold), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr, + ElementCompute const *beta_ptr, + ElementCompute threshold = ElementCompute(0) + ): alpha(0), beta(0), threshold(threshold), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { + + } + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + ElementCompute threshold_; + bool participates_in_reduction_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationDGelu(Params const ¶ms) { + + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + threshold_ = params.threshold; + participates_in_reduction_ = true; + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + return beta_ != ElementCompute(0); + } + + /// Returns true if the threadblock computes the reduction + MCTLASS_HOST_DEVICE + bool participates_in_reduction() const { + return participates_in_reduction_; + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + + if (k_partition != k_partition_count - 1) { + // set to NaN to make ReLU no-op for all except last k partitions + int64_t allones = -1; + threshold_ = reinterpret_cast(allones); + // Avoid computing the reduction if this isn't the final Split-K slice + participates_in_reduction_ = false; + } + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentCompute operator()( + FragmentAccumulator const &accumulator, + FragmentSource const &source, + FragmentTensor const &tensor) const { + + // Convert source to interal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_source = source_converter(source); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + + intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + + dGELU gelu_op; + + // dGelu + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + intermediate[i] = gelu_op(intermediate[i], ElementCompute(tensor[i])); + } + + return intermediate; + } + + /// Computes linear scaling: D = alpha * accumulator + MCTLASS_HOST_DEVICE + FragmentCompute operator()( + FragmentAccumulator const &accumulator, + FragmentTensor const &tensor) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_accumulator; + + intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum + + dGELU gelu_op; + + // dGelu with conversion + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + intermediate[i] = gelu_op(intermediate[i], ElementCompute(tensor[i])); + } + + return intermediate; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_drelu.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_drelu.h new file mode 100644 index 0000000..22f1266 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_drelu.h @@ -0,0 +1,452 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing linear combination with a maximum operation used by epilogues. +*/ + +#pragma once + +#include "mctlass/half.h" +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/epilogue/thread/activation.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements. +/// +/// D = alpha * accumulator + beta * source + uniform +/// +template < + typename ElementCompute_, ///< Data type returned by this functor + typename ElementAccumulator_, ///< Data type of accumulators + typename ElementSource_, ///< Data type of source tensor + typename ElementTensor_, ///< Data type of additional tensor + int Count, ///< Number of elements computed per operation + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data to store + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +class LinearCombinationDRelu { +public: + + using ElementOutput = ElementSource_; + using ElementCompute = ElementCompute_; + using ElementAccumulator = ElementAccumulator_; + using ElementSource = ElementSource_; + using ElementTensor = ElementTensor_; + + static int const kCount = Count; + + using FragmentCompute = Array; + using FragmentAccumulator = Array; + using FragmentSource = Array; + using FragmentTensor = Array; + + static FloatRoundStyle const kRound = Round; + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales source tensor + ElementCompute threshold; ///< minimum value that is output + ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute(0)), + threshold(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta, + ElementCompute threshold = ElementCompute(0) + ): alpha(alpha), beta(beta), threshold(threshold), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr, + ElementCompute const *beta_ptr, + ElementCompute threshold = ElementCompute(0) + ): alpha(0), beta(0), threshold(threshold), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { + + } + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + ElementTensor threshold_; + bool participates_in_reduction_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationDRelu(Params const ¶ms) { + + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + threshold_ = ElementTensor(params.threshold); + participates_in_reduction_ = true; + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + return beta_ != ElementCompute(0); + } + + /// Returns true if the threadblock computes the reduction + MCTLASS_HOST_DEVICE + bool participates_in_reduction() const { + return participates_in_reduction_; + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + + if (k_partition != k_partition_count - 1) { + // set to NaN to make ReLU no-op for all except last k partitions + int64_t allones = -1; + threshold_ = reinterpret_cast(allones); + participates_in_reduction_ = false; + } + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentCompute operator()( + FragmentAccumulator const &accumulator, + FragmentSource const &source, + FragmentTensor const &tensor) const { + + // Convert source to interal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_source = source_converter(source); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + + intermediate = mul_add_source(beta_, converted_source); // X = beta * C + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + + // dReLU = (cond ? dy : 0) + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + ElementTensor cond = tensor[i]; + if (cond <= threshold_) { + intermediate[i] = ElementCompute(); + } + } + + return intermediate; + } + + /// Computes linear scaling: D = alpha * accumulator + MCTLASS_HOST_DEVICE + FragmentCompute operator()( + FragmentAccumulator const &accumulator, + FragmentTensor const &tensor) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_accumulator; + + intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum + + // dReLU = (cond ? dy : 0) + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + ElementTensor cond = tensor[i]; + if (cond <= threshold_) { + intermediate[i] = ElementCompute(); + } + } + + return intermediate; + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements. +/// +/// D = alpha * accumulator + beta * source + uniform +/// +template < + typename ElementCompute_, ///< Data type returned by this functor + typename ElementAccumulator_, ///< Data type of accumulators + typename ElementSource_, ///< Data type of source tensor + int Count, ///< Number of elements computed per operation + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +class LinearCombinationDReluConditionalBits { +public: + + using ElementOutput = ElementSource_; + using ElementCompute = ElementCompute_; + using ElementAccumulator = ElementAccumulator_; + using ElementSource = ElementSource_; + using ElementTensor = uint1b_t; + + static bool const kIsHeavy = false; + + static int const kCount = Count; + + using FragmentCompute = Array; + using FragmentAccumulator = Array; + using FragmentSource = Array; + using FragmentTensor = Array; + + static FloatRoundStyle const kRound = Round; + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales source tensor + ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta + ): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr, + ElementCompute const *beta_ptr + ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { + + } + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + FragmentTensor predicate_mask_; + bool participates_in_reduction_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationDReluConditionalBits(Params const ¶ms) { + + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + participates_in_reduction_ = true; + predicate_mask_.clear(); + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + return beta_ != ElementCompute(0); + } + + /// Returns true if the threadblock computes the reduction + MCTLASS_HOST_DEVICE + bool participates_in_reduction() const { + return participates_in_reduction_; + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + predicate_mask_.clear(); + + if (k_partition) { + beta_ = ElementCompute(1); + } + + if (k_partition != k_partition_count - 1) { + // Avoid computing the reduction if this isn't the final Split-K slice + participates_in_reduction_ = false; + + bit_not not_op; + predicate_mask_ = not_op(predicate_mask_); + } + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_DEVICE + FragmentCompute operator()( + FragmentAccumulator const &accumulator, + FragmentSource const &source, + FragmentTensor const &tensor) const { + + // Convert source to interal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_source = source_converter(source); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + + intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + + bit_or or_op; + + FragmentTensor predicates = or_op(tensor, predicate_mask_); + + // Obtain from packed bits + bool conditions[kCount]; + UnpackPredicates unpack_predicates; + + unpack_predicates(conditions, predicates); + + // dReLU = (cond ? dy : 0) + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + if (!conditions[i]) { + intermediate[i] = ElementCompute(); + } + } + + return intermediate; + } + + /// Computes linear scaling: D = alpha * accumulator + MCTLASS_HOST_DEVICE + FragmentCompute operator()( + FragmentAccumulator const &accumulator, + FragmentTensor const &tensor) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_accumulator; + + intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum + + bit_or or_op; + + FragmentTensor predicates = or_op(tensor, predicate_mask_); + + // Obtain from packed bits + bool conditions[kCount]; + UnpackPredicates unpack_predicates; + + unpack_predicates(conditions, predicates); + + // dReLU = (cond ? dy : 0) + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + if (!conditions[i]) { + intermediate[i] = ElementCompute(); + } + } + + return intermediate; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_gelu.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_gelu.h new file mode 100644 index 0000000..6691284 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_gelu.h @@ -0,0 +1,70 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing linear combination with GELU operations used by epilogues. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/epilogue/thread/activation.h" +#include "mctlass/epilogue/thread/linear_combination_generic.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator followed by the GELU activation to an array of elements. +/// +/// D = gelu(alpha * accumulator + beta * source + uniform) +/// +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data to store + typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type + typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination + ScaleType::Kind Scale = ScaleType::Default, ///< Control Alpha and Beta scaling + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +using LinearCombinationGELU = LinearCombinationGeneric; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_generic.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_generic.h new file mode 100644 index 0000000..4d57eba --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_generic.h @@ -0,0 +1,208 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing linear combination operations used by epilogues. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/epilogue/thread/scale_type.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator followed by an activation function to an array of elements. +/// +/// D = activation(alpha * accumulator + beta * source + uniform) +/// +template < + template class ActivationFunctor, + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data to store + typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type + typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination + ScaleType::Kind Scale = ScaleType::Default, ///< Control Alpha and Beta scaling + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest, + bool IsHeavy = false +> +class LinearCombinationGeneric { +public: + + using ElementOutput = ElementOutput_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + + static bool const kIsHeavy = IsHeavy; + static int const kCount = Count; + static const ScaleType::Kind kScale = Scale; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using FragmentSource = Array; + using FragmentCompute = Array; + + static FloatRoundStyle const kRound = Round; + + /// Host-constructable parameters structure + using Params = typename ActivationFunctor::Params; + +private: + + // + // Data members + // + + Params params_; + bool skip_elementwise_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationGeneric(Params const ¶ms) { + params_ = params; + params_.alpha = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + params_.beta = (params.beta_ptr ? *params.beta_ptr : params.beta); + skip_elementwise_ = false; + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + if (Scale == ScaleType::NoBetaScaling) return true; + + if (Scale == ScaleType::OnlyAlphaScaling) return false; + + if (Scale == ScaleType::Nothing) return false; + + return params_.beta != ElementCompute(0); + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + params_.beta = ElementCompute(1); + } + + if (k_partition != k_partition_count - 1) { + skip_elementwise_ = true; + } + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentOutput const &source) const { + + // Convert source to interal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_source = source_converter(source); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + + FragmentCompute intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + ActivationFunctor activation; + + if (Scale == ScaleType::NoBetaScaling) { + intermediate = converted_source; + intermediate = mul_add_accumulator(params_.alpha, converted_accumulator, intermediate); // D = alpha * Accum + X + } else if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_add_source(params_.beta, converted_source); // X = beta * C + uniform + intermediate = mul_add_accumulator(params_.alpha, converted_accumulator, intermediate); // D = alpha * Accum + X + } + + intermediate = skip_elementwise_ ? intermediate : activation(intermediate, params_); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(intermediate); + } + + /// Computes linear scaling: D = alpha * accumulator + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + + FragmentCompute intermediate; + + multiplies mul_add_accumulator; + ActivationFunctor activation; + + if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_add_accumulator(params_.alpha, converted_accumulator); // D = alpha * Accum + } + + intermediate = skip_elementwise_ ? intermediate : activation(intermediate, params_); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(intermediate); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_hardswish.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_hardswish.h new file mode 100644 index 0000000..53acba5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_hardswish.h @@ -0,0 +1,69 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing linear combination with HardSwish operations used by epilogues. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/epilogue/thread/activation.h" +#include "mctlass/epilogue/thread/linear_combination_generic.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator followed by the HardSwish activation to an array of elements. +/// +/// D = hardswish(alpha * accumulator + beta * source + uniform) +/// +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data to store + typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type + typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination + ScaleType::Kind Scale = ScaleType::Default, ///< Control Alpha and Beta scaling + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +using LinearCombinationHardSwish = LinearCombinationGeneric; +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_leaky_relu.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_leaky_relu.h new file mode 100644 index 0000000..97bc350 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_leaky_relu.h @@ -0,0 +1,231 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/epilogue/thread/activation.h" +#include "mctlass/epilogue/thread/scale_type.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements. +/// +/// D = alpha * accumulator + beta * source + uniform +/// +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type + typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination + ScaleType::Kind Scale = ScaleType::Default, ///< Control Alpha and Beta scaling + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +class LinearCombinationLeakyRelu { +public: + + using ElementOutput = ElementOutput_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + + static int const kCount = Count; + static const ScaleType::Kind kScale = Scale; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using ComputeFragment = Array; + using FragmentSource = Array; + + static FloatRoundStyle const kRound = Round; + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta_bias; ///< scales bias tensor + ElementCompute leaky_alpha; ///< leaky_alpha + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta_bias(ElementCompute(0)), + leaky_alpha(ElementCompute(1)) + { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta_bias, + ElementCompute leaky_alpha = ElementCompute(1) + ): alpha(alpha), beta_bias(beta_bias), leaky_alpha(leaky_alpha) { + + } + + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_bias_; + ElementCompute leaky_alpha_recip_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationLeakyRelu(Params const ¶ms) { + alpha_ = (params.alpha); + beta_bias_ = (params.beta_bias); + leaky_alpha_recip_ = (ElementCompute(params.leaky_alpha)); + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + if (Scale == ScaleType::NoBetaScaling) return true; + + if (Scale == ScaleType::OnlyAlphaScaling) return false; + + if (Scale == ScaleType::Nothing) return false; + + return beta_bias_ != ElementCompute(0); + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition) { + if (k_partition) { + beta_bias_ = ElementCompute(1); + } + } + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_bias_ = ElementCompute(1); + } + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentOutput const &source) const { + + // Convert source to interal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + ComputeFragment converted_source = source_converter(source); + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + ComputeFragment intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + + LeakyReLU leakyrelu; + + if (Scale == ScaleType::NoBetaScaling) { + intermediate = converted_source; + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } else if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_add_source(beta_bias_, converted_source); // X = beta * C + uniform + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } + // Compute threshold optionally + intermediate = leakyrelu(intermediate, leaky_alpha_recip_); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(intermediate); + } + + /// Computes linear scaling: D = alpha * accumulator + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + ComputeFragment converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + ComputeFragment intermediate; + + multiplies mul_accumulator; + LeakyReLU leakyrelu; + //printf("in doing with bias"); + if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum + } + + // Compute threshold optionally + intermediate = leakyrelu(intermediate, leaky_alpha_recip_); + + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(intermediate); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_params.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_params.h new file mode 100644 index 0000000..9419a33 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_params.h @@ -0,0 +1,75 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief +*/ + +#pragma once + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +struct LinearCombinationParams { + uint64_t alpha_data[2]; + uint64_t beta_data[2]; + + MCTLASS_HOST_DEVICE + LinearCombinationParams() + : alpha_data {0lu, 0lu}, beta_data {0lu, 0lu} + { } + + template + MCTLASS_HOST_DEVICE + LinearCombinationParams(ElementCompute alpha, ElementCompute beta) + : alpha_data {0lu, 0lu}, beta_data {0lu, 0lu} + { +#if defined(__MACA_ARCH__) + reinterpret_cast(alpha_data) = alpha; + reinterpret_cast(beta_data) = beta; +#else + memcpy( alpha_data, &alpha, sizeof(ElementCompute) ); + memcpy( beta_data, &beta, sizeof(ElementCompute) ); +#endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_planar_complex.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_planar_complex.h new file mode 100644 index 0000000..50c9e1b --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_planar_complex.h @@ -0,0 +1,237 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing linear combination operations on planar-complex arrays +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/complex.h" +#include "mctlass/array_planar_complex.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to arrays of planar-complex elements. +/// +/// D = alpha * accumulator + beta * source + uniform +/// +/// Note, as with most MCTLASS components for planar complex, the template arguments describe +/// the underlying real data type. +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data to store + typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type + typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +class LinearCombinationPlanarComplex { +public: + + using ElementOutput = ElementOutput_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + + static int const kCount = Count; + + using FragmentOutput = ArrayPlanarComplex; + using FragmentAccumulator = ArrayPlanarComplex; + using ComputeFragment = ArrayPlanarComplex; + + static FloatRoundStyle const kRound = Round; + + /// Host-constructable parameters structure + struct Params { + + complex alpha; ///< scales accumulators + complex beta; ///< scales source tensor + complex const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + complex const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + complex alpha, + complex beta + ): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + complex const *alpha_ptr, + complex const *beta_ptr + ): alpha(complex()), beta(complex()), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { + + } + }; + +private: + + // + // Data members + // + + complex alpha_; + complex beta_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationPlanarComplex(Params const ¶ms) { + + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + return beta_.real() != ElementCompute(0) || beta_.imag() != ElementCompute(0); + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentOutput const &source) const { + + // Convert source to interal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + ComputeFragment converted_source( + source_converter(source.real), + source_converter(source.imag)); + + ComputeFragment converted_accumulator( + accumulator_converter(accumulator.real), + accumulator_converter(accumulator.imag)); + + // Perform binary operations + ComputeFragment intermediate; + + multiplies > mul_op; + multiply_add > mul_add_op; + + // complex multiply: I = beta * C + intermediate.real = mul_op(beta_.real(), converted_source.real); + intermediate.imag = mul_op(beta_.real(), converted_source.imag); + + intermediate.real = mul_add_op(-beta_.imag(), converted_source.imag, intermediate.real); + intermediate.imag = mul_add_op( beta_.imag(), converted_source.real, intermediate.imag); + + // complex multiply-add: I = alpha * AB + I + intermediate.real = mul_add_op(alpha_.real(), converted_accumulator.real, intermediate.real); + intermediate.imag = mul_add_op(alpha_.real(), converted_accumulator.imag, intermediate.imag); + + intermediate.real = mul_add_op(-alpha_.imag(), converted_accumulator.imag, intermediate.real); + intermediate.imag = mul_add_op( alpha_.imag(), converted_accumulator.real, intermediate.imag); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return FragmentOutput( + destination_converter(intermediate.real), + destination_converter(intermediate.imag)); + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + ComputeFragment converted_accumulator( + accumulator_converter(accumulator.real), + accumulator_converter(accumulator.imag)); + + // Perform binary operations + ComputeFragment intermediate; + + multiplies > mul_op; + multiply_add > mul_add_op; + + // complex multiply-add: I = alpha * AB + I + intermediate.real = mul_add_op(alpha_.real(), converted_accumulator.real); + intermediate.imag = mul_add_op(alpha_.real(), converted_accumulator.imag); + + intermediate.real = mul_add_op(-alpha_.imag(), converted_accumulator.imag, intermediate.real); + intermediate.imag = mul_add_op( alpha_.imag(), converted_accumulator.real, intermediate.imag); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return FragmentOutput( + destination_converter(intermediate.real), + destination_converter(intermediate.imag)); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_relu.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_relu.h new file mode 100644 index 0000000..7d3cfa0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_relu.h @@ -0,0 +1,572 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing linear combination with a maximum operation used by epilogues. +*/ + +#pragma once + +#include "mctlass/half.h" +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/epilogue/thread/activation.h" +#include "mctlass/epilogue/thread/scale_type.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +/// Single source of truth for whether to unroll for `LinearCombinationClamp()` +constexpr bool LinearCombinationReluIsHeavy() { + return false; +} + +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements. +/// +/// D = alpha * accumulator + beta * source + uniform +/// +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data to store + typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type + typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination + ScaleType::Kind Scale = ScaleType::Default, ///< Control Alpha and Beta scaling + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +class LinearCombinationRelu { +public: + + using ElementOutput = ElementOutput_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + + static int const kCount = Count; + static const ScaleType::Kind kScale = Scale; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using FragmentCompute = Array; + using FragmentScaleBias = Array; + using FragmentSource = Array; + + static FloatRoundStyle const kRound = Round; + + static bool const kIsHeavy = detail::LinearCombinationReluIsHeavy(); + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales source tensor + ElementCompute threshold; ///< minimum value that is output + ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute(0)), + threshold(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta = ElementCompute(0), + ElementCompute threshold = ElementCompute(0) + ): alpha(alpha), beta(beta), threshold(threshold), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr, + ElementCompute const *beta_ptr = nullptr, + ElementCompute threshold = ElementCompute(0) + ): alpha(0), beta(0), threshold(threshold), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { + + } + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + ElementCompute threshold_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationRelu(Params const ¶ms) { + + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + threshold_ = params.threshold; + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + if (Scale == ScaleType::NoBetaScaling) return true; + + if (Scale == ScaleType::OnlyAlphaScaling) return false; + + if (Scale == ScaleType::OnlyAlphaPerChannelScaling) return false; + + if (Scale == ScaleType::Nothing) return false; + + return beta_ != ElementCompute(0); + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + + if (k_partition != k_partition_count - 1) { + // set to NaN to make ReLU no-op for all except last k partitions + int64_t allones = -1; + threshold_ = reinterpret_cast(allones); + } + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentOutput const &source) const { + + // Convert source to interal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_source = source_converter(source); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + ReLu relu; + + if (Scale == ScaleType::NoBetaScaling) { + intermediate = converted_source; + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } else if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } + + // Compute threshold optionally + intermediate = relu(threshold_, intermediate); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(intermediate); + } + + /// Computes linear scaling: D = alpha * accumulator + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_accumulator; + ReLu relu; + + if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum + } + + // Compute threshold optionally + intermediate = relu(threshold_, intermediate); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(intermediate); + } + + /// Computes per-channel linear scaling and bias : D = scale * accumulator + bias + /// Scale and Bias are from input Fragment + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentScaleBias const &scale, + FragmentScaleBias const &bias) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform per-channel scale and bias + FragmentCompute intermediate; + + multiply_add mul_add_accumulator; + + if(Scale == ScaleType::OnlyAlphaPerChannelScaling) + intermediate = mul_add_accumulator(scale, converted_accumulator, bias); // D = scale * Accum + bias + else + intermediate = mul_add_accumulator(alpha_, converted_accumulator, bias); // D = alpha * Accum + bias + + ReLu relu; + + // Compute threshold optionally + intermediate = relu(threshold_, intermediate); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(intermediate); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Conditional guards to enable partial specialization for packed integers +#if defined(__MACA_ARCH__) + +/// Applies a linear combination operator to an array of elements. +/// +/// D = alpha * accumulator + beta * source + uniform +/// +/// Special handling for int types + +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + ScaleType::Kind Scale, ///< Control Alpha and Beta scaling + FloatRoundStyle Round +> +class LinearCombinationRelu { +public: + + using ElementOutput = ElementOutput_; + using ElementAccumulator = int; + using ElementCompute = float; + + static bool const kIsHeavy = detail::LinearCombinationReluIsHeavy(); + + static int const kCount = Count; + static const ScaleType::Kind kScale = Scale; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using FragmentCompute = Array; + using FragmentScaleBias = Array; + using FragmentSource = Array; + + static FloatRoundStyle const kRound = Round; + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales source tensor + ElementCompute threshold; ///< minimum value that is output + ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute(0)), + threshold(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta = ElementCompute(0), + ElementCompute threshold = ElementCompute(0) + ): alpha(alpha), beta(beta), threshold(threshold), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr, + ElementCompute const *beta_ptr = nullptr, + ElementCompute threshold = ElementCompute(0) + ): alpha(0), beta(0), threshold(threshold), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { + + } + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + ElementCompute threshold_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationRelu(Params const ¶ms) { + + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + threshold_ = params.threshold; + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + if (Scale == ScaleType::NoBetaScaling) return true; + + if (Scale == ScaleType::OnlyAlphaScaling) return false; + + if (Scale == ScaleType::OnlyAlphaPerChannelScaling) return false; + + if (Scale == ScaleType::Nothing) return false; + + return beta_ != ElementCompute(0); + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + + if (k_partition != k_partition_count - 1) { + // set to NaN to make ReLU no-op for all except last k partitions + int64_t allones = -1; + threshold_ = reinterpret_cast(allones); + } + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentOutput const &source) const { + + // Convert source to interal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_source = source_converter(source); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + ReLu relu; + + if (Scale == ScaleType::NoBetaScaling) { + intermediate = converted_source; + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } else if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } + + // Compute threshold optionally + intermediate = relu(threshold_, intermediate); + + if (platform::numeric_limits::is_integer) { + // Convert floats back to INT + FragmentAccumulator scaled_accumulator; + + NumericArrayConverter compute_converter; + + scaled_accumulator = compute_converter(intermediate); + + // Convert to destination numeric type + NumericArrayConverter + destination_converter; + + return destination_converter(scaled_accumulator); + } else { + NumericArrayConverter + destination_converter; + return destination_converter(intermediate); + } + } + + /// Computes linear scaling: D = alpha * accumulator + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_accumulator; + ReLu relu; + + if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum + } + + // Compute threshold optionally + intermediate = relu(threshold_, intermediate); + + if (platform::numeric_limits::is_integer) { + // Convert floats back to INT + FragmentAccumulator scaled_accumulator; + + NumericArrayConverter compute_converter; + + scaled_accumulator = compute_converter(intermediate); + + // Convert to destination numeric type + NumericArrayConverter + destination_converter; + + return destination_converter(scaled_accumulator); + } else { + NumericArrayConverter + destination_converter; + return destination_converter(intermediate); + } + } + + /// Computes per-channel linear scaling and bias : D = scale * accumulator + bias + /// Scale and Bias are from input Fragment + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentScaleBias const &scale, + FragmentScaleBias const &bias) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform per-channel scale and bias + FragmentCompute intermediate; + + multiply_add mul_add_accumulator; + + if(Scale == ScaleType::OnlyAlphaPerChannelScaling) + intermediate = mul_add_accumulator(scale, converted_accumulator, bias); // D = scale * Accum + bias + else + intermediate = mul_add_accumulator(alpha_, converted_accumulator, bias); // D = alpha * Accum + bias + + ReLu relu; + + // Compute threshold optionally + intermediate = relu(threshold_, intermediate); + + if (platform::numeric_limits::is_integer) { + // Convert floats back to INT + FragmentAccumulator scaled_accumulator; + + NumericArrayConverter compute_converter; + + scaled_accumulator = compute_converter(intermediate); + + // Convert to destination numeric type + NumericArrayConverter + destination_converter; + + return destination_converter(scaled_accumulator); + } else { + NumericArrayConverter + destination_converter; + return destination_converter(intermediate); + } + } +}; + +#endif // Conditional guards to enable partial specialization for packed integers + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_relu0.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_relu0.h new file mode 100644 index 0000000..ed7a420 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_relu0.h @@ -0,0 +1,543 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing linear combination with a relu operation used by epilogues. + This one only supports relu0 and tries to folding relu into other instructions. Thus, + serial splitk is not supported by this one. For example, relu can be folded into + hfma2/hmul2 for sm80+ +*/ + +#pragma once + +#include "mctlass/half.h" +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/epilogue/thread/activation.h" +#include "mctlass/epilogue/thread/scale_type.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +/// Single source of truth for whether to unroll for `LinearCombinationClamp()` +constexpr bool LinearCombinationRelu0IsHeavy() { + return false; +} + +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements. +/// +/// D = alpha * accumulator + beta * source + uniform +/// +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data to store + typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type + typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination + ScaleType::Kind Scale = ScaleType::Default, ///< Control Alpha and Beta scaling + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +class LinearCombinationRelu0 { +public: + + using ElementOutput = ElementOutput_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + + static int const kCount = Count; + static const ScaleType::Kind kScale = Scale; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using FragmentCompute = Array; + using FragmentScaleBias = Array; + using FragmentSource = Array; + + static FloatRoundStyle const kRound = Round; + + static bool const kIsHeavy = detail::LinearCombinationRelu0IsHeavy(); + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales source tensor + ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta = ElementCompute(0) + ): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr, + ElementCompute const *beta_ptr = nullptr + ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { + + } + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationRelu0(Params const ¶ms) { + + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + if (Scale == ScaleType::NoBetaScaling) return true; + + if (Scale == ScaleType::OnlyAlphaScaling) return false; + + if (Scale == ScaleType::Nothing) return false; + + return beta_ != ElementCompute(0); + } + + /// This is used for serial reduction which is not supported by Relu0 + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + assert(k_partition == 0); + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentOutput const &source) const { + + // Convert source to interal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_source = source_converter(source); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_add_source; + multiply_add_relu0 mul_add_relu0_accumulator; + ReLu relu; + + if (Scale == ScaleType::NoBetaScaling) { + intermediate = converted_source; + intermediate = mul_add_relu0_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } else if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + + // Compute threshold optionally + intermediate = relu(intermediate); + } else { + intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform + intermediate = mul_add_relu0_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(intermediate); + } + + /// Computes linear scaling: D = alpha * accumulator + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_accumulator; + ReLu relu; + + if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum + } + + // Compute threshold optionally + intermediate = relu(intermediate); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(intermediate); + } + + /// Computes per-channel linear scaling and bias : D = scale * accumulator + bias + /// Scale and Bias are from input Fragment + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentScaleBias const &scale, + FragmentScaleBias const &bias) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform per-channel scale and bias + FragmentCompute intermediate; + + multiply_add mul_add_accumulator; + + if(Scale == ScaleType::OnlyAlphaPerChannelScaling) + intermediate = mul_add_accumulator(scale, converted_accumulator, bias); // D = scale * Accum + bias + else + intermediate = mul_add_accumulator(alpha_, converted_accumulator, bias); // D = alpha * Accum + bias + + ReLu relu; + + // Compute threshold optionally + intermediate = relu(intermediate); + + // Convert to destination numeric type + NumericArrayConverter destination_converter; + + return destination_converter(intermediate); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Conditional guards to enable partial specialization for packed integers +#if defined(__MACA_ARCH__) + +/// Applies a linear combination operator to an array of elements. +/// +/// D = alpha * accumulator + beta * source + uniform +/// +/// Special handling for int types + +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + ScaleType::Kind Scale, ///< Control Alpha and Beta scaling + FloatRoundStyle Round +> +class LinearCombinationRelu0 { +public: + + using ElementOutput = ElementOutput_; + using ElementAccumulator = int; + using ElementCompute = float; + + static bool const kIsHeavy = detail::LinearCombinationRelu0IsHeavy(); + + static int const kCount = Count; + static const ScaleType::Kind kScale = Scale; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using FragmentCompute = Array; + using FragmentScaleBias = Array; + using FragmentSource = Array; + + static FloatRoundStyle const kRound = Round; + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales source tensor + ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta = ElementCompute(0) + ): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr, + ElementCompute const *beta_ptr = nullptr + ): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { + + } + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationRelu0(Params const ¶ms) { + + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + if (Scale == ScaleType::NoBetaScaling) return true; + + if (Scale == ScaleType::OnlyAlphaScaling) return false; + + if (Scale == ScaleType::Nothing) return false; + + return beta_ != ElementCompute(0); + } + + /// This is used for serial reduction which is not supported by Relu0 + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + assert(k_partition == 0); + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentOutput const &source) const { + + // Convert source to interal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_source = source_converter(source); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + ReLu relu; + + if (Scale == ScaleType::NoBetaScaling) { + intermediate = converted_source; + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } else if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + } + + // Compute threshold optionally + intermediate = relu(intermediate); + + if (platform::numeric_limits::is_integer) { + // Convert floats back to INT + FragmentAccumulator scaled_accumulator; + + NumericArrayConverter compute_converter; + + scaled_accumulator = compute_converter(intermediate); + + // Convert to destination numeric type + NumericArrayConverter + destination_converter; + + return destination_converter(scaled_accumulator); + } else { + NumericArrayConverter + destination_converter; + return destination_converter(intermediate); + } + } + + /// Computes linear scaling: D = alpha * accumulator + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_accumulator; + ReLu relu; + + if (Scale == ScaleType::Nothing) { + intermediate = converted_accumulator; + } else { + intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum + } + + // Compute threshold optionally + intermediate = relu(intermediate); + + if (platform::numeric_limits::is_integer) { + // Convert floats back to INT + FragmentAccumulator scaled_accumulator; + + NumericArrayConverter compute_converter; + + scaled_accumulator = compute_converter(intermediate); + + // Convert to destination numeric type + NumericArrayConverter + destination_converter; + + return destination_converter(scaled_accumulator); + } else { + NumericArrayConverter + destination_converter; + return destination_converter(intermediate); + } + } + + /// Computes per-channel linear scaling and bias : D = scale * accumulator + bias + /// Scale and Bias are from input Fragment + MCTLASS_HOST_DEVICE + FragmentOutput operator()( + FragmentAccumulator const &accumulator, + FragmentScaleBias const &scale, + FragmentScaleBias const &bias) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform per-channel scale and bias + FragmentCompute intermediate; + + multiply_add mul_add_accumulator; + + if(Scale == ScaleType::OnlyAlphaPerChannelScaling) + intermediate = mul_add_accumulator(scale, converted_accumulator, bias); // D = scale * Accum + bias + else + intermediate = mul_add_accumulator(alpha_, converted_accumulator, bias); // D = alpha * Accum + bias + + ReLu relu; + + // Compute threshold optionally + intermediate = relu(intermediate); + + if (platform::numeric_limits::is_integer) { + // Convert floats back to INT + FragmentAccumulator scaled_accumulator; + + NumericArrayConverter compute_converter; + + scaled_accumulator = compute_converter(intermediate); + + // Convert to destination numeric type + NumericArrayConverter + destination_converter; + + return destination_converter(scaled_accumulator); + } else { + NumericArrayConverter + destination_converter; + return destination_converter(intermediate); + } + } +}; + +#endif // Conditional guards to enable partial specialization for packed integers + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_residual_block.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_residual_block.h new file mode 100644 index 0000000..0f2063c --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_residual_block.h @@ -0,0 +1,301 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Epilogue functor specialized for residual blocks in deep neural network. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/epilogue/thread/detail.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +/// Models a residual block of the form: UnaryOp(BinaryOp(BinaryOp(ActivationOp(TensorOp(X) + bias), residual1), residual2)) +template class ActivationOp_, + template class BinaryOp1_, + template class UnaryOp_, + template class BinaryOp2_ = detail::NoOp, + bool StoreT_ = false, + typename ElementVector_ = ElementC_> +class LinearCombinationResidualBlock { +public: + static bool const kIsSingleSource = false; + + using ElementOutput = ElementC_; + using ElementC = ElementC_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + using ElementVector = ElementVector_; + static int const kElementsPerAccess = ElementsPerAccess; + static int const kCount = kElementsPerAccess; + + using UnaryOp = UnaryOp_>; + using BinaryOp1 = BinaryOp1_>; + using BinaryOp2 = BinaryOp2_>; + using ActivationOp = ActivationOp_>; + + using FragmentAccumulator = Array; + using FragmentCompute = Array; + using FragmentC = Array; + using FragmentOutput = Array; + + using ElementZ = ElementOutput_; + using ElementT = ElementZ; + using FragmentZ = Array; + using FragmentT = Array; + + static bool const kIsHeavy = true; + static bool const kStoreZ = true; + static bool const kStoreT = StoreT_; + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales residual input + ElementCompute const *alpha_ptr{nullptr}; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr{nullptr}; ///< pointer to residual scalar - if not null, loads it from memory + + MCTLASS_HOST_DEVICE + Params() : alpha(ElementCompute(1)), beta(ElementCompute(1)) {} + + MCTLASS_HOST_DEVICE + Params(ElementCompute alpha, ElementCompute beta) + : alpha(alpha), beta(beta) {} + + MCTLASS_HOST_DEVICE + Params(ElementCompute const *alpha_ptr, ElementCompute const *beta_ptr) + : alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {} + }; + +private: + + ElementCompute alpha_; + ElementCompute beta_; + bool skip_elementwise_; + +public: + + /// Constructor from Params + MCTLASS_HOST_DEVICE + LinearCombinationResidualBlock(Params const ¶ms) { + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + skip_elementwise_ = false; + } + + /// The "source" tensor corresponds to the residual input + MCTLASS_HOST_DEVICE + bool is_source_needed() const { return true; } + + /// Functionally required for serial reduction in the epilogue + /// IMPORTANT: Split-k is supported only when ActivationOp is Identity. + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + + if (k_partition != k_partition_count - 1) { + skip_elementwise_ = true; + } + } + + /// Applies the operation UnaryOp(BinaryOp(BinaryOp(ActivationOp(AB + bias), residual1), residual2)) + MCTLASS_HOST_DEVICE + void operator()(FragmentOutput &frag_Z, FragmentOutput &, FragmentAccumulator const &AB, + FragmentC const &residual1, FragmentC const &residual2, + FragmentCompute const &bias) const { + UnaryOp unary_op; + BinaryOp1 binary_op1; + BinaryOp2 binary_op2; + ActivationOp activation; + + FragmentCompute tmp_Accum = + NumericArrayConverter()(AB); + FragmentCompute tmp_residual1 = + NumericArrayConverter()(residual1); + FragmentCompute tmp_residual2 = + NumericArrayConverter()(residual2); + + FragmentCompute z = + binary_op2(binary_op1(activation(alpha_ * tmp_Accum + bias), beta_ * tmp_residual1), beta_ * tmp_residual2); + FragmentCompute result_Z = skip_elementwise_ ? z : unary_op(z); + + NumericArrayConverter convert_z; + frag_Z = convert_z(result_Z); + } + + /// Should never be called + MCTLASS_HOST_DEVICE + void operator()(FragmentOutput &, FragmentOutput &, FragmentAccumulator const &, + FragmentCompute const &) const {} +}; + +/// Models a residual block of the form: UnaryOp(BinaryOp(ActivationOp(TensorOp(X) + bias), residual)) +template class ActivationOp_, + template class BinaryOp1_, + template class UnaryOp_, + bool StoreT_, + typename ElementVector_> +class LinearCombinationResidualBlock { +public: + static bool const kIsSingleSource = true; + + using ElementOutput = ElementC_; + using ElementC = ElementC_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + using ElementVector = ElementVector_; + static int const kElementsPerAccess = ElementsPerAccess; + static int const kCount = kElementsPerAccess; + + using UnaryOp = UnaryOp_>; + using BinaryOp = BinaryOp1_>; + using ActivationOp = ActivationOp_>; + + using FragmentAccumulator = Array; + using FragmentCompute = Array; + using FragmentC = Array; + using FragmentOutput = Array; + + using ElementZ = ElementOutput_; + using ElementT = ElementZ; + using FragmentZ = Array; + using FragmentT = Array; + + static bool const kIsHeavy = true; + static bool const kStoreZ = true; + static bool const kStoreT = StoreT_; + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales residual input + ElementCompute const *alpha_ptr{nullptr}; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr{nullptr}; ///< pointer to residual scalar - if not null, loads it from memory + + MCTLASS_HOST_DEVICE + Params() : alpha(ElementCompute(1)), beta(ElementCompute(1)) {} + + MCTLASS_HOST_DEVICE + Params(ElementCompute alpha, ElementCompute beta) + : alpha(alpha), beta(beta) {} + + MCTLASS_HOST_DEVICE + Params(ElementCompute const *alpha_ptr, ElementCompute const *beta_ptr) + : alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {} + }; + +private: + + ElementCompute alpha_; + ElementCompute beta_; + bool skip_elementwise_; + +public: + + /// Constructor from Params + MCTLASS_HOST_DEVICE + LinearCombinationResidualBlock(Params const ¶ms) { + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + skip_elementwise_ = false; + } + + /// The "source" tensor corresponds to the residual input + MCTLASS_HOST_DEVICE + bool is_source_needed() const { return true; } + + /// Functionally required for serial reduction in the epilogue + /// IMPORTANT: Split-k is supported only when ActivationOp is Identity. + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + + if (k_partition != k_partition_count - 1) { + skip_elementwise_ = true; + } + } + + /// Applies the operation UnaryOp(BinaryOp(ActivationOp(AB + bias), residual)) + MCTLASS_HOST_DEVICE + void operator()(FragmentOutput &frag_Z, FragmentOutput &, FragmentAccumulator const &AB, + FragmentC const &residual, + FragmentCompute const &bias) const { + UnaryOp unary_op; + BinaryOp binary_op; + ActivationOp activation; + + FragmentCompute tmp_Accum = + NumericArrayConverter()(AB); + FragmentCompute tmp_residual = + NumericArrayConverter()(residual); + + FragmentCompute z = + binary_op(activation(alpha_ * tmp_Accum + bias), beta_ * tmp_residual); + FragmentCompute result_Z = skip_elementwise_ ? z : unary_op(z); + + NumericArrayConverter convert_z; + frag_Z = convert_z(result_Z); + } + + /// Should never be called + MCTLASS_HOST_DEVICE + void operator()(FragmentOutput &, FragmentOutput &, FragmentAccumulator const &, + FragmentCompute const &) const {} +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_sigmoid.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_sigmoid.h new file mode 100644 index 0000000..2d89b69 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_sigmoid.h @@ -0,0 +1,70 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing linear combination with Sigmoid operations used by epilogues. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/epilogue/thread/activation.h" +#include "mctlass/epilogue/thread/linear_combination_generic.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator followed by the Sigmoid activation, to an array of elements. +/// +/// D = sigmoid(alpha * accumulator + beta * source + uniform) +/// +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data to store + typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type + typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination + ScaleType::Kind Scale = ScaleType::Default, ///< Control Alpha and Beta scaling + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +using LinearCombinationSigmoid = LinearCombinationGeneric; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_silu.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_silu.h new file mode 100644 index 0000000..86d47e6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_silu.h @@ -0,0 +1,69 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing linear combination with SiLU operations used by epilogues. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/epilogue/thread/activation.h" +#include "mctlass/epilogue/thread/linear_combination_generic.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator folllowed by the SiLU activation to an array of elements. +/// +/// D = silu(alpha * accumulator + beta * source + uniform) +/// +template < + typename ElementOutput_, ///< Data type used to load and store tensors + int Count, ///< Number of elements computed per operation + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data to store + typename ElementAccumulator_ = ElementOutput_, ///< Accumulator data type + typename ElementCompute_ = ElementOutput_, ///< Data type used to compute linear combination + ScaleType::Kind Scale = ScaleType::Default, ///< Control Alpha and Beta scaling + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +using LinearCombinationSilu = LinearCombinationGeneric; +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_tensor_broadcast.hpp b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_tensor_broadcast.hpp new file mode 100644 index 0000000..978e82a --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_tensor_broadcast.hpp @@ -0,0 +1,252 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Functor performing linear combination operation, bias addition, and tensor-tensor + elementwise operations +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/numeric_types.h" +#include "mctlass/epilogue/thread/activation.h" +#include "mctlass/epilogue/thread/detail.hpp" +#include "mctlass/epilogue/thread/scale_type.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +namespace detail { + +/// Returns whether a source operand is needed for a combination of binary operation and scale +/// type. Simple specialized checks are made for cases in which 0 is an identity element of +/// the binary operation. +template +MCTLASS_HOST_DEVICE +bool is_binary_op_source_needed(ElementCompute scale) { + if constexpr (cute::is_same_v>) { + return false; + } + else if constexpr (cute::is_same_v> || cute::is_same_v>) { + // Cases for binary operators for which 0 is an identity element + if constexpr (Scale == ScaleType::NoBetaScaling) return true; + + if constexpr (Scale == ScaleType::OnlyAlphaScaling) return false; + + if constexpr (Scale == ScaleType::Nothing) return false; + + return scale != ElementCompute(0); + } + + return true; +} + +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/** Compute a tensor-tensor broadcast epilogue. + * + * @param ElementOutput_ Data type used to load and store tensors + * @param ElementAccumulator_ Accumulator data type + * @param ElementCompute_ Data type used to compute linear combination + * @param ElementBias_ Data type of Bias elements + * @param ActivationFunctor_ Fused Activation + * @param BinaryOp0_ Binary operation to perform on O0 and C0. detail::NoOp means no operation + * @param BinaryOp1_ Binary operation to perform on O1 and C1. detail::NoOp means no operation + * @param UnaryOp_ Unary operation to perform on final result + * @param Scale Controls the type of Alpha and Beta scaling to perform + * @param Round How values should be rounded in conversions + * @param ElementSource_ Data type used for source operands + * + * Computes the following: + * O0 = alpha * accumulator + bias + * O1 = BinaryOp0(O0, beta * C0) + * O2 = BinaryOp1(O1, beta * C1) + * D = UnaryOp(O2) + */ +template < + class ElementOutput_, + class ElementAccumulator_ = ElementOutput_, + class ElementCompute_ = ElementOutput_, + class ElementBias_ = ElementCompute_, + template class ActivationFunctor_ = Identity, + template class BinaryOp0_ = plus, + template class BinaryOp1_ = detail::NoOp, + template class UnaryOp_ = Identity, + ScaleType::Kind Scale = ScaleType::Default, + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest, + class ElementSource_ = ElementOutput_ +> +class LinearCombinationTensorBroadcast { +public: + + using ElementOutput = ElementOutput_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + using ElementBias = ElementBias_; + using ElementC = ElementSource_; + using ElementD = ElementOutput_; + using ElementScalingFactor = ElementAccumulator_; + + using UnaryOp = UnaryOp_; + using BinaryOp0 = BinaryOp0_; + using BinaryOp1 = BinaryOp1_; + using ActivationFunctor = ActivationFunctor_; + + static constexpr int kCount = 1; + static constexpr ScaleType::Kind kScale = Scale; + + using FragmentOutput = Array; + using FragmentAccumulator = Array; + using ComputeFragment = Array; + using FragmentBias = Array; + + static constexpr FloatRoundStyle kRound = Round; + using NoOpType = detail::NoOp; + static constexpr bool IsBinaryOp0Enabled = !cute::is_same_v; + static constexpr bool IsBinaryOp1Enabled = !cute::is_same_v; + static constexpr bool IsUnaryOpEnabled = !cute::is_same_v && !cute::is_same_v>; + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha{}; ///< scales accumulators + ElementCompute beta{}; ///< scales source tensor + ElementCompute const* alpha_ptr = nullptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const* beta_ptr = nullptr; ///< pointer to source scalar - if not null, loads it from memory + + // + // Methods + // + Params() = default; + + MCTLASS_HOST_DEVICE + Params(ElementCompute const* alpha_ptr, ElementCompute const* beta_ptr) + : alpha_ptr(alpha_ptr), + beta_ptr(beta_ptr) {} + + MCTLASS_HOST_DEVICE + Params(ElementCompute const* alpha_ptr) + : alpha_ptr(alpha_ptr) {} + + MCTLASS_HOST_DEVICE + Params(ElementCompute alpha, + ElementCompute beta) + : alpha(alpha), + beta(beta) {} + }; + +private: + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationTensorBroadcast(Params const& params) + : alpha_(params.alpha_ptr ? *params.alpha_ptr : params.alpha), + beta_(params.beta_ptr ? *params.beta_ptr : params.beta) {} + + /// Returns true if source 0 is needed + MCTLASS_HOST_DEVICE + bool is_source0_needed() const { + return detail::is_binary_op_source_needed(beta_); + } + + /// Returns true if source 1 is needed + MCTLASS_HOST_DEVICE + bool is_source1_needed() const { + return detail::is_binary_op_source_needed(beta_); + } + + // + // Specialization for scalar + // + MCTLASS_HOST_DEVICE + ElementD operator()(ElementAccumulator const accumulator, ElementC const source0, ElementC source1, ElementBias const bias) { + // Convert everything to Compute type, do compute, and then store to output type + NumericConverter accumulator_converter; + NumericConverter bias_converter; + NumericConverter source_converter; + NumericConverter destination_converter; + + ActivationFunctor act; + multiplies mul; + multiply_add madd; + + ElementCompute intermediate = accumulator_converter(accumulator); + intermediate = madd(alpha_, intermediate, bias_converter(bias)); + intermediate = act(intermediate); + + // Apply BinaryOp0, if needed + if constexpr (IsBinaryOp0Enabled) { + BinaryOp0 bin0; + ElementCompute converted_source = source_converter(source0); + intermediate = bin0(intermediate, mul(beta_, converted_source)); + } + + // Apply BinaryOp1, if needed + if constexpr (IsBinaryOp1Enabled) { + BinaryOp1 bin1; + ElementCompute converted_source = source_converter(source1); + intermediate = bin1(intermediate, mul(beta_, converted_source)); + } + + // Apply UnaryOp, if needed + if constexpr (IsUnaryOpEnabled) { + UnaryOp unary; + intermediate = unary(intermediate); + } + + return destination_converter(intermediate); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_with_elementwise.h b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_with_elementwise.h new file mode 100644 index 0000000..0c77dae --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/linear_combination_with_elementwise.h @@ -0,0 +1,234 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + + \brief Functor performing linear combination with elementwise +*/ + +#pragma once + +#include "mctlass/half.h" +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/constants.h" +#include "mctlass/fast_math.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/epilogue/thread/activation.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a linear combination operator to an array of elements. +/// +/// D = alpha * accumulator + beta * source + uniform +/// +template < + typename ElementCompute_, ///< Data type returned by this functor + typename ElementAccumulator_, ///< Data type of accumulators + typename ElementSource_, ///< Data type of source tensor + typename ElementTensor_, ///< Data type of additional tensor + int Count, ///< Number of elements computed per operation + ///< Usually it is 128/sizeof_bits, + ///< but we use 64 or 32 sometimes when there are not enough data to store + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +class LinearCombinationWithElementwise { +public: + + using ElementOutput = ElementSource_; + using ElementCompute = ElementCompute_; + using ElementAccumulator = ElementAccumulator_; + using ElementSource = ElementSource_; + using ElementTensor = ElementTensor_; + + static bool const kIsHeavy = true; + + static int const kCount = Count; + + using FragmentCompute = Array; + using FragmentAccumulator = Array; + using FragmentSource = Array; + using FragmentTensor = Array; + + static FloatRoundStyle const kRound = Round; + + /// Host-constructable parameters structure + struct Params { + + ElementCompute alpha; ///< scales accumulators + ElementCompute beta; ///< scales source tensor + ElementCompute threshold; ///< minimum value that is output + ElementCompute const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory + ElementCompute const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + alpha(ElementCompute(1)), + beta(ElementCompute(0)), + threshold(ElementCompute(0)), + alpha_ptr(nullptr), + beta_ptr(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute alpha, + ElementCompute beta, + ElementCompute threshold = ElementCompute(0) + ): alpha(alpha), beta(beta), threshold(threshold), alpha_ptr(nullptr), beta_ptr(nullptr) { + + } + + MCTLASS_HOST_DEVICE + Params( + ElementCompute const *alpha_ptr, + ElementCompute const *beta_ptr, + ElementCompute threshold = ElementCompute(0) + ): alpha(0), beta(0), threshold(threshold), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { + + } + }; + +private: + + // + // Data members + // + + ElementCompute alpha_; + ElementCompute beta_; + ElementCompute threshold_; + bool participates_in_reduction_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + LinearCombinationWithElementwise(Params const ¶ms) { + + alpha_ = (params.alpha_ptr ? *params.alpha_ptr : params.alpha); + beta_ = (params.beta_ptr ? *params.beta_ptr : params.beta); + threshold_ = params.threshold; + participates_in_reduction_ = true; + } + + /// Returns true if source is needed + MCTLASS_HOST_DEVICE + bool is_source_needed() const { + return beta_ != ElementCompute(0); + } + + /// Returns true if the threadblock computes the reduction + MCTLASS_HOST_DEVICE + bool participates_in_reduction() const { + return participates_in_reduction_; + } + + /// Functionally required for serial reduction in the epilogue + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { + if (k_partition) { + beta_ = ElementCompute(1); + } + + if (k_partition != k_partition_count - 1) { + // set to NaN to make ReLU no-op for all except last k partitions + int64_t allones = -1; + threshold_ = reinterpret_cast(allones); + // Avoid computing the reduction if this isn't the final Split-K slice + participates_in_reduction_ = false; + } + } + + /// Computes linear scaling: D = alpha * accumulator + beta * source + MCTLASS_HOST_DEVICE + FragmentCompute operator()( + FragmentAccumulator const &accumulator, + FragmentSource const &source, + FragmentTensor const &tensor) const { + + // Convert source to interal compute numeric type + NumericArrayConverter source_converter; + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_source = source_converter(source); + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_add_source; + multiply_add mul_add_accumulator; + + intermediate = mul_add_source(beta_, converted_source); // X = beta * C + uniform + intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate); // D = alpha * Accum + X + + return intermediate; + } + + /// Computes linear scaling: D = alpha * accumulator + MCTLASS_HOST_DEVICE + FragmentCompute operator()( + FragmentAccumulator const &accumulator, + FragmentTensor const &tensor) const { + + // Convert source to interal compute numeric type + NumericArrayConverter accumulator_converter; + + FragmentCompute converted_accumulator = accumulator_converter(accumulator); + + // Perform binary operations + FragmentCompute intermediate; + + multiplies mul_accumulator; + + intermediate = mul_accumulator(alpha_, converted_accumulator); // D = alpha * Accum + + return intermediate; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/reduction_op.h b/csrc/mctlass/include/mctlass/epilogue/thread/reduction_op.h new file mode 100644 index 0000000..9519e81 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/reduction_op.h @@ -0,0 +1,97 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Functor performing reduction operations used by epilogues. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Applies a reduction sum to an array of elements. +/// +/// +template < + typename Element_, ///< Data type used to load and store tensors + int Count ///< Number of elements computed per operation +> +class ReductionOpPlus { +public: + + using Element = Element_; + static int const kCount = Count; + + using Fragment = Array; + using Operator = plus; + + /// Host-constructable parameters structure + struct Params { }; + +private: + + /// reduction operator + Operator operator_; + +public: + + /// Constructs the function object, possibly loading from pointers in host memory + MCTLASS_HOST_DEVICE + ReductionOpPlus(Params const ¶ms) { + + } + + /// Computes Compute => + MCTLASS_HOST_DEVICE + Fragment operator()( + Fragment const &lhs, + Fragment const &rhs) const { + + return operator_(lhs, rhs); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/thread/scale_type.h b/csrc/mctlass/include/mctlass/epilogue/thread/scale_type.h new file mode 100644 index 0000000..ae3e796 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/thread/scale_type.h @@ -0,0 +1,62 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Enum defines the behaviors of the epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specifies internal data type for computation +struct ScaleType { + enum Kind { + Default, // alpha x C + beta x D + NoBetaScaling, // alpha x C + D + OnlyAlphaScaling, // alpha x C + OnlyAlphaPerChannelScaling, // alpha_vec x C + Nothing // C + }; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_complex_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_complex_tensor_op.h new file mode 100644 index 0000000..3d18a23 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_complex_tensor_op.h @@ -0,0 +1,255 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped complex GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/epilogue/thread/linear_combination.h" +#include "mctlass/epilogue/thread/linear_combination_relu.h" +#include "mctlass/epilogue/thread/linear_combination_gelu.h" +#include "mctlass/epilogue/thread/linear_combination_sigmoid.h" +#include "mctlass/epilogue/thread/linear_combination_planar_complex.h" + +#include "mctlass/epilogue/thread/conversion_op.h" +#include "mctlass/epilogue/thread/reduction_op.h" + +#include "mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h" + +#include "mctlass/epilogue/warp/fragment_iterator_complex_tensor_op.h" +#include "mctlass/epilogue/warp/fragment_iterator_gaussian_complex_tensor_op.h" +#include "mctlass/epilogue/warp/tile_iterator_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_thread_map_tensor_op.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator.h" +#include "mctlass/epilogue/threadblock/shared_load_iterator.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Specialization and defines sensible defaults for epilogues for complex*complex case +// 4 real-valued mma operations (Complex) +// A = (ar + j ai), B (br +j bi), D = AB +// D = dr + j di = (ar*br - ai*bi) + j (ar*bi + ai*br) +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Epilouge Shape + typename Shape_, + /// Warp-level mma operator + typename WarpMmaTensorOp_, + /// Number of k partitions + int PartitionsK, + /// Epilogue output operator + typename OutputOp_, + /// Elements accessed by inner-most loop of AccumulatorFragmentIterator::load() + int ElementsPerAccess, + /// Multiply-add operator + /// Selects between (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator_ = arch::OpMultiplyAddComplex +> +struct DefaultEpilogueComplexTensorOp { + + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + using Operator = Operator_; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapTensorOp< + Shape, + typename WarpMmaTensorOp::Shape, + kPartitionsK, + ElementOutput, + kElementsPerAccess + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + OutputTileThreadMap, + ElementOutput + >; + + using AccumulatorFragmentIterator = mctlass::epilogue::warp::FragmentIteratorComplexTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC + >; + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + ElementAccumulator, + LayoutC + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + typename OutputTileThreadMap::CompactedThreadMap, + ElementAccumulator + >; + + /// Hard-coded padding elements added + using Padding = mctlass::MatrixShape<0, 0>; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaTensorOp, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Partial specialization and defines sensible defaults for epilogues for complex*complex case +// 3 real-valued mma operations (Gaussian Complex) +// A = (ar + j ai), B = (br +j bi), D = AB +// P1 = (ar + ai) * br, P2 = - ar * (br - bi), P3 = ai * (br + bi) +// D = dr + j di = (P1 - P3) + j (P1 + P2) +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + typename Shape_, + typename WarpMmaTensorOp_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess +> +struct DefaultEpilogueComplexTensorOp { + + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + using Operator = arch::OpMultiplyAddGaussianComplex; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapTensorOp< + Shape, + typename WarpMmaTensorOp::Shape, + kPartitionsK, + ElementOutput, + kElementsPerAccess + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + OutputTileThreadMap, + ElementOutput + >; + + using AccumulatorFragmentIterator = mctlass::epilogue::warp::FragmentIteratorGaussianComplexTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC + >; + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + ElementAccumulator, + LayoutC + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + typename OutputTileThreadMap::CompactedThreadMap, + ElementAccumulator + >; + + /// Hard-coded padding elements added + using Padding = mctlass::MatrixShape<0, 0>; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaTensorOp, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_complex_tensor_op_blas3.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_complex_tensor_op_blas3.h new file mode 100644 index 0000000..bdb8ec5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_complex_tensor_op_blas3.h @@ -0,0 +1,264 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped complex GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/epilogue/thread/linear_combination.h" +#include "mctlass/epilogue/thread/linear_combination_relu.h" +#include "mctlass/epilogue/thread/linear_combination_gelu.h" +#include "mctlass/epilogue/thread/linear_combination_sigmoid.h" +#include "mctlass/epilogue/thread/linear_combination_planar_complex.h" + +#include "mctlass/epilogue/thread/conversion_op.h" +#include "mctlass/epilogue/thread/reduction_op.h" + +#include "mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h" + +#include "mctlass/epilogue/warp/fragment_iterator_complex_tensor_op.h" +#include "mctlass/epilogue/warp/fragment_iterator_gaussian_complex_tensor_op.h" +#include "mctlass/epilogue/warp/tile_iterator_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_thread_map_tensor_op.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_blas3.h" +#include "mctlass/epilogue/threadblock/shared_load_iterator.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Specialization and defines sensible defaults for epilogues for complex*complex case +// 4 real-valued mma operations (Complex) +// A = (ar + j ai), B (br +j bi), D = AB +// D = dr + j di = (ar*br - ai*bi) + j (ar*bi + ai*br) +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Epilogue Shape + typename Shape_, + /// Warp-level mma operator + typename WarpMmaTensorOp_, + /// Number of k partitions + int PartitionsK, + /// Epilogue output operator + typename OutputOp_, + /// Elements accessed by inner-most loop of AccumulatorFragmentIterator::load() + int ElementsPerAccess, + /// Multiply-add operator + /// Selects between (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator_ = arch::OpMultiplyAddComplex, + /// Is for a symmetric kernel + BlasMode BlasMode_ = BlasMode::kGemm +> +struct DefaultEpilogueComplexTensorOpBlas3 { + + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + using Operator = Operator_; + static BlasMode const kBlasMode = BlasMode_; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapTensorOp< + Shape, + typename WarpMmaTensorOp::Shape, + kPartitionsK, + ElementOutput, + kElementsPerAccess + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIteratorBlas3< + OutputTileThreadMap, + ElementOutput + , kBlasMode + >; + + using AccumulatorFragmentIterator = mctlass::epilogue::warp::FragmentIteratorComplexTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC + >; + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + ElementAccumulator, + LayoutC + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + typename OutputTileThreadMap::CompactedThreadMap, + ElementAccumulator + >; + + /// Hard-coded padding elements added + using Padding = mctlass::MatrixShape<0, 0>; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaTensorOp, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Partial specialization and defines sensible defaults for epilogues for complex*complex case +// 3 real-valued mma operations (Gaussian Complex) +// A = (ar + j ai), B = (br +j bi), D = AB +// P1 = (ar + ai) * br, P2 = - ar * (br - bi), P3 = ai * (br + bi) +// D = dr + j di = (P1 - P3) + j (P1 + P2) +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + typename Shape_, + typename WarpMmaTensorOp_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess, + BlasMode BlasMode_ +> +struct DefaultEpilogueComplexTensorOpBlas3 { + + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + using Operator = arch::OpMultiplyAddGaussianComplex; + static BlasMode const kBlasMode = BlasMode_; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapTensorOp< + Shape, + typename WarpMmaTensorOp::Shape, + kPartitionsK, + ElementOutput, + kElementsPerAccess + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIteratorBlas3< + OutputTileThreadMap, + ElementOutput, + kBlasMode + >; + + using AccumulatorFragmentIterator = mctlass::epilogue::warp::FragmentIteratorGaussianComplexTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC + >; + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + ElementAccumulator, + LayoutC + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + typename OutputTileThreadMap::CompactedThreadMap, + ElementAccumulator + >; + + /// Hard-coded padding elements added + using Padding = mctlass::MatrixShape<0, 0>; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaTensorOp, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_direct_store.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_direct_store.h new file mode 100644 index 0000000..477be95 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_direct_store.h @@ -0,0 +1,74 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Direct store epilogue +*/ + +#pragma once + +//////////////////////////////////////////////////////////////////////////////// + +#include "mctlass/epilogue/threadblock/epilogue_direct_store.h" +#include "mctlass/epilogue/threadblock/direct_store_epilogue_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Given a properly constructed epilogue, returns a direct store epilogue +template +struct DefaultEpilogueDirectStore { + + using OutputTileIterator = DirectStoreEpilogueIterator; + + using Epilogue = EpilogueDirectStore< + typename EpilogueTensorOp::Shape, + typename EpilogueTensorOp::WarpMmaOperator, + EpilogueTensorOp::kPartitionsK, + OutputTileIterator, + typename EpilogueTensorOp::AccumulatorFragmentIterator, + typename EpilogueTensorOp::WarpTileIterator, + typename EpilogueTensorOp::SharedLoadIterator, + typename EpilogueTensorOp::OutputOp + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_planar_complex.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_planar_complex.h new file mode 100644 index 0000000..3a75464 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_planar_complex.h @@ -0,0 +1,241 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Constructs a default epilogue for planar complex outputs. + + This template reuses components for real-valued epilogues and applies them to planar complex + output matrices. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/array_planar_complex.h" + +#include "mctlass/arch/arch.h" + +#include "mctlass/epilogue/thread/linear_combination_planar_complex.h" +#include "mctlass/epilogue/threadblock/default_epilogue_simt.h" +#include "mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op.h" + +#include "mctlass/epilogue/threadblock/epilogue_planar_complex.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues. +template < + typename ThreadblockShape_, + typename WarpMma_, + typename OpcodeClass_, + typename ArchTag_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess +> +struct DefaultEpiloguePlanarComplex; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues. +template < + typename ThreadblockShape_, + typename WarpMmaOperator_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess +> +struct DefaultEpiloguePlanarComplex< + ThreadblockShape_, + WarpMmaOperator_, + arch::OpClassTensorOp, + arch::Sm70, + PartitionsK, + OutputOp_, + ElementsPerAccess> { + + using RealEpilogue = DefaultEpilogueVoltaTensorOp< + ThreadblockShape_, + WarpMmaOperator_, + PartitionsK, + OutputOp_, + ElementsPerAccess + >; + + using Epilogue = EpiloguePlanarComplex< + ThreadblockShape_, + WarpMmaOperator_, + PartitionsK, + typename RealEpilogue::OutputTileIterator, + typename RealEpilogue::AccumulatorFragmentIterator, + typename RealEpilogue::WarpTileIterator, + typename RealEpilogue::SharedLoadIterator, + OutputOp_, + typename RealEpilogue::Padding + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues. +template < + typename ThreadblockShape_, + typename WarpMmaOperator_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess +> +struct DefaultEpiloguePlanarComplex< + ThreadblockShape_, + WarpMmaOperator_, + arch::OpClassTensorOp, + arch::Sm75, + PartitionsK, + OutputOp_, + ElementsPerAccess> { + + using RealEpilogue = DefaultEpilogueTensorOp< + ThreadblockShape_, + WarpMmaOperator_, + PartitionsK, + OutputOp_, + ElementsPerAccess + >; + + using Epilogue = EpiloguePlanarComplex< + ThreadblockShape_, + WarpMmaOperator_, + PartitionsK, + typename RealEpilogue::OutputTileIterator, + typename RealEpilogue::AccumulatorFragmentIterator, + typename RealEpilogue::WarpTileIterator, + typename RealEpilogue::SharedLoadIterator, + OutputOp_, + typename RealEpilogue::Padding + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues. +template < + typename ThreadblockShape_, + typename WarpMmaOperator_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess +> +struct DefaultEpiloguePlanarComplex< + ThreadblockShape_, + WarpMmaOperator_, + arch::OpClassTensorOp, + arch::Sm80, + PartitionsK, + OutputOp_, + ElementsPerAccess> { + + using RealEpilogue = DefaultEpilogueTensorOp< + ThreadblockShape_, + WarpMmaOperator_, + PartitionsK, + OutputOp_, + ElementsPerAccess + >; + + using Epilogue = EpiloguePlanarComplex< + ThreadblockShape_, + WarpMmaOperator_, + PartitionsK, + typename RealEpilogue::OutputTileIterator, + typename RealEpilogue::AccumulatorFragmentIterator, + typename RealEpilogue::WarpTileIterator, + typename RealEpilogue::SharedLoadIterator, + OutputOp_, + typename RealEpilogue::Padding + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues. +template < + typename ThreadblockShape_, + typename WarpMmaOperator_, + typename ArchTag_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess +> +struct DefaultEpiloguePlanarComplex< + ThreadblockShape_, + WarpMmaOperator_, + arch::OpClassSimt, + ArchTag_, + PartitionsK, + OutputOp_, + ElementsPerAccess> { + + using RealEpilogue = DefaultEpilogueSimt< + ThreadblockShape_, + WarpMmaOperator_, + OutputOp_, + ElementsPerAccess + >; + + using Epilogue = EpiloguePlanarComplex< + ThreadblockShape_, + WarpMmaOperator_, + PartitionsK, + typename RealEpilogue::OutputTileIterator, + typename RealEpilogue::AccumulatorFragmentIterator, + typename RealEpilogue::WarpTileIterator, + typename RealEpilogue::SharedLoadIterator, + OutputOp_, + typename RealEpilogue::Padding + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_simt.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_simt.h new file mode 100644 index 0000000..e77c96c --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_simt.h @@ -0,0 +1,422 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using SIMT. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" + +#include "mctlass/arch/mma.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/epilogue/thread/linear_combination.h" +#include "mctlass/epilogue/thread/linear_combination_clamp.h" +#include "mctlass/epilogue/thread/linear_combination_relu.h" +#include "mctlass/epilogue/thread/linear_combination_gelu.h" +#include "mctlass/epilogue/thread/linear_combination_sigmoid.h" +#include "mctlass/epilogue/thread/linear_combination_planar_complex.h" +#include "mctlass/epilogue/thread/conversion_op.h" +#include "mctlass/epilogue/thread/reduction_op.h" + +#include "mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h" + +#include "mctlass/epilogue/warp/fragment_iterator_simt.h" +#include "mctlass/epilogue/warp/tile_iterator_simt.h" +#include "mctlass/epilogue/threadblock/default_thread_map_simt.h" +#include "mctlass/transform/pitch_linear_thread_map.h" + +#include "mctlass/epilogue/threadblock/predicated_tile_iterator.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_strided_dgrad.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_affine.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_direct_conv.h" +#include "mctlass/epilogue/threadblock/shared_load_iterator.h" +#include "mctlass/epilogue/threadblock/shared_load_iterator_pitch_liner.h" +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/threadblock/epilogue_depthwise.h" + +#include "mctlass/layout/permute.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for SimtOps. +template < + typename Shape_, + typename WarpMmaSimt_, + typename OutputOp_, + int ElementsPerAccess, + bool ScatterD = false, + typename PermuteDLayout = layout::NoPermute +> +struct DefaultEpilogueSimt { + + using Shape = Shape_; + using WarpMmaSimt = WarpMmaSimt_; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + static const int kPartitionsK = Shape::kK / WarpMmaSimt::Shape::kK; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaSimt::LayoutC; + using ElementAccumulator = typename WarpMmaSimt::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapSimt< + Shape, + typename WarpMmaSimt::Shape, + typename WarpMmaSimt::Policy, + kPartitionsK, + ElementOutput, + kElementsPerAccess + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + OutputTileThreadMap, + ElementOutput, + ScatterD, + PermuteDLayout + >; + + using AccumulatorFragmentIterator = mctlass::epilogue::warp::FragmentIteratorSimt< + typename WarpMmaSimt::Shape, + typename WarpMmaSimt::ThreadMma, + layout::RowMajor, + typename WarpMmaSimt::Policy + >; + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorSimt< + typename WarpMmaSimt::Shape, + typename WarpMmaSimt::ThreadMma, + ElementAccumulator, + layout::RowMajor, + typename WarpMmaSimt::Policy + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + typename OutputTileThreadMap::CompactedThreadMap, + ElementAccumulator + >; + + /// Hard-coded padding elements added + using Padding = typename WarpTileIterator::Padding; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaSimt, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for SimtOps. +template < + typename Shape_, + typename WarpMmaSimt_, + typename OutputOp_, + int ElementsPerAccess +> +struct DefaultEpilogueSimtStridedDgrad { + + using Shape = Shape_; + using WarpMmaSimt = WarpMmaSimt_; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + static const int kPartitionsK = Shape::kK / WarpMmaSimt::Shape::kK; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaSimt::LayoutC; + using ElementAccumulator = typename WarpMmaSimt::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapSimt< + Shape, + typename WarpMmaSimt::Shape, + typename WarpMmaSimt::Policy, + kPartitionsK, + ElementOutput, + kElementsPerAccess + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIteratorStridedDgrad< + OutputTileThreadMap, + ElementOutput + >; + + using AccumulatorFragmentIterator = mctlass::epilogue::warp::FragmentIteratorSimt< + typename WarpMmaSimt::Shape, + typename WarpMmaSimt::ThreadMma, + layout::RowMajor, + typename WarpMmaSimt::Policy + >; + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorSimt< + typename WarpMmaSimt::Shape, + typename WarpMmaSimt::ThreadMma, + ElementAccumulator, + layout::RowMajor, + typename WarpMmaSimt::Policy + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + typename OutputTileThreadMap::CompactedThreadMap, + ElementAccumulator + >; + + /// Hard-coded padding elements added + using Padding = typename WarpTileIterator::Padding; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaSimt, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for SimtOps. +template < + int Rank, + typename Shape_, + typename WarpMmaSimt_, + typename OutputOp_, + int ElementsPerAccess +> +struct DefaultEpilogueSimtAffineRankN { + + using Shape = Shape_; + using WarpMmaSimt = WarpMmaSimt_; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + static const int kPartitionsK = Shape::kK / WarpMmaSimt::Shape::kK; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaSimt::LayoutC; + using ElementAccumulator = typename WarpMmaSimt::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapSimt< + Shape, + typename WarpMmaSimt::Shape, + typename WarpMmaSimt::Policy, + kPartitionsK, + ElementOutput, + kElementsPerAccess + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIteratorAffineRankN< + OutputTileThreadMap, + ElementOutput, + Rank + >; + + using AccumulatorFragmentIterator = mctlass::epilogue::warp::FragmentIteratorSimt< + typename WarpMmaSimt::Shape, + typename WarpMmaSimt::ThreadMma, + layout::RowMajor, + typename WarpMmaSimt::Policy + >; + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorSimt< + typename WarpMmaSimt::Shape, + typename WarpMmaSimt::ThreadMma, + ElementAccumulator, + layout::RowMajor, + typename WarpMmaSimt::Policy + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + typename OutputTileThreadMap::CompactedThreadMap, + ElementAccumulator + >; + + /// Hard-coded padding elements added + using Padding = typename WarpTileIterator::Padding; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaSimt, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for SimtOps. +template , + typename ThreadBlockOutputShape_ = mctlass::conv::TensorNHWCShape<1, 1, 1, 1> > +struct DefaultDirectConvEpilogueSimt { + using Shape = Shape_; + using WarpMmaSimt = WarpMmaSimt_; + using WarpShape = typename WarpMmaSimt::Shape; + using OutputOp = OutputOp_; + using ThreadOutputShape = ThreadOutputShape_; + using ThreadBlockOutputShape = ThreadBlockOutputShape_; + static int const kElementsPerAccess = ElementsPerAccess_; + + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaSimt::LayoutC; + using ElementAccumulator = typename WarpMmaSimt::ElementC; + + /// Number of threads total + using WarpCount = gemm::GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN + >; + + static int const kWarpSize = mctlass::gemm::warp::WarpSize::value; + + static int const kThreads = WarpCount::kCount * kWarpSize; + + // + // Thread map + // + + using OutputTileThreadMap = mctlass::transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIteratorDirectConv< + OutputTileThreadMap, + ElementOutput, + ThreadOutputShape, + ThreadBlockOutputShape + >; + + using AccumulatorFragmentIterator = mctlass::epilogue::warp::FragmentIteratorSimt< + typename WarpMmaSimt::Shape, + typename WarpMmaSimt::ThreadMma, + layout::RowMajor, + typename WarpMmaSimt::Policy + >; + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorSimtDirect2dConv< + typename WarpMmaSimt::Shape, + ThreadOutputShape, + ThreadBlockOutputShape, + typename WarpMmaSimt::ThreadMma, + ElementAccumulator, + layout::RowMajor, + typename WarpMmaSimt::Policy + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIteratorPitchLiner< + OutputTileThreadMap, + ElementAccumulator + >; + + /// Hard-coded padding elements added + using Padding = typename WarpTileIterator::Padding; + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::EpilogueDepthwise< + Shape, + ThreadOutputShape, + ThreadBlockOutputShape, + WarpMmaSimt, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_tensor_op.h new file mode 100644 index 0000000..ff0e48e --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_tensor_op.h @@ -0,0 +1,848 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" + +#include "mctlass/platform/platform.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/epilogue/thread/linear_combination.h" +#include "mctlass/epilogue/thread/linear_combination_clamp.h" +#include "mctlass/epilogue/thread/linear_combination_relu.h" +#include "mctlass/epilogue/thread/linear_combination_relu0.h" +#include "mctlass/epilogue/thread/linear_combination_gelu.h" +#include "mctlass/epilogue/thread/linear_combination_sigmoid.h" +#include "mctlass/epilogue/thread/linear_combination_hardswish.h" +#include "mctlass/epilogue/thread/linear_combination_planar_complex.h" + +#include "mctlass/epilogue/thread/conversion_op.h" +#include "mctlass/epilogue/thread/reduction_op.h" + +#include "mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h" + +#include "mctlass/epilogue/warp/fragment_iterator_tensor_op.h" +#include "mctlass/epilogue/warp/fragment_iterator_complex_tensor_op.h" +#include "mctlass/epilogue/warp/tile_iterator_tensor_op.h" +#include "mctlass/epilogue/warp/tile_iterator_tensor_op_mixed.h" +#include "mctlass/epilogue/threadblock/default_thread_map_tensor_op.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_strided_dgrad.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_affine.h" +#include "mctlass/epilogue/threadblock/shared_load_iterator.h" +#include "mctlass/epilogue/threadblock/shared_load_iterator_mixed.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/threadblock/interleaved_epilogue.h" + +#include "mctlass/layout/permute.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +template < + typename ElementOutput, + typename ElementAccumulator, + int ElementsPerAccess, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename ThreadMap +> +struct DefaultIteratorsTensorOp { + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorTensorOp< + WarpShape, + InstructionShape, + ElementAccumulator, + layout::RowMajor + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + ThreadMap, + ElementAccumulator + >; + + static int const kFragmentsPerIteration = 1; +}; + +/// Partial specialization for float <= float x 4 +template < + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename ThreadMap +> +struct DefaultIteratorsTensorOp { + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorTensorOp< + WarpShape, + InstructionShape, + float, + layout::RowMajor + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + ThreadMap, + float + >; + + static int const kFragmentsPerIteration = 2; +}; + +/// Partial specialization for int32_t <= int32_t x 4 +template < + int ElementsPerAccess, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename ThreadMap +> +struct DefaultIteratorsTensorOp { + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorTensorOp< + WarpShape, + InstructionShape, + int32_t, + layout::RowMajor + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + ThreadMap, + int32_t + >; + + static int const kFragmentsPerIteration = 1; +}; + +/// Partial specialization for float <= int32_t x 4 +template < + int ElementsPerAccess, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename ThreadMap +> +struct DefaultIteratorsTensorOp { + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorTensorOp< + WarpShape, + InstructionShape, + int32_t, + layout::RowMajor + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + ThreadMap, + int32_t + >; + + static int const kFragmentsPerIteration = 1; +}; + +/// Partial specialization for half <= float x 8 epilogues avoids shared memory bank conflicts. +template < + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename ThreadMap +> +struct DefaultIteratorsTensorOp< + half_t, + float, + 8, + ThreadblockShape, + WarpShape, + InstructionShape, + ThreadMap> { + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorTensorOpMixed< + WarpShape, + InstructionShape, + float, + 32, + 16, + 8, + 8 + >; + +using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIteratorMixed< + ThreadMap, + float, + 32, + 16, + 8, + 8 + >; + + static int const kFragmentsPerIteration = 2; +}; + +/// Partial specialization for half <= int32_t x 8 epilogues avoids shared memory bank conflicts. +template < + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename ThreadMap +> +struct DefaultIteratorsTensorOp< + half_t, + int32_t, + 8, + ThreadblockShape, + WarpShape, + InstructionShape, + ThreadMap> { + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorTensorOpMixed< + WarpShape, + InstructionShape, + int32_t, + 32, + 16, + 8, + 8 + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIteratorMixed< + ThreadMap, + int32_t, + 32, + 16, + 8, + 8 + >; + + static int const kFragmentsPerIteration = 2; +}; + +/// Partial specialization for int8/int4b_t <= int32 x 16/8 epilogues avoids shared memory bank conflicts. +/// Threadblock::kN = 256 still has bank conflicts. +template < + typename ElementOutput, + int ElementsPerAccess, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename ThreadMap +> +struct DefaultIteratorsTensorOp< + ElementOutput, + int32_t, + ElementsPerAccess, + ThreadblockShape, + WarpShape, + InstructionShape, + ThreadMap> { + + static_assert(platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value, + "ElementOutput needs to be 4 or 8 bit (unsigned) int."); + + static_assert((ElementsPerAccess == 16 || ElementsPerAccess == 8), + "ElementsPerAccess needs to be 16 or 8."); + + using WarpTileIteratorMixed = mctlass::epilogue::warp::TileIteratorTensorOpMixed< + WarpShape, + InstructionShape, + int32_t, + 32, + mctlass::sizeof_bits::value, + ElementsPerAccess, + 8 + >; + + using WarpTileIteratorNotMixed = mctlass::epilogue::warp::TileIteratorTensorOp< + WarpShape, + InstructionShape, + int32_t, + layout::RowMajor + >; + + using WarpTileIterator = typename platform::conditional< + (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8), + WarpTileIteratorNotMixed, + WarpTileIteratorMixed>::type; + + using SharedLoadIteratorMixed = mctlass::epilogue::threadblock::SharedLoadIteratorMixed< + ThreadMap, + int32_t, + 32, + mctlass::sizeof_bits::value, + ElementsPerAccess, + 8 + >; + + using SharedLoadIteratorNotMixed = mctlass::epilogue::threadblock::SharedLoadIterator< + ThreadMap, + int32_t + >; + + using SharedLoadIterator = typename platform::conditional< + (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8), + SharedLoadIteratorNotMixed, + SharedLoadIteratorMixed>::type; + + static int const kFragmentsPerIteration = 1; +}; + +/// Partial specialization for float_e4m3_t <= float x 16/8 epilogues avoids shared memory bank conflicts. +/// Threadblock::kN = 256 still has bank conflicts. +template < + int ElementsPerAccess, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename ThreadMap +> +struct DefaultIteratorsTensorOp< + mctlass::float_e4m3_t, + float, + ElementsPerAccess, + ThreadblockShape, + WarpShape, + InstructionShape, + ThreadMap> { + + using ElementOutput = mctlass::float_e4m3_t; + + static_assert((ElementsPerAccess == 16 || ElementsPerAccess == 8), + "ElementsPerAccess needs to be 16 or 8."); + + using WarpTileIteratorMixed = mctlass::epilogue::warp::TileIteratorTensorOpMixed< + WarpShape, + InstructionShape, + float, + 32, + mctlass::sizeof_bits::value, + ElementsPerAccess, + 8 + >; + + using WarpTileIteratorNotMixed = mctlass::epilogue::warp::TileIteratorTensorOp< + WarpShape, + InstructionShape, + float, + layout::RowMajor + >; + + using WarpTileIterator = typename platform::conditional< + (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8), + WarpTileIteratorNotMixed, + WarpTileIteratorMixed>::type; + + using SharedLoadIteratorMixed = mctlass::epilogue::threadblock::SharedLoadIteratorMixed< + ThreadMap, + float, + 32, + mctlass::sizeof_bits::value, + ElementsPerAccess, + 8 + >; + + using SharedLoadIteratorNotMixed = mctlass::epilogue::threadblock::SharedLoadIterator< + ThreadMap, + float + >; + + using SharedLoadIterator = typename platform::conditional< + (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8), + SharedLoadIteratorNotMixed, + SharedLoadIteratorMixed>::type; + + static int const kFragmentsPerIteration = 1; +}; + +/// Partial specialization for float_e5m2_t <= float x 16/8 epilogues avoids shared memory bank conflicts. +/// Threadblock::kN = 256 still has bank conflicts. +template < + int ElementsPerAccess, + typename ThreadblockShape, + typename WarpShape, + typename InstructionShape, + typename ThreadMap +> +struct DefaultIteratorsTensorOp< + mctlass::float_e5m2_t, + float, + ElementsPerAccess, + ThreadblockShape, + WarpShape, + InstructionShape, + ThreadMap> { + + using ElementOutput = mctlass::float_e5m2_t; + + static_assert((ElementsPerAccess == 16 || ElementsPerAccess == 8), + "ElementsPerAccess needs to be 16 or 8."); + + using WarpTileIteratorMixed = mctlass::epilogue::warp::TileIteratorTensorOpMixed< + WarpShape, + InstructionShape, + float, + 32, + mctlass::sizeof_bits::value, + ElementsPerAccess, + 8 + >; + + using WarpTileIteratorNotMixed = mctlass::epilogue::warp::TileIteratorTensorOp< + WarpShape, + InstructionShape, + float, + layout::RowMajor + >; + + using WarpTileIterator = typename platform::conditional< + (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8), + WarpTileIteratorNotMixed, + WarpTileIteratorMixed>::type; + + using SharedLoadIteratorMixed = mctlass::epilogue::threadblock::SharedLoadIteratorMixed< + ThreadMap, + float, + 32, + mctlass::sizeof_bits::value, + ElementsPerAccess, + 8 + >; + + using SharedLoadIteratorNotMixed = mctlass::epilogue::threadblock::SharedLoadIterator< + ThreadMap, + float + >; + + using SharedLoadIterator = typename platform::conditional< + (ThreadblockShape::kN == 256) || (ThreadblockShape::kN == 128 && ElementsPerAccess == 8), + SharedLoadIteratorNotMixed, + SharedLoadIteratorMixed>::type; + + static int const kFragmentsPerIteration = 1; +}; + +} // namespace detail + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for TensorOps. +template < + typename Shape_, + typename WarpMmaTensorOp_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess, + bool ScatterD = false, + typename PermuteDLayout = layout::NoPermute +> +struct DefaultEpilogueTensorOp { + + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapTensorOp< + Shape, + typename WarpMmaTensorOp::Shape, + kPartitionsK, + ElementOutput, + kElementsPerAccess + >::Type; + + static bool const UseCUDAStore = platform::is_same::value; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + OutputTileThreadMap, + ElementOutput, + ScatterD, + PermuteDLayout, + UseCUDAStore + >; + + using AccumulatorFragmentIterator = typename platform::conditional::value, + mctlass::epilogue::warp::FragmentIteratorComplexTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC>, + mctlass::epilogue::warp::FragmentIteratorTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC> >::type; + + /// Support several implementations depending on structure of epilogue + using DefaultIterators = detail::DefaultIteratorsTensorOp< + ElementOutput, + ElementAccumulator, + kElementsPerAccess, + Shape, + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename OutputTileThreadMap::CompactedThreadMap + >; + + using WarpTileIterator = typename DefaultIterators::WarpTileIterator; + using SharedLoadIterator = typename DefaultIterators::SharedLoadIterator; + + /// Hard-coded padding elements added + using Padding = mctlass::MatrixShape<0, 64 / sizeof_bits::value * 4>; + + static int const kFragmentsPerIteration = (kPartitionsK == 1 ? DefaultIterators::kFragmentsPerIteration : 1); + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaTensorOp, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding, + kFragmentsPerIteration + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for TensorOps. +template < + typename Shape_, + typename WarpMmaTensorOp_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess +> +struct DefaultEpilogueTensorOpStridedDgrad { + + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapTensorOp< + Shape, + typename WarpMmaTensorOp::Shape, + kPartitionsK, + ElementOutput, + kElementsPerAccess + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIteratorStridedDgrad< + OutputTileThreadMap, + ElementOutput + >; + + using AccumulatorFragmentIterator = typename platform::conditional::value, + mctlass::epilogue::warp::FragmentIteratorComplexTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC>, + mctlass::epilogue::warp::FragmentIteratorTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC> >::type; + + /// Support several implementations depending on structure of epilogue + using DefaultIterators = detail::DefaultIteratorsTensorOp< + ElementOutput, + ElementAccumulator, + kElementsPerAccess, + Shape, + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename OutputTileThreadMap::CompactedThreadMap + >; + + using WarpTileIterator = typename DefaultIterators::WarpTileIterator; + using SharedLoadIterator = typename DefaultIterators::SharedLoadIterator; + + /// Hard-coded padding elements added + using Padding = mctlass::MatrixShape<0, 64 / sizeof_bits::value * 4>; + + static int const kFragmentsPerIteration = (kPartitionsK == 1 ? DefaultIterators::kFragmentsPerIteration : 1); + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaTensorOp, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding, + kFragmentsPerIteration + >; +}; + + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for TensorOps. +template < + int Rank, + typename Shape_, + typename WarpMmaTensorOp_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess +> +struct DefaultEpilogueTensorOpAffineRankN { + + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapTensorOp< + Shape, + typename WarpMmaTensorOp::Shape, + kPartitionsK, + ElementOutput, + kElementsPerAccess + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIteratorAffineRankN< + OutputTileThreadMap, + ElementOutput, + Rank + >; + + // Map to the row major iterator since the iterator selection for affineN is the same. + using AccumulatorFragmentIterator = typename platform::conditional::value, + mctlass::epilogue::warp::FragmentIteratorComplexTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + layout::RowMajor>, + mctlass::epilogue::warp::FragmentIteratorTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + layout::RowMajor> >::type; + + /// Support several implementations depending on structure of epilogue + using DefaultIterators = detail::DefaultIteratorsTensorOp< + ElementOutput, + ElementAccumulator, + kElementsPerAccess, + Shape, + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename OutputTileThreadMap::CompactedThreadMap + >; + + using WarpTileIterator = typename DefaultIterators::WarpTileIterator; + using SharedLoadIterator = typename DefaultIterators::SharedLoadIterator; + + /// Hard-coded padding elements added + using Padding = mctlass::MatrixShape<0, 64 / sizeof_bits::value * 4>; + + static int const kFragmentsPerIteration = (kPartitionsK == 1 ? DefaultIterators::kFragmentsPerIteration : 1); + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaTensorOp, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding, + kFragmentsPerIteration + >; +}; + +//////////////////////////////////////////////////////////////////////////////// +/// Defines sensible defaults for epilogues for TensorOps which uses +/// intereleaved output layout. For this case, shared memory is not needed. +template +struct DefaultInterleavedEpilogueTensorOp { + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + using OutputTileThreadMap = typename mctlass::epilogue::threadblock:: + DefaultInterleavedThreadMapTensorOp< + Shape, typename WarpMmaTensorOp::Shape, kPartitionsK, ElementOutput, + kElementsPerAccess, InterleavedK>::Type; + + using OutputTileIterator = + mctlass::epilogue::threadblock::InterleavedPredicatedTileIterator< + OutputTileThreadMap, ElementOutput, InterleavedK>; + + using AccumulatorFragmentIterator = + mctlass::epilogue::warp::FragmentIteratorTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC>; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::InterleavedEpilogue< + Shape, WarpMmaTensorOp, kPartitionsK, OutputTileIterator, + AccumulatorFragmentIterator, OutputOp, InterleavedK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for TensorOps which uses +/// intereleaved output layout. For this case, shared memory is not needed. +template +struct DefaultInterleavedConvEpilogue { + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + + using ElementOutput = typename OutputOp::ElementOutput; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + using OutputTileThreadMap = typename mctlass::epilogue::threadblock:: + DefaultInterleavedConvThreadMapTensorOp< + Shape, typename WarpMmaTensorOp::Shape, kPartitionsK, ElementOutput, + kElementsPerAccess, InterleavedK>::Type; + + using OutputTileIterator = + mctlass::epilogue::threadblock::InterleavedConvPredicatedTileIterator< + OutputTileThreadMap, ElementOutput, InterleavedK>; + + using AccumulatorFragmentIterator = + mctlass::epilogue::warp::FragmentIteratorTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + // can reuse the gemm version here to do element selection + layout::ColumnMajorInterleaved>; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::InterleavedEpilogue< + Shape, WarpMmaTensorOp, kPartitionsK, OutputTileIterator, + AccumulatorFragmentIterator, OutputOp, InterleavedK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_tensor_op_blas3.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_tensor_op_blas3.h new file mode 100644 index 0000000..5ece71a --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_tensor_op_blas3.h @@ -0,0 +1,175 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/epilogue/thread/linear_combination.h" +#include "mctlass/epilogue/thread/linear_combination_clamp.h" +#include "mctlass/epilogue/thread/linear_combination_relu.h" +#include "mctlass/epilogue/thread/linear_combination_gelu.h" +#include "mctlass/epilogue/thread/linear_combination_sigmoid.h" +#include "mctlass/epilogue/thread/linear_combination_planar_complex.h" + +#include "mctlass/epilogue/thread/conversion_op.h" +#include "mctlass/epilogue/thread/reduction_op.h" + +#include "mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h" + +#include "mctlass/epilogue/warp/fragment_iterator_tensor_op.h" +#include "mctlass/epilogue/warp/fragment_iterator_complex_tensor_op.h" +#include "mctlass/epilogue/warp/tile_iterator_tensor_op.h" +#include "mctlass/epilogue/warp/tile_iterator_tensor_op_mixed.h" +#include "mctlass/epilogue/threadblock/default_thread_map_tensor_op.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_blas3.h" +#include "mctlass/epilogue/threadblock/shared_load_iterator.h" +#include "mctlass/epilogue/threadblock/shared_load_iterator_mixed.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/threadblock/interleaved_epilogue.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for TensorOps. +template < + typename Shape_, + typename WarpMmaTensorOp_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess, + /// Is for a symmetric kernel + BlasMode BlasMode_ = BlasMode::kGemm +> +struct DefaultEpilogueTensorOpBlas3 { + + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + static BlasMode const kBlasMode = BlasMode_; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapTensorOp< + Shape, + typename WarpMmaTensorOp::Shape, + kPartitionsK, + ElementOutput, + kElementsPerAccess + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIteratorBlas3< + OutputTileThreadMap, + ElementOutput, + kBlasMode + >; + + using AccumulatorFragmentIterator = typename std::conditional::value, + mctlass::epilogue::warp::FragmentIteratorComplexTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC>, + mctlass::epilogue::warp::FragmentIteratorTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC> >::type; + + /// Support several implementations depending on structure of epilogue + using DefaultIterators = detail::DefaultIteratorsTensorOp< + ElementOutput, + ElementAccumulator, + kElementsPerAccess, + Shape, + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename OutputTileThreadMap::CompactedThreadMap + >; + + using WarpTileIterator = typename DefaultIterators::WarpTileIterator; + using SharedLoadIterator = typename DefaultIterators::SharedLoadIterator; + + /// Hard-coded padding elements added + using Padding = mctlass::MatrixShape<0, 64 / sizeof_bits::value * 4>; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaTensorOp, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_tensor_op_row_broadcast.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_tensor_op_row_broadcast.h new file mode 100644 index 0000000..b4241d0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_tensor_op_row_broadcast.h @@ -0,0 +1,183 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" + +#include "mctlass/platform/platform.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/epilogue/thread/linear_combination.h" +#include "mctlass/epilogue/thread/linear_combination_clamp.h" +#include "mctlass/epilogue/thread/linear_combination_relu.h" +#include "mctlass/epilogue/thread/linear_combination_relu0.h" +#include "mctlass/epilogue/thread/linear_combination_gelu.h" +#include "mctlass/epilogue/thread/linear_combination_sigmoid.h" +#include "mctlass/epilogue/thread/linear_combination_hardswish.h" +#include "mctlass/epilogue/thread/linear_combination_planar_complex.h" + +#include "mctlass/epilogue/thread/conversion_op.h" +#include "mctlass/epilogue/thread/reduction_op.h" + +#include "mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h" + +#include "mctlass/epilogue/warp/fragment_iterator_tensor_op.h" +#include "mctlass/epilogue/warp/fragment_iterator_complex_tensor_op.h" +#include "mctlass/epilogue/warp/tile_iterator_tensor_op.h" +#include "mctlass/epilogue/warp/tile_iterator_tensor_op_mixed.h" +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_thread_map_tensor_op.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_row_broadcast.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_strided_dgrad.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_affine.h" +#include "mctlass/epilogue/threadblock/shared_load_iterator.h" +#include "mctlass/epilogue/threadblock/shared_load_iterator_mixed.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/threadblock/interleaved_epilogue.h" + +#include "mctlass/layout/permute.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +template < + typename Shape_, + typename WarpMmaTensorOp_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess, + bool ScatterD = false, + typename PermuteDLayout = layout::NoPermute +> +struct DefaultEpilogueTensorOpRowBroadcast { + + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapTensorOp< + Shape, + typename WarpMmaTensorOp::Shape, + kPartitionsK, + ElementOutput, + kElementsPerAccess + >::Type; + + static bool const UseCUDAStore = platform::is_same::value; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIteratorRowBroadcast< + OutputTileThreadMap, + ElementOutput, + ScatterD, + PermuteDLayout, + UseCUDAStore + >; + + using AccumulatorFragmentIterator = typename platform::conditional::value, + mctlass::epilogue::warp::FragmentIteratorComplexTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC>, + mctlass::epilogue::warp::FragmentIteratorTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC> >::type; + + /// Support several implementations depending on structure of epilogue + using DefaultIterators = detail::DefaultIteratorsTensorOp< + ElementOutput, + ElementAccumulator, + kElementsPerAccess, + Shape, + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename OutputTileThreadMap::CompactedThreadMap + >; + + using WarpTileIterator = typename DefaultIterators::WarpTileIterator; + using SharedLoadIterator = typename DefaultIterators::SharedLoadIterator; + + /// Hard-coded padding elements added + using Padding = mctlass::MatrixShape<0, 64 / sizeof_bits::value * 4>; + + static int const kFragmentsPerIteration = (kPartitionsK == 1 ? DefaultIterators::kFragmentsPerIteration : 1); + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaTensorOp, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding, + kFragmentsPerIteration + >; +}; + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h new file mode 100644 index 0000000..7bac21e --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h @@ -0,0 +1,337 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops on Volta. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/epilogue/thread/linear_combination.h" +#include "mctlass/epilogue/thread/linear_combination_clamp.h" +#include "mctlass/epilogue/thread/linear_combination_relu.h" +#include "mctlass/epilogue/thread/linear_combination_gelu.h" +#include "mctlass/epilogue/thread/linear_combination_sigmoid.h" +#include "mctlass/epilogue/thread/linear_combination_planar_complex.h" + +#include "mctlass/epilogue/thread/conversion_op.h" +#include "mctlass/epilogue/thread/reduction_op.h" + +#include "mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_strided_dgrad.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_affine.h" +#include "mctlass/epilogue/threadblock/shared_load_iterator.h" + +#include "mctlass/epilogue/warp/fragment_iterator_volta_tensor_op.h" +#include "mctlass/epilogue/warp/tile_iterator_volta_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_thread_map_volta_tensor_op.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" + +#include "mctlass/layout/permute.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for TensorOps. +template < + typename Shape_, + typename WarpMmaTensorOp_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess, + bool ScatterD = false, + typename PermuteDLayout = layout::NoPermute +> +struct DefaultEpilogueVoltaTensorOp { + + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapVoltaTensorOp< + Shape, + typename WarpMmaTensorOp::Shape, + kPartitionsK, + ElementOutput, + kElementsPerAccess, + ElementAccumulator + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + OutputTileThreadMap, + ElementOutput, + ScatterD, + PermuteDLayout + >; + + using AccumulatorFragmentIterator = mctlass::epilogue::warp::FragmentIteratorVoltaTensorOp< + typename WarpMmaTensorOp::Shape, + gemm::GemmShape<32, 32, 4>, + ElementAccumulator, + LayoutC + >; + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorVoltaTensorOp< + typename WarpMmaTensorOp::Shape, + gemm::GemmShape<32, 32, 4>, + ElementAccumulator, + LayoutC + >; + + static int const kSharedMemAlignment = sizeof_bits::value * WarpTileIterator::kElementsPerAccess / 8; + + static_assert(kSharedMemAlignment == 8, "Shared memory alignment must be 8B"); + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + typename OutputTileThreadMap::CompactedThreadMap, + ElementAccumulator, + kSharedMemAlignment + >; + + /// Hard-coded padding elements added + using Padding = typename WarpTileIterator::Padding; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaTensorOp, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for TensorOps. +template < + typename Shape_, + typename WarpMmaTensorOp_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess +> +struct DefaultEpilogueVoltaTensorOpStridedDgrad { + + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapVoltaTensorOp< + Shape, + typename WarpMmaTensorOp::Shape, + kPartitionsK, + ElementOutput, + kElementsPerAccess, + ElementAccumulator + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIteratorStridedDgrad< + OutputTileThreadMap, + ElementOutput + >; + + using AccumulatorFragmentIterator = mctlass::epilogue::warp::FragmentIteratorVoltaTensorOp< + typename WarpMmaTensorOp::Shape, + gemm::GemmShape<32, 32, 4>, + ElementAccumulator, + LayoutC + >; + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorVoltaTensorOp< + typename WarpMmaTensorOp::Shape, + gemm::GemmShape<32, 32, 4>, + ElementAccumulator, + LayoutC + >; + + static int const kSharedMemAlignment = sizeof_bits::value * WarpTileIterator::kElementsPerAccess / 8; + + static_assert(kSharedMemAlignment == 8, "Shared memory alignment must be 8B"); + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + typename OutputTileThreadMap::CompactedThreadMap, + ElementAccumulator, + kSharedMemAlignment + >; + + /// Hard-coded padding elements added + using Padding = typename WarpTileIterator::Padding; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaTensorOp, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for TensorOps. +template < + int Rank, + typename Shape_, + typename WarpMmaTensorOp_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess +> +struct DefaultEpilogueVoltaTensorOpAffineRankN { + + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapVoltaTensorOp< + Shape, + typename WarpMmaTensorOp::Shape, + kPartitionsK, + ElementOutput, + kElementsPerAccess, + ElementAccumulator + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIteratorAffineRankN< + OutputTileThreadMap, + ElementOutput, + Rank + >; + + using AccumulatorFragmentIterator = mctlass::epilogue::warp::FragmentIteratorVoltaTensorOp< + typename WarpMmaTensorOp::Shape, + gemm::GemmShape<32, 32, 4>, + ElementAccumulator, + LayoutC + >; + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorVoltaTensorOp< + typename WarpMmaTensorOp::Shape, + gemm::GemmShape<32, 32, 4>, + ElementAccumulator, + LayoutC + >; + + static int const kSharedMemAlignment = sizeof_bits::value * WarpTileIterator::kElementsPerAccess / 8; + + static_assert(kSharedMemAlignment == 8, "Shared memory alignment must be 8B"); + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + typename OutputTileThreadMap::CompactedThreadMap, + ElementAccumulator, + kSharedMemAlignment + >; + + /// Hard-coded padding elements added + using Padding = typename WarpTileIterator::Padding; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaTensorOp, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_with_broadcast.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_with_broadcast.h new file mode 100644 index 0000000..52382a5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_with_broadcast.h @@ -0,0 +1,245 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/threadblock/epilogue_with_broadcast.h" +#include "mctlass/epilogue/threadblock/epilogue_streamk_with_broadcast.h" + +#include "mctlass/layout/permute.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for TensorOps. +template < + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename ElementOutput, + typename ElementTensor, + typename ElementVector, + typename OutputOp, + int ElementsPerAccess, + bool ScatterD = false, + typename PermuteDLayout = layout::NoPermute +> +struct DefaultEpilogueWithBroadcastTensorOp { + + /// Use defaults related to the existing epilogue + using Base = DefaultEpilogueTensorOp< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputOp, + ElementsPerAccess + >; + + // + // Stores the result z = (y = GEMM(A, B, C), broadcast) + // + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + ElementOutput, + ScatterD, + PermuteDLayout + >; + + // + // Additional tensor tile iterator - stores t = Elementwise(z) + // + using TensorTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + ElementTensor + >; + + /// Define the epilogue + using Epilogue = EpilogueWithBroadcast< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputTileIterator, + TensorTileIterator, + ElementVector, + typename Base::AccumulatorFragmentIterator, + typename Base::WarpTileIterator, + typename Base::SharedLoadIterator, + OutputOp, + typename Base::Padding, + Base::kFragmentsPerIteration + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for streamk epilogues for TensorOps. +template < + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename ElementOutput, + typename ElementTensor, + typename ElementVector, + typename OutputOp, + int ElementsPerAccess, + bool ScatterD = false, + typename PermuteDLayout = layout::NoPermute +> +struct DefaultStreamkEpilogueWithBroadcastTensorOp { + + /// Use defaults related to the existing epilogue + using Base = DefaultEpilogueTensorOp< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputOp, + ElementsPerAccess + >; + + // + // Stores the result z = (y = GEMM(A, B, C), broadcast) + // + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + ElementOutput, + ScatterD, + PermuteDLayout + >; + + // + // Additional tensor tile iterator - stores t = Elementwise(z) + // + using TensorTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + ElementTensor + >; + + /// Define the epilogue + using Epilogue = EpilogueStreamkWithBroadcast< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputTileIterator, + TensorTileIterator, + ElementVector, + typename Base::AccumulatorFragmentIterator, + typename Base::WarpTileIterator, + typename Base::SharedLoadIterator, + OutputOp, + typename Base::Padding, + Base::kFragmentsPerIteration + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for VoltaTensorOps. +template < + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename ElementOutput, + typename ElementTensor, + typename ElementVector, + typename OutputOp, + int ElementsPerAccess +> +struct DefaultEpilogueWithBroadcastVoltaTensorOp { + + /// Use defaults related to the existing epilogue + using Base = DefaultEpilogueVoltaTensorOp< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputOp, + ElementsPerAccess + >; + + // + // Stores the result z = (y = GEMM(A, B, C), broadcast) + // + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + ElementOutput + >; + + // + // Additional tensor tile iterator - stores t = Elementwise(z) + // + using TensorTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + ElementTensor + >; + + /// Define the epilogue + using Epilogue = EpilogueWithBroadcast< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputTileIterator, + TensorTileIterator, + ElementVector, + typename Base::AccumulatorFragmentIterator, + typename Base::WarpTileIterator, + typename Base::SharedLoadIterator, + OutputOp, + typename Base::Padding + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_with_reduction.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_with_reduction.h new file mode 100644 index 0000000..1c4fc7b --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_with_reduction.h @@ -0,0 +1,177 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/threadblock/epilogue_with_reduction.h" + +#include "mctlass/layout/permute.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for TensorOps. +template < + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename ElementOutput, + typename OutputOp, + typename ReductionOp, + int ElementsPerAccess, + bool ScatterD = false, + typename PermuteDLayout = layout::NoPermute +> +struct DefaultEpilogueWithReductionTensorOp { + + /// Use defaults related to the existing epilogue + using Base = DefaultEpilogueTensorOp< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputOp, + ElementsPerAccess + >; + + /// Additional tensor tile iterator + using TensorTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + typename OutputOp::ElementTensor + >; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + ElementOutput, + ScatterD, + PermuteDLayout + >; + + /// Define the epilogue + using Epilogue = EpilogueWithReduction< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputTileIterator, + TensorTileIterator, + typename WarpMmaTensorOp::ElementC, + typename Base::AccumulatorFragmentIterator, + typename Base::WarpTileIterator, + typename Base::SharedLoadIterator, + typename Base::OutputOp, + ReductionOp, + typename Base::Padding + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for TensorOps. +template < + typename Shape, + typename WarpMmaTensorOp, + int PartitionsK, + typename ElementOutput, + typename OutputOp, + typename ReductionOp, + int ElementsPerAccess, + bool ScatterD = false, + typename PermuteDLayout = layout::NoPermute +> +struct DefaultEpilogueWithReductionVoltaTensorOp { + + /// Use defaults related to the existing epilogue + using Base = DefaultEpilogueVoltaTensorOp< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputOp, + ElementsPerAccess + >; + + /// Additional tensor tile iterator + using TensorTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + typename OutputOp::ElementTensor + >; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + typename Base::OutputTileThreadMap, + ElementOutput, + ScatterD, + PermuteDLayout + >; + + /// Define the epilogue + using Epilogue = EpilogueWithReduction< + Shape, + WarpMmaTensorOp, + PartitionsK, + OutputTileIterator, + TensorTileIterator, + typename WarpMmaTensorOp::ElementC, + typename Base::AccumulatorFragmentIterator, + typename Base::WarpTileIterator, + typename Base::SharedLoadIterator, + typename Base::OutputOp, + ReductionOp, + typename Base::Padding + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h new file mode 100644 index 0000000..133442a --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h @@ -0,0 +1,165 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using WMMA. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/epilogue/thread/linear_combination.h" +#include "mctlass/epilogue/thread/linear_combination_clamp.h" +#include "mctlass/epilogue/thread/linear_combination_relu.h" +#include "mctlass/epilogue/thread/linear_combination_gelu.h" +#include "mctlass/epilogue/thread/linear_combination_sigmoid.h" +#include "mctlass/epilogue/thread/linear_combination_planar_complex.h" + +#include "mctlass/epilogue/thread/conversion_op.h" +#include "mctlass/epilogue/thread/reduction_op.h" + +#include "mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h" + +#include "mctlass/epilogue/warp/fragment_iterator_wmma_tensor_op.h" +#include "mctlass/epilogue/warp/tile_iterator_wmma_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_thread_map_wmma_tensor_op.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator.h" +#include "mctlass/epilogue/threadblock/shared_load_iterator.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" + +#include "mctlass/layout/permute.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines sensible defaults for epilogues for WMMA TensorOps. +template < + typename Shape_, + typename WarpMmaTensorOp_, + int PartitionsK, + typename OutputOp_, + int ElementsPerAccess, + bool ScatterD = false, + typename PermuteDLayout = layout::NoPermute +> +struct DefaultEpilogueWmmaTensorOp { + + using Shape = Shape_; + using WarpMmaTensorOp = WarpMmaTensorOp_; + static int const kPartitionsK = PartitionsK; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = ElementsPerAccess; + + using ElementOutput = typename OutputOp::ElementOutput; + using LayoutC = typename WarpMmaTensorOp::LayoutC; + using ElementAccumulator = typename WarpMmaTensorOp::ElementC; + + // + // Thread map + // + + using OutputTileThreadMap = typename mctlass::epilogue::threadblock::DefaultThreadMapWmmaTensorOp< + Shape, + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + kPartitionsK, + ElementOutput, + kElementsPerAccess + >::Type; + + using OutputTileIterator = mctlass::epilogue::threadblock::PredicatedTileIterator< + OutputTileThreadMap, + ElementOutput, + ScatterD, + PermuteDLayout + >; + + using AccumulatorFragmentIterator = mctlass::epilogue::warp::FragmentIteratorWmmaTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::ElementC, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC + >; + + using WarpTileIterator = mctlass::epilogue::warp::TileIteratorWmmaTensorOp< + typename WarpMmaTensorOp::Shape, + typename WarpMmaTensorOp::Policy::Operator::Shape, + typename WarpMmaTensorOp::Policy::Operator::FragmentC, + LayoutC + >; + + using SharedLoadIterator = mctlass::epilogue::threadblock::SharedLoadIterator< + typename OutputTileThreadMap::CompactedThreadMap, + ElementAccumulator + >; + + /// Hard-coded padding elements added + using Padding = typename WarpTileIterator::Padding; + + // + // Define the epilogue + // + using Epilogue = mctlass::epilogue::threadblock::Epilogue< + Shape, + WarpMmaTensorOp, + kPartitionsK, + OutputTileIterator, + AccumulatorFragmentIterator, + WarpTileIterator, + SharedLoadIterator, + OutputOp, + Padding + >; +}; + + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_simt.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_simt.h new file mode 100644 index 0000000..883101f --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_simt.h @@ -0,0 +1,127 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief + +*/ + +#pragma once + +#include "predicated_tile_iterator.h" +#include "mctlass/gemm/gemm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines the optimal thread map for SIMT accumulator layouts +template < + typename ThreadblockShape_, + typename WarpShape_, + typename MmaSimtPolicy_, + int PartitionsK, + typename Element_, + int ElementsPerAccess +> +struct DefaultThreadMapSimt { + + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using MmaSimtPolicy = MmaSimtPolicy_; + static int const kPartitionsK = PartitionsK; + using Element = Element_; + static int const kElementsPerAccess = ElementsPerAccess; + + // + // Definitions + // + + struct Detail { + + static int const kWarpSize = 32; + + static_assert( + !(ThreadblockShape::kM % WarpShape::kM) && + !(ThreadblockShape::kN % WarpShape::kN), "Divisibility"); + + /// Number of warps + using WarpCount = gemm::GemmShape< + ThreadblockShape::kM / WarpShape::kM, + ThreadblockShape::kN / WarpShape::kN, + kPartitionsK + >; + + /// Computes number of thread-level matrix multiplies are needed to span a warp + static int const kGroupCount = + WarpShape::kM / (MmaSimtPolicy::WarpShape::kRow * MmaSimtPolicy::LaneMmaShape::kM); + + /// Number of participating threads + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Number of iterations + static int const kIterations = MmaSimtPolicy::LaneMmaShape::kM * kGroupCount; + }; + + // + // ThreadMap + // + + /// ThreadMap to be used by epilogue::PredicatedTileIterator satisfying concept OutputTileThreadMap + using Type = OutputTileOptimalThreadMap< + OutputTileShape< // Shape + ThreadblockShape::kN, + 1, + MmaSimtPolicy::WarpShape::kRow, + Detail::WarpCount::kM, + 1>, + OutputTileShape< // Count + 1, + MmaSimtPolicy::LaneMmaShape::kM, + Detail::kGroupCount, + 1, + Detail::kIterations>, + Detail::kThreads, + kElementsPerAccess, + sizeof_bits::value + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_tensor_op.h new file mode 100644 index 0000000..0f47dcf --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_tensor_op.h @@ -0,0 +1,208 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief + +*/ + +#pragma once + +#include "predicated_tile_iterator.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/layout/pitch_linear.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines the optimal thread map for TensorOp accumulator layouts +template < + typename ThreadblockShape_, + typename WarpShape_, + int PartitionsK, + typename Element_, + int ElementsPerAccess +> +struct DefaultThreadMapTensorOp { + + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + static int const kPartitionsK = PartitionsK; + using Element = Element_; + static int const kElementsPerAccess = ElementsPerAccess; + + // + // Definitions + // + + struct Detail { + + /// Tensor Operations fundamentally perform operations on 8 rows + static int const kTensorOpRows = 8; + static int const kWarpSize = 32; + + static_assert( + !(ThreadblockShape::kM % WarpShape::kM) && + !(ThreadblockShape::kN % WarpShape::kN), "Divisibility"); + + /// Number of warps + using WarpCount = gemm::GemmShape< + ThreadblockShape::kM / WarpShape::kM, + ThreadblockShape::kN / WarpShape::kN, + kPartitionsK + >; + + /// Number of participating threads + static int const kThreads = WarpCount::kCount * kWarpSize; + }; + + // + // ThreadMap + // + + /// ThreadMap to be used by epilogue::PredicatedTileIterator satisfying concept OutputTileThreadMap + using Type = OutputTileOptimalThreadMap < + OutputTileShape, + OutputTileShape<1, WarpShape::kM / Detail::kTensorOpRows, 1, 1, WarpShape::kM / Detail::kTensorOpRows>, + Detail::kThreads, + kElementsPerAccess, + sizeof_bits::value + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines the optimal thread map for TensorOp accumulator layouts +template +struct DefaultInterleavedThreadMapTensorOp { + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + static int const kPartitionsK = PartitionsK; + using Element = Element_; + static int const kElementsPerAccess = ElementsPerAccess; + static int const kInterleavedK = InterleavedK; + + // + // Definitions + // + + struct Detail { + /// Tensor Operations fundamentally perform operations on 8 rows + static int const kTensorOpRows = 8; + static int const kWarpSize = 64; + + static_assert(!(ThreadblockShape::kM % WarpShape::kM) && + !(ThreadblockShape::kN % WarpShape::kN), + "Divisibility"); + + /// Number of warps + using WarpCount = + gemm::GemmShape; + + /// Number of participating threads + static int const kThreads = WarpCount::kCount * kWarpSize; + }; + + // + // ThreadMap + // + + /// ThreadMap to be used by epilogue::PredicatedTileIterator satisfying concept + /// InterleavedOutputTileThreadMap + using Type = InterleavedOutputTileThreadMap< + layout::PitchLinearShape, + layout::PitchLinearShape, + Detail::kThreads, kElementsPerAccess, sizeof_bits::value>; +}; + + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines the optimal thread map for TensorOp accumulator layouts +template +struct DefaultInterleavedConvThreadMapTensorOp { + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + static int const kPartitionsK = PartitionsK; + using Element = Element_; + static int const kElementsPerAccess = ElementsPerAccess; + static int const kInterleavedK = InterleavedK; + + // + // Definitions + // + + struct Detail { + /// Tensor Operations fundamentally perform operations on 8 rows + static int const kTensorOpRows = 8; + static int const kWarpSize = 64; + + static_assert(!(ThreadblockShape::kM % WarpShape::kM) && + !(ThreadblockShape::kN % WarpShape::kN), + "Divisibility"); + + /// Number of warps + using WarpCount = + gemm::GemmShape; + + /// Number of participating threads + static int const kThreads = WarpCount::kCount * kWarpSize; + }; + + // + // ThreadMap + // + + /// ThreadMap to be used by epilogue::MaskedTileIterator satisfying concept + /// InterleavedOutputTileThreadMap + using Type = InterleavedConvOutputTileThreadMap< + MatrixShape, + MatrixShape, + Detail::kThreads, kElementsPerAccess, sizeof_bits::value>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_volta_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_volta_tensor_op.h new file mode 100644 index 0000000..c6aabfd --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_volta_tensor_op.h @@ -0,0 +1,228 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief + +*/ + +#pragma once + +#include "predicated_tile_iterator.h" +#include "mctlass/gemm/gemm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines the optimal thread map for TensorOp accumulator layouts +template < + typename ThreadblockShape, + typename WarpShape, + int PartitionsK, + typename ElementOutput, + int ElementsPerAccess, + typename ElementAccumulator +> +struct DefaultThreadMapVoltaTensorOp; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines the optimal thread map for TensorOp accumulator layouts +template < + typename ThreadblockShape_, + typename WarpShape_, + int PartitionsK, + typename ElementOutput_, + int ElementsPerAccess +> +struct DefaultThreadMapVoltaTensorOp< + ThreadblockShape_, + WarpShape_, + PartitionsK, + ElementOutput_, + ElementsPerAccess, + half_t> { + + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + static int const kPartitionsK = PartitionsK; + using ElementOutput = ElementOutput_; + static int const kElementsPerAccess = ElementsPerAccess; + using ElementAccumulator = half_t; + + // + // Definitions + // + + struct Detail { + + static int const kTensorOpRows = 16; + static int const kWarpSize = 32; + static int const kInterleavedTilesM = WarpShape::kM / 32; + + static_assert( + !(ThreadblockShape::kM % WarpShape::kM) && + !(ThreadblockShape::kN % WarpShape::kN), "Divisibility"); + + /// Number of warps + using WarpCount = gemm::GemmShape< + ThreadblockShape::kM / WarpShape::kM, + ThreadblockShape::kN / WarpShape::kN, + kPartitionsK + >; + + /// Number of participating threads + static int const kThreads = WarpCount::kCount * kWarpSize; + + using Shape = mctlass::epilogue::threadblock::OutputTileShape< + ThreadblockShape::kN, // column + 4, // row + 4, // group + WarpCount::kM, // cluster + 1 // tile + >; + + /// Number of iterations per subspace + using Count = mctlass::epilogue::threadblock::OutputTileShape< + 1, // column + 2, // row + kInterleavedTilesM, // group + 1, // cluster + WarpShape::kM / kTensorOpRows // iterations + >; + }; + + // + // ThreadMap + // + + /// ThreadMap to be used by epilogue::PredicatedTileIterator satisfying concept OutputTileThreadMap + using Type = OutputTileOptimalThreadMap < + typename Detail::Shape, + typename Detail::Count, + Detail::kThreads, + kElementsPerAccess, + sizeof_bits::value + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines the optimal thread map for TensorOp accumulator layouts +template < + typename ThreadblockShape_, + typename WarpShape_, + int PartitionsK, + typename ElementOutput_, + int ElementsPerAccess +> +struct DefaultThreadMapVoltaTensorOp< + ThreadblockShape_, + WarpShape_, + PartitionsK, + ElementOutput_, + ElementsPerAccess, + float> { + + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + static int const kPartitionsK = PartitionsK; + using ElementOutput = ElementOutput_; + static int const kElementsPerAccess = ElementsPerAccess; + using ElementAccumulator = float; + + // + // Definitions + // + + struct Detail { + + static int const kTensorOpRows = 16; + static int const kWarpSize = 32; + static int const kInterleavedTilesM = WarpShape::kM / 32; + + static_assert( + !(ThreadblockShape::kM % WarpShape::kM) && + !(ThreadblockShape::kN % WarpShape::kN), "Divisibility"); + + /// Number of warps + using WarpCount = gemm::GemmShape< + ThreadblockShape::kM / WarpShape::kM, + ThreadblockShape::kN / WarpShape::kN, + kPartitionsK + >; + + /// Number of participating threads + static int const kThreads = WarpCount::kCount * kWarpSize; + + using Shape = mctlass::epilogue::threadblock::OutputTileShape< + ThreadblockShape::kN, // column + 4, // row + 4, // group + WarpCount::kM, // cluster + 1 // tile + >; + + /// Number of iterations per subspace + using Count = mctlass::epilogue::threadblock::OutputTileShape< + 1, // column + 2, // row + kInterleavedTilesM, // group + 1, // cluster + WarpShape::kM / kTensorOpRows // iterations + >; + }; + + // + // ThreadMap + // + + /// ThreadMap to be used by epilogue::PredicatedTileIterator satisfying concept OutputTileThreadMap + using Type = OutputTileOptimalThreadMap < + typename Detail::Shape, + typename Detail::Count, + Detail::kThreads, + kElementsPerAccess, + sizeof_bits::value + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_wmma_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_wmma_tensor_op.h new file mode 100644 index 0000000..917e952 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/default_thread_map_wmma_tensor_op.h @@ -0,0 +1,113 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief + +*/ + +#pragma once + +#include "predicated_tile_iterator.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/layout/pitch_linear.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Defines the optimal thread map for Wmma TensorOp accumulator layouts +template < + typename ThreadblockShape_, + typename WarpShape_, + typename InstructionShape_, + int PartitionsK, + typename Element_, + int ElementsPerAccess +> +struct DefaultThreadMapWmmaTensorOp { + + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + static int const kPartitionsK = PartitionsK; + using Element = Element_; + static int const kElementsPerAccess = ElementsPerAccess; + + // + // Definitions + // + + struct Detail { + + /// Wmma Tensor Operations fundamentally perform operations on InstructionShape::kM rows + static int const kTensorOpRows = InstructionShape::kM; + static int const kWarpSize = 32; + + static_assert( + !(ThreadblockShape::kM % WarpShape::kM) && + !(ThreadblockShape::kN % WarpShape::kN), "Divisibility"); + + /// Number of warps + using WarpCount = gemm::GemmShape< + ThreadblockShape::kM / WarpShape::kM, + ThreadblockShape::kN / WarpShape::kN, + kPartitionsK + >; + + /// Number of participating threads + static int const kThreads = WarpCount::kCount * kWarpSize; + }; + + // + // ThreadMap + // + + /// ThreadMap to be used by epilogue::PredicatedTileIterator satisfying concept OutputTileThreadMap + using Type = OutputTileOptimalThreadMap < + OutputTileShape, + OutputTileShape<1, WarpShape::kM / Detail::kTensorOpRows, 1, 1, WarpShape::kM / Detail::kTensorOpRows>, + Detail::kThreads, + kElementsPerAccess, + sizeof_bits::value + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/direct_store_epilogue_iterator.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/direct_store_epilogue_iterator.h new file mode 100644 index 0000000..8420839 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/direct_store_epilogue_iterator.h @@ -0,0 +1,142 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/epilogue/threadblock/output_tile_thread_map.h" +#include "mctlass/arch/arch.h" +#include "mctlass/arch/memory.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////// + +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template +class DirectStoreEpilogueIterator { +public: + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = 1; + + /// Uses a non-template class + struct Params : PredicatedTileIteratorParams { + using Base = PredicatedTileIteratorParams; + + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Layout const &layout) { + stride = layout.stride(0) * sizeof(Element); + } + + MCTLASS_HOST_DEVICE + Params(Base const &base) : + Base(base) { } + }; + +public: + + // + // Data members + // + + Element *pointer; // pointer to the output matrix + + LongIndex stride; // stride in elements between rows + + TensorCoord extent; // extent of output matrix + + int thread_idx; // thread index + + TensorCoord threadblock_offset; + +public: + + /// Constructor + MCTLASS_DEVICE + DirectStoreEpilogueIterator( + PredicatedTileIteratorParams const & params, + Element *pointer_, + TensorCoord extent_, + int thread_idx_, + TensorCoord threadblock_offset_ = TensorCoord(), + int const * indices = nullptr + ): + pointer(pointer_), + stride(params.stride / sizeof(Element)), + extent(extent_), + thread_idx(thread_idx_), + threadblock_offset(threadblock_offset_) + { + + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue.h new file mode 100644 index 0000000..961f74a --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue.h @@ -0,0 +1,652 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + + The shared memory resource is time-sliced across warps. +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/vector.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/tensor_coord.h" +#include "mctlass/aligned_buffer.h" +#include "mctlass/functional.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/arch/mma.h" + +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_iterator.h" +#include "mctlass/epilogue/threadblock/epilogue_base.h" +#include "mctlass/epilogue/threadblock/epilogue_base_streamk.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// +template +MCTLASS_DEVICE +T shfl_down_sync(unsigned long mask, T var, unsigned int laneDelta, int width = 64) { + return var; +} + +template +MCTLASS_DEVICE +T shfl_up_sync(unsigned long mask, T var, unsigned int laneDelta, int width = 64) { + return var; +} + +MCTLASS_DEVICE +int32_t shfl_down_sync(unsigned long mask, int32_t var, unsigned int laneDelta, int width = 64) { + return __shfl_down_sync(mask, var, laneDelta, width); +} + +MCTLASS_DEVICE +int32_t shfl_up_sync(unsigned long mask, int32_t var, unsigned int laneDelta, int width = 64) { + return __shfl_up_sync(mask, var, laneDelta, width); +} + +MCTLASS_DEVICE +float shfl_down_sync(unsigned long mask, float var, unsigned int laneDelta, int width = 64) { + return __shfl_down_sync(mask, var, laneDelta, width); +} + +MCTLASS_DEVICE +float shfl_up_sync(unsigned long mask, float var, unsigned int laneDelta, int width = 64) { + return __shfl_up_sync(mask, var, laneDelta, width); +} + +MCTLASS_DEVICE +mctlass::half_t shfl_down_sync(unsigned long mask, mctlass::half_t var, unsigned int laneDelta, int width = 64) { + auto temp_value = __shfl_down_sync(mask, var, laneDelta, width); + return mctlass::half_t(temp_value); +} + +MCTLASS_DEVICE +mctlass::half_t shfl_up_sync(unsigned long mask, mctlass::half_t var, unsigned int laneDelta, int width = 64) { + auto temp_value = __shfl_up_sync(mask, var, laneDelta, width); + return mctlass::half_t(temp_value); +} + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator +template < + typename Shape_, ///< Shape of threadblock tile (concept: GemmShape) + typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp) + int PartitionsK, ///< Number of partitions of the K dimension + typename OutputTileIterator_, ///< Tile iterator reading and writing output tensors + typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators + typename WarpTileIterator_, ///< Warp-scoped tile iterator writing accumulators to SMEM + typename SharedLoadIterator_, ///< Threadblock-scoped tile iterator loading from SMEM + typename OutputOp_, ///< Output operator + typename Padding_, ///< Padding added to SMEM allocation to avoid bank conflicts (concept: MatrixShape) + int FragmentsPerPartition = 1, ///< Used to coarsten the epilogue granularity + int IterationsUnroll = ///< Used to reduce binary size when epilogue op is large + (!IsEpilogueFunctorHeavy::value) +> +class Epilogue : + public EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_, + FragmentsPerPartition>, + public EpilogueBaseStreamK< + Shape_, + PartitionsK, + WarpMmaOperator_, + AccumulatorFragmentIterator_> +{ + +public: + + using Base = EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_, + FragmentsPerPartition>; + + using BaseStreamK = EpilogueBaseStreamK< + Shape_, + PartitionsK, + WarpMmaOperator_, + AccumulatorFragmentIterator_>; + + using Shape = Shape_; + using WarpMmaOperator = WarpMmaOperator_; + static int const kPartitionsK = PartitionsK; + using OutputTileIterator = OutputTileIterator_; + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using WarpTileIterator = WarpTileIterator_; + using SharedLoadIterator = SharedLoadIterator_; + using OutputOp = OutputOp_; + using Padding = Padding_; + using Layout = layout::RowMajor; + using LongIndex = typename Layout::LongIndex; + + /// Number of warps per block + using WarpCount = typename Base::WarpCount; + + /// Number of threads per block + static int const kBlockThreads = 64 * WarpCount::kCount; + + /// Per-thread accumulator tile type + using AccumulatorTile = typename Base::AccumulatorTile; + + /// Numerical accumulation element type + using ElementAccumulator = typename WarpMmaOperator::ElementC; + + /// Fragment type used by the accumulator tile's fragment iterator + using AccumulatorFragment = typename AccumulatorFragmentIterator::Fragment; + + /// Output element + using ElementOutput = typename OutputTileIterator::Element; + + /// Output access size + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + /// Tensor reference to destination tensor + using TensorRef = typename OutputTileIterator::TensorRef; + + /// Tensor reference to sync tensor + using SyncTensorRef = typename mctlass::TensorRef; + + /// Const tensor reference to source tensor + using ConstTensorRef = typename OutputTileIterator::ConstTensorRef; + + /// Vector type used by the global output iterator + using OutputAccessType = Array< + typename OutputTileIterator::Element, OutputTileIterator::kElementsPerAccess>; + + /// Vector type used by the shared output iterator + using AccumulatorAccessType = Array; + + static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1 ? Base::kFragmentsPerIteration : kPartitionsK; + static int constexpr kSmemPointerOffset = Base::SharedStorage::StorageShape::kCount / kSmemTiles; + + +public: + + static_assert(SharedLoadIterator::Fragment::kElements == OutputTileIterator::Fragment::kElements, + "Mismatch between shared load iterator and output tile iterator."); + + static_assert(OutputTileIterator::kElementsPerAccess, "OutputTileIterator::kElementsPerAccess must not be zero."); + + static_assert(!(OutputTileIterator::Fragment::kElements % OutputTileIterator::kElementsPerAccess), + "Divisibility"); + + static_assert(kPartitionsK == 1 || Base::kFragmentsPerIteration == 1, "One of these must be exactly 1."); + +public: + + /// Aspect for when epilogue source is not needed + struct SourceAspectNotNeeded + { + /// Constructor + MCTLASS_DEVICE + SourceAspectNotNeeded() + {} + + /// Invoke the output functor over each vector of output + MCTLASS_DEVICE + void apply_output_operator( + typename OutputTileIterator::Fragment &output_fragment, + OutputOp const &output_op, + typename SharedLoadIterator::Fragment const &aligned_accum_fragment) + { + OutputAccessType *output_frag_ptr = + reinterpret_cast(&output_fragment); + + AccumulatorAccessType const *compute_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) + { + // Call the output operator + output_frag_ptr[i] = output_op(compute_frag_ptr[i]); + } + } + }; + + + /// Aspect for when epilogue source is needed + struct SourceAspectNeeded + { + OutputTileIterator source_iterator; + + typename OutputTileIterator::Fragment source_fragment; + + /// Invoke the output functor over each vector of output + MCTLASS_DEVICE + static void apply_output_operator( + typename OutputTileIterator::Fragment &output_fragment, + OutputOp const &output_op, + typename SharedLoadIterator::Fragment const &aligned_accum_fragment, + typename OutputTileIterator::Fragment const &source_fragment) + { + OutputAccessType *output_frag_ptr = + reinterpret_cast(&output_fragment); + + AccumulatorAccessType const *compute_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + OutputAccessType const *source_frag_ptr = + reinterpret_cast(&source_fragment); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) + { + // Call the output operator + output_frag_ptr[i] = output_op(compute_frag_ptr[i], source_frag_ptr[i]); + } + } + + /// Constructor + MCTLASS_DEVICE + SourceAspectNeeded(OutputTileIterator source_iterator) : + source_iterator(source_iterator) + { + source_fragment.clear(); + } + + /// Invoke the output functor over each vector of output + MCTLASS_DEVICE + void apply_output_operator( + typename OutputTileIterator::Fragment &output_fragment, + OutputOp const &output_op, + typename SharedLoadIterator::Fragment const &aligned_accum_fragment) + { + // Load addend source fragment from global memory + source_iterator.load(source_fragment); + ++source_iterator; + + apply_output_operator(output_fragment, output_op, aligned_accum_fragment, source_fragment); + } + }; + +private: + + /// Loads fragment from shared memory aligned with output tensor + SharedLoadIterator shared_load_iterator_; + + /// Thread index in the threadblock + int thread_idx; + + /// Warp index in the threadblock + int warp_idx; + +public: + + /// Constructor + MCTLASS_DEVICE + Epilogue( + typename Base::SharedStorage &shared_storage, ///< Shared storage object + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx) ///< Id of thread within warp + : + Base(shared_storage, thread_idx, warp_idx, lane_idx), + BaseStreamK(thread_idx), + shared_load_iterator_(shared_storage.reference(), thread_idx), + thread_idx(thread_idx), + warp_idx(warp_idx) + {} + + /// Aggregates the accumulator sets shared by peer blocks in the global workspace, + /// performing epilogue computations, writing to output + MCTLASS_DEVICE + void reduce( + int peer_idx_begin, + int peer_idx_end, + int reduce_fragment_idx, + void *element_workspace, + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + OutputTileIterator source_iterator) ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + { + // Reduce peer accumulator fragments into one fragment + AccumulatorFragment accum_fragment; + BaseStreamK::reduce(accum_fragment, peer_idx_begin, peer_idx_end, reduce_fragment_idx, element_workspace); + + // Store fragment to shared memory + this->warp_tile_iterator_.store(accum_fragment); + + __syncthreads(); + + // Initialize/load source-fragment data + typename OutputTileIterator::Fragment source_fragment; + source_fragment.clear(); + + if (output_op.is_source_needed()) + { + source_iterator += reduce_fragment_idx; + source_iterator.load(source_fragment); + } + + // Load fragment from shared memory + typename SharedLoadIterator::Fragment aligned_accum_fragment; + shared_load_iterator_.load(aligned_accum_fragment); + + // Add fragments shared by other k partitions + if (kPartitionsK > 1) + { + plus add_fragments; + + MCTLASS_PRAGMA_UNROLL + for ( int i = 1; i < kPartitionsK; ++i) { + typename SharedLoadIterator::Fragment aligned_addend_fragment; + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + shared_load_iterator_.load(aligned_addend_fragment); + aligned_accum_fragment = add_fragments(aligned_accum_fragment, aligned_addend_fragment); + } + } + + // Compute the output result + typename OutputTileIterator::Fragment output_fragment; + + // Apply the output operator + SourceAspectNeeded::apply_output_operator( + output_fragment, + output_op, + aligned_accum_fragment, + source_fragment); + + // Store the final result + destination_iterator += reduce_fragment_idx; + destination_iterator.store(output_fragment); + } + + /// Perform the epilogue computations and stream the result to global memory. + MCTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators) ///< Complete warp-level accumulator tile + { + operator()(output_op, destination_iterator, accumulators, SourceAspectNotNeeded()); + } + + /// Streams the result to global memory + MCTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + + // Special process for m16n8k32 int32_t cases + // For type mctlass::complex,when in AccumulatorTile,the real part and the imag part are stored separately. + if ((mctlass::platform::is_same::value == true && WarpTileIterator::WarpShape::kK >= 64) || + (mctlass::platform::is_same>::value && mctlass::platform::is_same::value) || + (mctlass::platform::is_same::value && (mctlass::platform::is_same::value || mctlass::platform::is_same::value) && + mctlass::platform::is_same::value && + WarpTileIterator::WarpShape::kN >= 64) || + (mctlass::platform::is_same::value && mctlass::platform::is_same::value)) { + using Element = typename AccumulatorTile::Element; + for (int i = 0; i < accumulators.size() / 2; i += 2) { + + Element x_real = accumulators[2 * i + 0]; + Element y_real = accumulators[2 * i + 1]; + Element z_real = accumulators[2 * (i + 1) + 0]; + Element w_real = accumulators[2 * (i + 1) + 1]; + + Element x0 = shfl_down_sync(UINT64_MAX, x_real, 32); + Element y0 = shfl_down_sync(UINT64_MAX, y_real, 32); + Element z0 = shfl_up_sync(UINT64_MAX, z_real, 32); + Element w0 = shfl_up_sync(UINT64_MAX, w_real, 32); + if (__lane_id() < 32) { + + accumulators[2 * (i + 1) + 0] = x0; + accumulators[2 * (i + 1) + 1] = y0; + } + else { + accumulators[2 * i + 0] = z0; + accumulators[2 * i + 1] = w0; + } + } + } + + if (output_op.is_source_needed()) { + operator()(output_op, destination_iterator, accumulators, SourceAspectNeeded(source_iterator)); + } + else { + operator()(output_op, destination_iterator, accumulators, SourceAspectNotNeeded()); + } + } + + /// Perform the epilogue computations and stream the result to global memory. Implements + /// two alternative codepaths, depending on whether the output op requires addend data to be loaded. + // MCTLASS_DEVICE + // void operator()( + // OutputOp const &output_op, ///< Output operator + // OutputTileIterator destination_iterator, ///< Tile iterator for destination + // AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + // OutputTileIterator source_iterator ) ///< Tile iterator for addend source + // { + // if (output_op.is_source_needed()) + // { + // operator()(output_op, destination_iterator, accumulators, SourceAspectNeeded(source_iterator)); + // } + // else + // { + // operator()(output_op, destination_iterator, accumulators, SourceAspectNotNeeded()); + // } + // } + + + /// Perform the epilogue computations and stream the result to global memory. Implements a + /// single codepath, regardless of whether the output op requires addend data to be loaded + MCTLASS_DEVICE + void unified( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator ) ///< Tile iterator for addend source + { + if (!output_op.is_source_needed()) + { + source_iterator.clear_mask(); + __syncthreads(); // Dummy (CUDA 11.0) + } + + operator()(output_op, destination_iterator, accumulators, SourceAspectNeeded(source_iterator)); + } + +private: + + template + struct acc2smem; + + template + struct acc2smem> { + template + MCTLASS_DEVICE static void helper(AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator &warp_tile_iterator) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { + ++accum_fragment_iterator; + } + + MCTLASS_PRAGMA_UNROLL + for (int p = 0; p < Base::kFragmentsPerIteration; ++p) { + typename AccumulatorFragmentIterator::Fragment accum_fragment; + + accum_fragment_iterator.load(accum_fragment); + ++accum_fragment_iterator; + + warp_tile_iterator.store(accum_fragment); + if (p < Base::kFragmentsPerIteration - 1) { + warp_tile_iterator.add_pointer_offset(kSmemPointerOffset); + } + } + + if (Base::kFragmentsPerIteration > 1) { + warp_tile_iterator.add_pointer_offset(kSmemPointerOffset * + (1 - Base::kFragmentsPerIteration)); + } + } + + MCTLASS_DEVICE + static void push(size_t pos, + AccumulatorFragmentIterator const &iterator_begin, + WarpTileIterator &warp_tile_iterator) { + int dummy[] = {(pos == Seq) && (helper(iterator_begin, warp_tile_iterator), 0)...}; + } + }; + + /// Streams the result to global memory + template + MCTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + SourceAspect source) + { + + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + + //#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations / Base::kFragmentsPerIteration : 1) + //mxcc not support using unroll with parentheses, + //and we would support with unroll with OutputTileIterator::kIterations / Base::kFragmentsPerIteration someday + #pragma unroll 1 + for (int iter = 0; iter < OutputTileIterator::kIterations; iter += Base::kFragmentsPerIteration) { + + // + // Convert and store fragment + // + + __syncthreads(); + + + acc2smem>::push( + iter, accum_fragment_iterator, this->warp_tile_iterator_); + + + // + // Load fragments from shared memory + // + + __syncthreads(); + + MCTLASS_PRAGMA_UNROLL + for (int p = 0; p < Base::kFragmentsPerIteration; ++p) + { + typename SharedLoadIterator::Fragment aligned_accum_fragment; + shared_load_iterator_.load(aligned_accum_fragment); + + if (p < Base::kFragmentsPerIteration - 1) + { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + } + else if (kPartitionsK > 1) + { + plus add_fragments; + + MCTLASS_PRAGMA_UNROLL + for ( int i = 1; i < kPartitionsK; ++i) { + typename SharedLoadIterator::Fragment aligned_accum_fragment_addend; + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + shared_load_iterator_.load(aligned_accum_fragment_addend); + aligned_accum_fragment = add_fragments(aligned_accum_fragment, aligned_accum_fragment_addend); + } + + shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * kSmemPointerOffset); + } + + // + // Compute the output result + // + + typename OutputTileIterator::Fragment output_fragment; + + source.apply_output_operator(output_fragment, output_op, aligned_accum_fragment); + + + // + // Store the final result + // + + destination_iterator.store(output_fragment); + ++destination_iterator; + } + + if (Base::kFragmentsPerIteration > 1) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset * (1 - Base::kFragmentsPerIteration)); + } + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_base.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_base.h new file mode 100644 index 0000000..f8361ee --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_base.h @@ -0,0 +1,240 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#if !defined(__MACACC_RTC__) +#include +#include +#endif + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/vector.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/tensor_coord.h" +#include "mctlass/aligned_buffer.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/transform/pitch_linear_thread_map.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +// +// This is used for metaprogramming epilogue functors. If they define +// `static bool const kIsHeavy = true;`, then the epilogue functor itself is +// not inlined. This results in smaller code and is advantageous if the epilogue +// functor consists of many instructions. +// +// If the epilogue functor does not define `kIsHeavy` or if it is `false`, then +// the behavior from MCTLASS 2.5 and before is retained. The epilogue is fully +// unrolled and inlined. +// + +template +struct TypeSink { typedef void type; }; + +template using TypeSinkT = typename TypeSink::type; + +template struct IsEpilogueFunctorHeavy { + static bool const value = false; +}; + +template struct IsEpilogueFunctorHeavy > { + static bool const value = T::kIsHeavy; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Base class for epilogues defining warp-level +template < + typename Shape_, ///< Shape of threadblock tile (concept: GemmShape) + typename WarpShape_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp) + int PartitionsK, ///< Number of partitions of the K dimension + typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators + typename WarpTileIterator_, ///< Warp-scoped tile iterator writing accumulators to SMEM + typename Padding_, ///< Padding added to SMEM allocation to avoid bank conflicts (concept: MatrixShape) + int FragmentsPerIteration = 1 +> +class EpilogueBase { +public: + + using Shape = Shape_; + using WarpShape = WarpShape_; + static int const kPartitionsK = PartitionsK; + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using WarpTileIterator = WarpTileIterator_; + using Padding = Padding_; + + /// Output layout is always row-major + using Layout = layout::RowMajor; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename AccumulatorFragmentIterator::AccumulatorTile; + + /// Accumulator element + using ElementAccumulator = typename AccumulatorTile::Element; + + /// Number of warps + using WarpCount = gemm::GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + kPartitionsK + >; + + /// Use this to control the granularity of one epilogue 'iteration' + static int const kFragmentsPerIteration = FragmentsPerIteration; + +public: + + /// Shared storage allocation needed by the epilogue + struct SharedStorage { + + // + // Type definitions + // + + /// Element type of shared memory + using Element = typename WarpTileIterator::Element; + + /// Tensor reference to shared memory allocation + using TensorRef = typename WarpTileIterator::TensorRef; + + /// Layout of shared memory allocation + using Layout = typename WarpTileIterator::Layout; + + /// Logical shape of the shared memory tile written to by all warps. + using Shape = MatrixShape< + WarpCount::kM * WarpTileIterator::Shape::kRow * WarpCount::kK, + WarpCount::kN * WarpTileIterator::Shape::kColumn + >; + + /// Shape of the shared memory allocation for the epilogue + using StorageShape = MatrixShape< + (Shape::kRow + Padding::kRow) * kFragmentsPerIteration, + Shape::kColumn + Padding::kColumn + >; + + // + // Data members + // + + AlignedBuffer storage; + + // + // Methods + // + + /// Returns a pointer to the shared memory buffer + MCTLASS_DEVICE + Element *data() { + return storage.data(); + } + + /// Returns a tensor reference to the shared memory buffer + MCTLASS_DEVICE + TensorRef reference() { + return TensorRef( + storage.data(), + Layout::packed({StorageShape::kRow, StorageShape::kColumn})); + } + }; + +protected: + + // + // Data members + // + + SharedStorage &shared_storage_; + + /// Stores a warp's fragment of accumulators to SMEM + WarpTileIterator warp_tile_iterator_; + +public: + + /// Constructor + MCTLASS_DEVICE + EpilogueBase( + SharedStorage &shared_storage, ///< Shared storage object + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx ///< Id of thread within warp + ): + shared_storage_(shared_storage), + warp_tile_iterator_(shared_storage.reference(), lane_idx) { + + // Compute warp location within threadblock tile by mapping the warp_id to three coordinates: + // + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_k = warp_idx / (WarpCount::kM * WarpCount::kN); + int warp_mn = warp_idx % (WarpCount::kM * WarpCount::kN); + int warp_m = warp_mn % WarpCount::kM; + int warp_n = warp_mn / WarpCount::kM; + + MatrixCoord warp_offset{warp_k * WarpCount::kM + warp_m, warp_n}; + + warp_tile_iterator_.add_tile_offset(warp_offset); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_base_streamk.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_base_streamk.h new file mode 100644 index 0000000..92e0036 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_base_streamk.h @@ -0,0 +1,197 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Basic subset of epilogue functionality for supporting StreamK decompositions +*/ + + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/functional.h" +#include "mctlass/block_striped.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + + +/// StreamK epilogue functionality for cross-block accumulator fragment reduction +template < + typename Shape, ///< Shape of threadblock tile (concept: GemmShape) + int PartitionsK, + typename WarpMmaOperator, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp) + typename AccumulatorFragmentIterator> ///< Iterator for enumerating fragments within the per-thread tile of raw accumulators +class EpilogueBaseStreamK +{ + +protected: + + /// The per-thread tile of raw accumulators + using AccumulatorTile = typename AccumulatorFragmentIterator::AccumulatorTile; + + /// Number of warps + using WarpCount = gemm::GemmShape< + Shape::kM / WarpMmaOperator::Shape::kM, + Shape::kN / WarpMmaOperator::Shape::kN, + PartitionsK>; + + /// Number of threads per block + static int const kBlockThreads = 32 * WarpCount::kCount; + + /// Numerical accumulation element type + using ElementAccumulator = typename WarpMmaOperator::ElementC; + + /// Fragment type used by the accumulator tile's fragment iterator + using AccumulatorFragment = typename AccumulatorFragmentIterator::Fragment; + +public: + + /// Number of AccumulatorTile fragments per thread + static int const kAccumulatorFragments = AccumulatorFragmentIterator::Policy::kIterations; + +protected: + + /// Number of AccumulatorTile fragments per block output tile + static int const kOutputTileFragments = kBlockThreads * kAccumulatorFragments; + + /// Block-striped transfer utility for sharing AccumulatorFragment + using BlockStripedT = BlockStriped; + + /// AccumulatorFragment stride in the shared workspace between different peer blocks (each thread block can share accumulators for up to two block output tiles) + static const int kPeerFragmentStride = kOutputTileFragments * 2; + +public: + + /// Workspace bytes per thread block + static size_t const kWorkspaceBytesPerBlock =sizeof(AccumulatorFragment) * kPeerFragmentStride; + +public: + + /// Thread index in the threadblock + int thread_idx; + +public: + + /// Constructor + MCTLASS_DEVICE + EpilogueBaseStreamK( + int thread_idx) ///< ID of a thread within the threadblock + : + thread_idx(thread_idx) + {} + + + /// Aggregates the accumulator sets shared by peer blocks in the global workspace + MCTLASS_DEVICE + void reduce( + AccumulatorFragment &accum_fragment, ///< [out] sum of all shared accumulator fragments for these peer partials + int peer_idx_begin, + int peer_idx_end, + int reduce_fragment_idx, + void *workspace_ptr) + { + plus add_fragments; + + AccumulatorFragment *fragment_workspace = reinterpret_cast(workspace_ptr); + + int fragment_offset = (peer_idx_begin * kPeerFragmentStride) + (reduce_fragment_idx * kBlockThreads); + + // Load first peer fragment + BlockStripedT::load(accum_fragment, fragment_workspace + fragment_offset, this->thread_idx); + + fragment_offset += kPeerFragmentStride; // Move to next peer + fragment_offset += kOutputTileFragments; // Move to the set of fragments for this peer's "non-started" output tile + + // Reduce fragments from additional peers + #pragma unroll 2 + for (; fragment_offset < peer_idx_end * kPeerFragmentStride; fragment_offset += kPeerFragmentStride) + { + // Load peer fragment + AccumulatorFragment addend_fragment; + BlockStripedT::load(addend_fragment, fragment_workspace + fragment_offset, this->thread_idx); + + // Add peer fragment + accum_fragment = add_fragments(accum_fragment, addend_fragment); + } + } + + + /// Shares the accumulator set with peers in the global workspace + MCTLASS_DEVICE + void share( + int peer_idx, + void *workspace_ptr, + AccumulatorTile const &accumulators, + bool started_tile) ///< Whether this thread block computed the first work volume for the current output tile + { + AccumulatorFragment *fragment_workspace = reinterpret_cast(workspace_ptr); + + int fragment_offset = peer_idx * kPeerFragmentStride; + + if (!started_tile) { + // Move to the set of fragments for the "non-started" output tile + fragment_offset += kOutputTileFragments; + } + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // Convert raw accumulator tile to fragments and store + MCTLASS_PRAGMA_UNROLL + for (int iter = 0; iter < kAccumulatorFragments; ++iter) + { + // Acquire reordered accumulator fragment + AccumulatorFragment accum_fragment; + accum_fragment_iterator.load(accum_fragment); + ++accum_fragment_iterator; + + // Store accumulator fragment + BlockStripedT::store(fragment_workspace + fragment_offset, accum_fragment, this->thread_idx); + + fragment_offset += kBlockThreads; + } + } + +}; + + + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_depthwise.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_depthwise.h new file mode 100644 index 0000000..a95d6ae --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_depthwise.h @@ -0,0 +1,335 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for Depthwise convoltuion + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/epilogue/thread/conversion_op.h" +#include "mctlass/epilogue/thread/linear_combination.h" +#include "mctlass/epilogue/thread/reduction_op.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/numeric_types.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator +template +class EpilogueDepthwise { + public: + using Shape = Shape_; + using WarpShape = typename WarpMmaOperator_::Shape; + using ThreadOutputShape = ThreadOutputShape_; + using ThreadBlockOutputShape = ThreadBlockOutputShape_; + using WarpMmaOperator = WarpMmaOperator_; + using OutputTileIterator = OutputTileIterator_; + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using WarpTileIterator = WarpTileIterator_; + using SharedLoadIterator = SharedLoadIterator_; + using OutputOp = OutputOp_; + using Padding = Padding_; + + using Layout = layout::RowMajor; + using LongIndex = typename Layout::LongIndex; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename AccumulatorFragmentIterator::AccumulatorTile; + + /// Accumulator element + using ElementAccumulator = typename WarpTileIterator::Element; + + /// Output element + using ElementOutput = typename OutputTileIterator::Element; + + /// Output access size + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + /// Tensor reference to destination tensor + using TensorRef = typename OutputTileIterator::TensorRef; + + /// Tensor reference to sync tensor + using SyncTensorRef = typename mctlass::TensorRef; + + /// Const tensor reference to source tensor + using ConstTensorRef = typename OutputTileIterator::ConstTensorRef; + + /// Array type used to output + using OutputAccessType = + Array; + + /// Array type used by output functor + using AccumulatorAccessType = + Array; + + /// Number of warps + using WarpCount = + gemm::GemmShape; + + public: + static_assert(SharedLoadIterator::Fragment::kElements == + OutputTileIterator::Fragment::kElements, + "Mismatch between shared load iterator and output tile iterator."); + + static_assert(OutputTileIterator::kElementsPerAccess, + "OutputTileIterator::kElementsPerAccess must not be zero."); + + static_assert(!(OutputTileIterator::Fragment::kElements % OutputTileIterator::kElementsPerAccess), + "Divisibility"); + + /// Shared storage allocation needed by the epilogue + struct SharedStorage { + // + // Type definitions + // + + /// Element type of shared memory + using Element = typename WarpTileIterator::Element; + + /// Tensor reference to shared memory allocation + using TensorRef = typename WarpTileIterator::TensorRef; + + /// Layout of shared memory allocation + using Layout = typename WarpTileIterator::Layout; + + /// Logical shape of the shared memory tile written to by all warps. + using Shape = MatrixShape; + + /// Shape of the shared memory allocation for the epilogue + using StorageShape = MatrixShape; + + // + // Data members + // + + AlignedBuffer storage; + + // + // Methods + // + + /// Returns a pointer to the shared memory buffer + MCTLASS_DEVICE + Element *data() { return storage.data(); } + + /// Returns a tensor reference to the shared memory buffer + MCTLASS_DEVICE + TensorRef reference() { + return TensorRef(storage.data(), Layout::packed({StorageShape::kRow, StorageShape::kColumn})); + } + }; + + private: + /// Loads fragment from shared memory aligned with output tensor + SharedLoadIterator shared_load_iterator_; + + /// Stores a warp's fragment of accumulators to SMEM + WarpTileIterator warp_tile_iterator_; + + LongIndex warp_offset; + int thread_idx; + int warp_idx; + int lane_idx; + int warp_m, warp_n; // warp coordinates within a cta + int tid_m, tid_n; // thread coordinates within a warp + + public: + /// Constructor + MCTLASS_DEVICE + EpilogueDepthwise(SharedStorage &shared_storage, ///< Shared storage object + int thread_idx_, ///< ID of a thread within the threadblock + int warp_idx_, ///< ID of warp within threadblock + int lane_idx_ ///< Id of thread within warp + ) + : thread_idx(thread_idx_), + warp_idx(warp_idx_), + lane_idx(lane_idx_), + shared_load_iterator_(shared_storage.reference(), thread_idx_), + warp_tile_iterator_(shared_storage.reference(), thread_idx_, lane_idx_) {} + + /// Streams the result to global memory + MCTLASS_DEVICE + void operator()(OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator, ///< Threadblock tile coordinate in GEMM (in + ///< units of threadblock tiles) + const int smem_base_offset) { ///< SMEM base offset for epilogue operation + // initiate the smem base offset for different output tile. + warp_tile_iterator_.set_smem_base_address(smem_base_offset); + + shared_load_iterator_.set_smem_base_address(smem_base_offset); + + if (!output_op.is_source_needed()) { + compute_source_not_needed_(output_op, destination_iterator, accumulators); + } else { + compute_source_needed_(output_op, destination_iterator, accumulators, source_iterator); + } + } + + private: + /// Streams the result to global memory + MCTLASS_DEVICE + void compute_source_needed_( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + + typename OutputTileIterator::Fragment source_fragment; + + source_fragment.clear(); + + source_iterator.load(source_fragment); + + // store to smem + warp_tile_iterator_.store(accumulators); + + __syncthreads(); + + typename SharedLoadIterator::Fragment aligned_accum_fragment; + + // load from smem + shared_load_iterator_.load(aligned_accum_fragment); + + typename OutputTileIterator::Fragment output_fragment; + + apply_output_operator_(output_fragment, output_op, aligned_accum_fragment, source_fragment); + + // Store to GMEM + destination_iterator.store(output_fragment); + } + + /// Streams the result to global memory + MCTLASS_DEVICE + void compute_source_not_needed_( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + + // store to smem + warp_tile_iterator_.store(accumulators); + + __syncthreads(); + + typename SharedLoadIterator::Fragment aligned_accum_fragment; + + // load from smem + shared_load_iterator_.load(aligned_accum_fragment); + + typename OutputTileIterator::Fragment output_fragment; + + apply_output_operator_source_not_needed_(output_fragment, output_op, aligned_accum_fragment); + + // Store to GMEM + destination_iterator.store(output_fragment); + } + + /// Helper to invoke the output functor over each vector of output + MCTLASS_DEVICE + void apply_output_operator_( + typename OutputTileIterator::Fragment &output_fragment, + OutputOp const &output_op, ///< Output operator + typename SharedLoadIterator::Fragment const &aligned_accum_fragment, + typename OutputTileIterator::Fragment const &source_fragment) { + + OutputAccessType *output_frag_ptr = + reinterpret_cast(&output_fragment); + + AccumulatorAccessType const *compute_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + OutputAccessType const *source_frag_ptr = + reinterpret_cast(&source_fragment); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + // Call the output operator + output_frag_ptr[i] = output_op(compute_frag_ptr[i], source_frag_ptr[i]); + } + } + + /// Helper to invoke the output functor over each vector of output + MCTLASS_DEVICE + void apply_output_operator_source_not_needed_( + typename OutputTileIterator::Fragment &output_fragment, + OutputOp const &output_op, ///< Output operator + typename SharedLoadIterator::Fragment const &aligned_accum_fragment) { + OutputAccessType *output_frag_ptr = reinterpret_cast(&output_fragment); + + AccumulatorAccessType const *compute_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + // Call the output operator + output_frag_ptr[i] = output_op(compute_frag_ptr[i]); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_direct_store.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_direct_store.h new file mode 100644 index 0000000..f099b8f --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_direct_store.h @@ -0,0 +1,374 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs and convolution using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/epilogue/thread/linear_combination.h" +#include "mctlass/epilogue/thread/conversion_op.h" +#include "mctlass/epilogue/thread/reduction_op.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator +template < + typename Shape_, ///< Shape of threadblock tile (concept: GemmShape) + typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp) + int PartitionsK, ///< Number of partitions of the K dimension + typename OutputTileIterator_, ///< Tile iterator reading and writing output tensors + typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators + typename WarpTileIterator_, ///< Warp-scoped tile iterator writing accumulators to SMEM + typename SharedLoadIterator_, ///< Threadblock-scoped tile iterator loading from SMEM + typename OutputOp_ ///< Output operator +> +class EpilogueDirectStore { +public: + + using Shape = Shape_; + using WarpMmaOperator = WarpMmaOperator_; + using WarpShape = typename WarpMmaOperator_::Shape; + static int const kPartitionsK = PartitionsK; + using OutputTileIterator = OutputTileIterator_; + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using WarpTileIterator = WarpTileIterator_; + using OutputOp = OutputOp_; + using Padding = MatrixShape<0, 0>; + + using Layout = layout::RowMajor; + using LongIndex = typename Layout::LongIndex; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename AccumulatorFragmentIterator::AccumulatorTile; + + /// Accumulator element + using ElementAccumulator = typename WarpTileIterator::Element; + + /// Output element + using ElementOutput = typename OutputTileIterator::Element; + + /// Output access size + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + /// Tensor reference to destination tensor + using TensorRef = typename OutputTileIterator::TensorRef; + + /// Tensor reference to sync tensor + using SyncTensorRef = typename mctlass::TensorRef; + + /// Const tensor reference to source tensor + using ConstTensorRef = typename OutputTileIterator::ConstTensorRef; + + /// Array type used to output + using OutputAccessType = Array< + typename OutputTileIterator::Element, OutputTileIterator::kElementsPerAccess>; + + /// Array type used by output functor + using AccumulatorAccessType = Array; + + /// Number of warps + using WarpCount = gemm::GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + kPartitionsK + >; + + /// Use this to control the granularity of one epilogue 'iteration' + static int const kFragmentsPerIteration = 1; + + static int constexpr kSmemTiles = 1; + static int constexpr kSmemPointerOffset = 0; + + /// Shared storage allocation needed by the epilogue + struct SharedStorage { } ; + +private: + + // Assume accumulator tile is multipile interleaved 32x32 tile. + static int const kElementsPerPartial = 4; + using EleShapePerPatial = typename platform::conditional< + platform::is_same::value, + MatrixShape<2, 2>, + MatrixShape<1, 4> >::type; + static int const kElementsPerMma = 8; + static int const kAccumulatorPatials = 2; + using QuadShapePerPatialMma = MatrixShape<4, 4>; + + static_assert(OutputOp::kCount >= 2, + "The direct store epilogue for Tensor Ops requires the output functor have kCount >= 2."); + +private: + + LongIndex warp_offset; + int thread_idx; + int warp_idx; + int lane_idx; + int warp_m, warp_n; // warp coordinates within a cta + int tid_m, tid_n; // thread coordinates within a warp + +public: + + /// Constructor + MCTLASS_DEVICE + EpilogueDirectStore( + SharedStorage &shared_storage, ///< Shared storage object + int thread_idx_, ///< ID of a thread within the threadblock + int warp_idx_, ///< ID of warp within threadblock + int lane_idx_ ///< Id of thread within warp + ): + thread_idx(thread_idx_), + warp_idx(warp_idx_), + lane_idx(lane_idx_) + { + + // warp offsetting calculations + warp_offset = warp_idx * WarpShape::kM * WarpShape::kN; + int warp_id_mn = warp_idx % (WarpCount::kM * WarpShape::kN); + warp_m = warp_id_mn % WarpCount::kM; + warp_n = warp_id_mn / WarpCount::kM; + MatrixCoord warp_offset_coord(warp_m*WarpShape::kM, warp_n*WarpShape::kN); + + // thread offsetting calculations + //int quad = (lane_idx >> 2); + //int lane_in_quad = (lane_idx & 3); + const int quad = ((lane_idx >> 4) << 2); + const int lane_in_quad = (lane_idx & 0x7); + + // this seems to be te correct layout + tid_m = quad; + //tid_n = 2 * lane_in_quad; + tid_n = lane_in_quad; + + } + + /// Streams the result to global memory + MCTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + + if (!output_op.is_source_needed()) { + compute_source_not_needed_(output_op, destination_iterator, accumulators); + } + else { + compute_source_needed_(output_op, destination_iterator, accumulators, source_iterator); + } + } + +private: + + /// Streams the result to global memory + MCTLASS_DEVICE + void compute_source_needed_( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + + const int kAccumBlockN = 4;//2; + const int kThreadsM = 16;//8; + const int kThreadsN = 2;//4; + const int kBlockM = WarpShape::kM / kThreadsM; + + /// Array type used to output + using OutputAccessType = AlignedArray; + + /// Array type passed to the output operator - unused elements are optimized away + using OutputFragmentType = Array; + + /// Array type used by output functor + using AccumulatorAccessType = Array; + + /// Array type used by output functor + using AccumulatorFragmentType = Array; + + AccumulatorAccessType const *accumulator_pair = reinterpret_cast(&accumulators); + + MCTLASS_PRAGMA_UNROLL + for (int accum_m_idx = 0; accum_m_idx < WarpShape::kM / kThreadsM; accum_m_idx++) { + + int accum_m = kThreadsM * accum_m_idx; + int mL = destination_iterator.threadblock_offset.row() + WarpShape::kM * warp_m + tid_m + accum_m; + int nL_base = destination_iterator.threadblock_offset.column() + WarpShape::kN * warp_n + tid_n; + + //ElementOutput *output_ptr = destination_iterator.pointer + mL * destination_iterator.stride; + //ElementOutput *source_ptr = source_iterator.pointer + mL * source_iterator.stride; + + int const kIterationsN = WarpShape::kN / kThreadsN / kAccumBlockN; + + MCTLASS_PRAGMA_UNROLL + for (int accum_n_idx = 0; accum_n_idx < kIterationsN; accum_n_idx++) { + + int accum_idx = accum_m_idx + kBlockM * accum_n_idx; + //int accum_n = kThreadsM * accum_n_idx; + int accum_n = (kThreadsM / 2) * accum_n_idx; + + // mL and nL are logical coordinate in 2D mapping of epilogue's 4D output + int nL = nL_base + accum_n; + + //bool guard = (mL < destination_iterator.extent.row()) && (nL < destination_iterator.extent.column()); + bool guard = (nL < destination_iterator.extent.column()); + + AccumulatorFragmentType accum_fragment; + reinterpret_cast(accum_fragment) = accumulator_pair[accum_idx]; + + OutputFragmentType output_fragment; + + if(guard) { + // reinterpret_cast(output_fragment) = + // *reinterpret_cast(source_ptr + nL); + for (int i = 0; i < OutputOp::kCount; ++i) { + const bool guard_row = ((mL + i) < destination_iterator.extent.row()); + if (guard_row) { + ElementOutput *source_ptr = source_iterator.pointer + (mL + i) * source_iterator.stride; + output_fragment[i] = source_ptr[nL]; + } + } + } + + // Perform output operator + output_fragment = output_op(accum_fragment, output_fragment); + + if(guard) { + // Store + //*reinterpret_cast(output_ptr + nL) = reinterpret_cast(output_fragment); + for (int i = 0; i < OutputOp::kCount; ++i) { + const bool guard_row = ((mL + i) < destination_iterator.extent.row()); + if (guard_row) { + ElementOutput *output_ptr = destination_iterator.pointer + (mL + i) * destination_iterator.stride; + output_ptr[nL] = output_fragment[i]; + } + } + } + } + } + } + + /// Streams the result to global memory + MCTLASS_DEVICE + void compute_source_not_needed_( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + + const int kAccumBlockN = 4;//2; + const int kThreadsM = 16;//8; + const int kThreadsN = 2;//4; + const int kBlockM = WarpShape::kM / kThreadsM; + + /// Array type used to output + using OutputAccessType = AlignedArray; + + /// Array type passed to the output operator - unused elements are optimized away + using OutputFragmentType = Array; + + /// Array type used by output functor + using AccumulatorAccessType = Array; + + /// Array type used by output functor + using AccumulatorFragmentType = Array; + + AccumulatorAccessType const *accumulator_pair = reinterpret_cast(&accumulators); + + MCTLASS_PRAGMA_UNROLL + for (int accum_m_idx = 0; accum_m_idx < WarpShape::kM / kThreadsM; accum_m_idx++) { + + int accum_m = kThreadsM * accum_m_idx; + int mL = destination_iterator.threadblock_offset.row() + WarpShape::kM * warp_m + tid_m + accum_m; + int nL_base = destination_iterator.threadblock_offset.column() + WarpShape::kN * warp_n + tid_n; + + //ElementOutput *output_ptr = destination_iterator.pointer + mL * destination_iterator.stride; + + int const kIterationsN = WarpShape::kN / kThreadsN / kAccumBlockN; + + MCTLASS_PRAGMA_UNROLL + for (int accum_n_idx = 0; accum_n_idx < kIterationsN; accum_n_idx++) { + + int accum_idx = accum_m_idx + kBlockM * accum_n_idx; + // int accum_n = kThreadsM * accum_n_idx; + int accum_n = (kThreadsM / 2) * accum_n_idx; + + // mL and nL are logical coordinate in 2D mapping of epilogue's 4D output + int nL = nL_base + accum_n; + // bool guard = (mL < destination_iterator.extent.row()) && (nL < destination_iterator.extent.column()); + bool guard = (nL < destination_iterator.extent.column()); + + AccumulatorFragmentType accum_fragment; + reinterpret_cast(accum_fragment) = accumulator_pair[accum_idx]; + + OutputFragmentType output_fragment; + + // Perform output operator + output_fragment = output_op(accum_fragment); + if(guard) { + // Store + // *reinterpret_cast(output_ptr + nL) = + // reinterpret_cast(output_fragment); + for (int i = 0; i < OutputOp::kCount; ++i) { + + const bool guard_row = ((mL + i) < destination_iterator.extent.row()); + if(guard_row) { + ElementOutput *output_ptr = destination_iterator.pointer + (mL + i) * destination_iterator.stride; + output_ptr[nL] = output_fragment[i]; + } + } + } + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_gemm_k_reduction.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_gemm_k_reduction.h new file mode 100644 index 0000000..c260b25 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_gemm_k_reduction.h @@ -0,0 +1,212 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/vector.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/tensor_coord.h" +#include "mctlass/aligned_buffer.h" +#include "mctlass/functional.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_iterator.h" + +#include "mctlass/epilogue/threadblock/epilogue_base.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator.h" +#include "mctlass/numeric_types.h" + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator +template < + typename ElementAccumulator_, + typename ElementOutput_, + typename ThreadBlockShape_, ///< Shape of threadblock tile (concept: GemmShape) + typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp) + bool ReduceKForA_ +> +class EpilogueGemmKReduction { + +public: + + using ThreadBlockShape = ThreadBlockShape_; + using WarpMmaOperator = WarpMmaOperator_; + using WarpShape = typename WarpMmaOperator::Shape; + using Layout = layout::RowMajor; + using LongIndex = typename Layout::LongIndex; + + /// Accumulator element + using ElementAccumulator = ElementAccumulator_; + + /// Output element + using ElementOutput = ElementOutput_; + + /// Output access size + static int const kElementsPerAccess = 1; + + static bool const kReduceKForA = ReduceKForA_; + + static int const kThreadBlockSize = kReduceKForA ? ThreadBlockShape::kM : ThreadBlockShape::kN; + + static int const kWarpSize = kReduceKForA ? WarpShape::kM : WarpShape::kN; + + static int const kIterations = kWarpSize / 8; + + using FragmentAccumulator = Array; + +private: + + int thread_offset_; + ElementOutput* pointer_; + int col_; +public: + + /// Constructor + MCTLASS_DEVICE + EpilogueGemmKReduction( + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx, ///< Id of thread within warp + int threadblock_offset, + ElementOutput* pointer + ) + { + col_ = lane_idx % 4; + thread_offset_ = threadblock_offset * kThreadBlockSize + + warp_idx * kWarpSize + + lane_idx / 4 + col_ * 8; + + pointer_ = pointer + LongIndex(thread_offset_); + } + + /// Streams the result to global memory + MCTLASS_DEVICE + void operator()( + int size, + FragmentAccumulator &gemm_k_with_reduction_accumulation, + bool LoadForSerialSplitK + ) { + bool guard[kIterations / 4]; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kIterations / 4; ++i) { + guard[i] = ((thread_offset_ + i * 32) < size); + } + + Array source; + source.clear(); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kIterations / 4; ++i) { + ElementOutput *source_ptr = reinterpret_cast(&source); + mctlass::arch::global_load( + source_ptr[i], + (void *)(pointer_ + i * 32), + guard[i] && LoadForSerialSplitK); + + } + + FragmentAccumulator sum = gemm_k_with_reduction_accumulation; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kIterations; ++i) { + sum[i] += __shfl_xor_sync(0xffffffff, sum[i], 1); + sum[i] += __shfl_xor_sync(0xffffffff, sum[i], 2); + } + + Array intermediate; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kIterations / 4; ++i) { + if (col_ == 0) { + intermediate[i] = sum[0 + i * 4]; + } + + if (col_ == 1) { + intermediate[i] = sum[1 + i * 4]; + } + + if (col_ == 2) { + intermediate[i] = sum[2 + i * 4]; + } + + if (col_ == 3) { + intermediate[i] = sum[3 + i * 4]; + } + } + + NumericArrayConverter source_converter; + Array converted_source = source_converter(source); + + plus> plus_source; + intermediate = plus_source(intermediate, converted_source); + + NumericArrayConverter converter; + Array result = converter(intermediate); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kIterations / 4; ++i) { + mctlass::arch::global_store(result[i], + (void *)(pointer_ + i * 32), guard[i]); + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_planar_complex.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_planar_complex.h new file mode 100644 index 0000000..9f6b40e --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_planar_complex.h @@ -0,0 +1,441 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/array_planar_complex.h" +#include "mctlass/layout/vector.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/tensor_coord.h" +#include "mctlass/aligned_buffer.h" +#include "mctlass/functional.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_iterator.h" + +#include "mctlass/epilogue/threadblock/epilogue_base.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator for planar-complex output representations. +/// +/// Note, as with most MCTLASS components for planar complex, the template arguments describe +/// the underlying real data type. +template < + typename Shape_, ///< Shape of threadblock tile (concept: GemmShape) + typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp) + int PartitionsK, ///< Number of partitions of the K dimension + typename OutputTileIterator_, ///< Tile iterator reading and writing output tensors + typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators + typename WarpTileIterator_, ///< Warp-scoped tile iterator writing accumulators to SMEM + typename SharedLoadIterator_, ///< Threadblock-scoped tile iterator loading from SMEM + typename OutputOp_, ///< Output operator + typename Padding_ ///< Padding added to SMEM allocation to avoid bank conflicts (concept: MatrixShape) +> +class EpiloguePlanarComplex { +public: + + using Shape = Shape_; + using WarpMmaOperator = WarpMmaOperator_; + static int const kPartitionsK = PartitionsK; + using OutputTileIterator = OutputTileIterator_; + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using WarpTileIterator = WarpTileIterator_; + using SharedLoadIterator = SharedLoadIterator_; + using OutputOp = OutputOp_; + using Padding = Padding_; + + /// Output layout is always row-major + using Layout = layout::RowMajor; + using LongIndex = typename Layout::LongIndex; + + /// The complete warp-level accumulator tile + using AccumulatorTile = ArrayPlanarComplex< + typename WarpMmaOperator::FragmentC::Element, + WarpMmaOperator::FragmentC::kElements + >; + + /// Accumulator element + using ElementAccumulator = typename WarpTileIterator::Element; + + /// Output element + using ElementOutput = typename OutputTileIterator::Element; + + /// Output access size + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + /// Tensor reference to destination tensor + using TensorRef = typename OutputTileIterator::TensorRef; + + /// Tensor reference to sync tensor + using SyncTensorRef = typename mctlass::TensorRef; + + /// Const tensor reference to source tensor + using ConstTensorRef = typename OutputTileIterator::ConstTensorRef; + + /// Array type used to output + using OutputAccessType = Array< + typename OutputTileIterator::Element, OutputTileIterator::kElementsPerAccess>; + + /// Array type used by output functor + using AccumulatorAccessType = Array; + + /// Shape of each warp-level operation + using WarpShape = typename WarpMmaOperator::Shape; + + /// Number of warps + using WarpCount = gemm::GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + kPartitionsK + >; + + /// Shared memory allocation + struct SharedStorage { + + // + // Type definitions + // + + /// Element type of shared memory + using Element = typename WarpTileIterator::Element; + + /// Tensor reference to shared memory allocation + using TensorRef = typename WarpTileIterator::TensorRef; + + /// Layout of shared memory allocation + using Layout = typename WarpTileIterator::Layout; + + /// Logical shape of the shared memory tile written to by all warps. + using Shape = MatrixShape< + WarpCount::kM * WarpTileIterator::Shape::kRow * WarpCount::kK, + WarpCount::kN * WarpTileIterator::Shape::kColumn + >; + + /// Shape of the shared memory allocation for the epilogue + using StorageShape = MatrixShape< + Shape::kRow + Padding::kRow, + Shape::kColumn + Padding::kColumn + >; + + static int const kImaginaryStride = StorageShape::kCount; + + // + // Data members + // + + AlignedBuffer storage; + + // + // Methods + // + + /// Returns a pointer to the shared memory buffer + MCTLASS_DEVICE + Element *data() { + return storage.data(); + } + + /// Returns a tensor reference to the shared memory buffer + MCTLASS_DEVICE + TensorRef reference() { + return TensorRef( + storage.data(), + Layout::packed({StorageShape::kRow, StorageShape::kColumn})); + } + }; + +private: + + // + // Data members + // + + SharedStorage &shared_storage_; + + /// Loads fragment from shared memory aligned with output tensor + SharedLoadIterator shared_load_iterator_; + + /// Stores a warp's fragment of accumulators to SMEM + WarpTileIterator warp_tile_iterator_; + +public: + + /// Constructor + MCTLASS_DEVICE + EpiloguePlanarComplex( + SharedStorage &shared_storage, ///< Shared storage object + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx ///< Id of thread within warp + ): + shared_storage_(shared_storage), + shared_load_iterator_(shared_storage.reference(), thread_idx), + warp_tile_iterator_(shared_storage.reference(), lane_idx) { + + // Compute warp location within threadblock tile by mapping the warp_id to three coordinates: + // + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_k = warp_idx / (WarpCount::kM * WarpCount::kN); + int warp_mn = warp_idx % (WarpCount::kM * WarpCount::kN); + int warp_m = warp_mn % WarpCount::kM; + int warp_n = warp_mn / WarpCount::kM; + + MatrixCoord warp_offset{warp_k * WarpCount::kM + warp_m, warp_n}; + + warp_tile_iterator_.add_tile_offset(warp_offset); + } + + /// Streams the result to global memory + MCTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator_real, ///< Tile iterator for destination + OutputTileIterator destination_iterator_imag, ///< Tile iterator for destination + AccumulatorTile &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator_real, ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + OutputTileIterator source_iterator_imag) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + + typename OutputTileIterator::Fragment source_fragment_real; + typename OutputTileIterator::Fragment source_fragment_imag; + + if (!output_op.is_source_needed()) { + source_iterator_real.clear_mask(); + source_iterator_imag.clear_mask(); + } + + source_fragment_real.clear(); + source_fragment_imag.clear(); + + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator_real(accumulators.real); + AccumulatorFragmentIterator accum_fragment_iterator_imag(accumulators.imag); + + //Special process for mma<16,8,16> and warpshape=<32,32,32> + if (mctlass::platform::is_same::value == true && + WarpShape::kM == 32 && WarpShape::kN == 32 && WarpShape::kK == 32) { + for (int i = 0; i < accumulators.real.size() / 2; i += 2) { + + ElementAccumulator x_real = accumulators.real[2 * i + 0]; + ElementAccumulator y_real = accumulators.real[2 * i + 1]; + ElementAccumulator z_real = accumulators.real[2 * (i + 1) + 0]; + ElementAccumulator w_real = accumulators.real[2 * (i + 1) + 1]; + + ElementAccumulator x_imag = accumulators.imag[2 * i + 0]; + ElementAccumulator y_imag = accumulators.imag[2 * i + 1]; + ElementAccumulator z_imag = accumulators.imag[2 * (i + 1) + 0]; + ElementAccumulator w_imag = accumulators.imag[2 * (i + 1) + 1]; + + float x0 = __shfl_down_sync(UINT64_MAX, float(x_real), 32); + float y0 = __shfl_down_sync(UINT64_MAX, float(y_real), 32); + float z0 = __shfl_up_sync(UINT64_MAX, float(z_real), 32); + float w0 = __shfl_up_sync(UINT64_MAX, float(w_real), 32); + + float x1 = __shfl_down_sync(UINT64_MAX, float(x_imag), 32); + float y1 = __shfl_down_sync(UINT64_MAX, float(y_imag), 32); + float z1 = __shfl_up_sync(UINT64_MAX, float(z_imag), 32); + float w1 = __shfl_up_sync(UINT64_MAX, float(w_imag), 32); + + if (__lane_id() < 32) { + accumulators.real[2 * (i + 1) + 0] = ElementAccumulator(x0); + accumulators.real[2 * (i + 1) + 1] = ElementAccumulator(y0); + + accumulators.imag[2 * (i + 1) + 0] = ElementAccumulator(x1); + accumulators.imag[2 * (i + 1) + 1] = ElementAccumulator(y1); + } + else { + accumulators.real[2 * i + 0] = ElementAccumulator(z0); + accumulators.real[2 * i + 1] = ElementAccumulator(w0); + accumulators.imag[2 * i + 0] = ElementAccumulator(z1); + accumulators.imag[2 * i + 1] = ElementAccumulator(w1); + } + } + } + // + // Iterate over accumulator tile + // + MCTLASS_PRAGMA_UNROLL + for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) + { + + // + // Load the source + // + + source_iterator_real.load(source_fragment_real); + source_iterator_imag.load(source_fragment_imag); + + ++source_iterator_real; + ++source_iterator_imag; + + // + // Convert and store fragment + // + + __syncthreads(); + + typename AccumulatorFragmentIterator::Fragment accum_fragment_real; + typename AccumulatorFragmentIterator::Fragment accum_fragment_imag; + + accum_fragment_iterator_real.load(accum_fragment_real); + accum_fragment_iterator_imag.load(accum_fragment_imag); + + ++accum_fragment_iterator_real; + ++accum_fragment_iterator_imag; + + this->warp_tile_iterator_.store(accum_fragment_real); + this->warp_tile_iterator_.store_with_pointer_offset(accum_fragment_imag, SharedStorage::kImaginaryStride); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + typename SharedLoadIterator::Fragment aligned_accum_fragment_real[kPartitionsK]; + typename SharedLoadIterator::Fragment aligned_accum_fragment_imag[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment_real[0]); + shared_load_iterator_.load_with_pointer_offset(aligned_accum_fragment_imag[0], SharedStorage::kImaginaryStride); + + // If the number of k-slices is > 1 - perform a reduction amongst the k-slices + static_assert(kPartitionsK == 1, "Sliced-K not supported for planar complex at this time"); + + // + // Compute the output result + // + + typename OutputTileIterator::Fragment output_fragment_real; + typename OutputTileIterator::Fragment output_fragment_imag; + + apply_output_operator_( + output_fragment_real, + output_fragment_imag, + output_op, + aligned_accum_fragment_real[0], + aligned_accum_fragment_imag[0], + source_fragment_real, + source_fragment_imag); + + // + // Store the final result + // + + destination_iterator_real.store(output_fragment_real); + destination_iterator_imag.store(output_fragment_imag); + + ++destination_iterator_real; + ++destination_iterator_imag; + } + } + +private: + + /// Helper to invoke the output functor over each vector of output + MCTLASS_DEVICE + void apply_output_operator_( + typename OutputTileIterator::Fragment &output_fragment_real, + typename OutputTileIterator::Fragment &output_fragment_imag, + OutputOp const &output_op, ///< Output operator + typename SharedLoadIterator::Fragment const &aligned_accum_fragment_real, + typename SharedLoadIterator::Fragment const &aligned_accum_fragment_imag, + typename OutputTileIterator::Fragment const &source_fragment_real, + typename OutputTileIterator::Fragment const &source_fragment_imag) { + + OutputAccessType *output_frag_real_ptr = + reinterpret_cast(&output_fragment_real); + + OutputAccessType *output_frag_imag_ptr = + reinterpret_cast(&output_fragment_imag); + + AccumulatorAccessType const *compute_frag_real_ptr = + reinterpret_cast(&aligned_accum_fragment_real); + + AccumulatorAccessType const *compute_frag_imag_ptr = + reinterpret_cast(&aligned_accum_fragment_imag); + + OutputAccessType const *source_frag_real_ptr = + reinterpret_cast(&source_fragment_real); + + OutputAccessType const *source_frag_imag_ptr = + reinterpret_cast(&source_fragment_imag); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + + // Call the output operator + auto result_fragment = output_op( + make_ArrayPlanarComplex(compute_frag_real_ptr[i], compute_frag_imag_ptr[i]), + make_ArrayPlanarComplex(source_frag_real_ptr[i], source_frag_imag_ptr[i]) + ); + + output_frag_real_ptr[i] = result_fragment.real; + output_frag_imag_ptr[i] = result_fragment.imag; + } + } + +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_smem_accumulator.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_smem_accumulator.h new file mode 100644 index 0000000..a38912c --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_smem_accumulator.h @@ -0,0 +1,230 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMM/CONV to store accumulator in shared memory after + applying scale, bias loaded from global memory and element-wise operations. + + This Epilogue is typically used in fused GEMM/CONV to stage the intermediate accumulator. + +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/vector.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/tensor_coord.h" +#include "mctlass/aligned_buffer.h" +#include "mctlass/functional.h" + +#include "mctlass/epilogue/warp/fragment_iterator_tensor_op.h" +#include "mctlass/epilogue/warp/tile_iterator_tensor_op.h" +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator +template < + typename SmemTileIterator_, ///< Shared memory Tile iterator to output to shared memory + typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators + typename ScaleBiasIterator_, ///< Iterator to load scale and bias from global memory + typename OutputOp_ ///< Output operator +> +class EpilogueSmemAccumulator { + +public: + + using SmemTileIterator = SmemTileIterator_; + + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + + using ScaleBiasIterator = ScaleBiasIterator_; + + using OutputOp = OutputOp_; + + /// Fragment of accumulator tile + using FragmentAccumulator = typename AccumulatorFragmentIterator::Fragment; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename AccumulatorFragmentIterator::AccumulatorTile; + + /// Fragment of Scale and Bias loaded from global memory + using FragmentScaleBias = typename ScaleBiasIterator::Fragment; + + static const bool PerChannelScale = (OutputOp::kScale == + epilogue::thread::ScaleType::OnlyAlphaPerChannelScaling); + + /// Constructor + MCTLASS_DEVICE + EpilogueSmemAccumulator() {} + + /// Streams the result to shared memory + MCTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + SmemTileIterator smem_iterator, ///< Tile iterator for destination in shared memory + AccumulatorTile const &accumulator, ///< Complete warp-level accumulator tile + ScaleBiasIterator scale_iterator, ///< iterator for scale vector in global memory + ScaleBiasIterator bias_iterator) { ///< iterator for bias vector in global memory + + + // Fragment to load scale bias from global memory + FragmentScaleBias tb_frag_scale; + FragmentScaleBias tb_frag_bias; + + /// Fragment Iterator to load slice of accumulator tile + AccumulatorFragmentIterator frag_iterator_accum(accumulator); + FragmentAccumulator tb_frag_accum; + + /// Epilogue output fragment + typename SmemTileIterator::Fragment tb_frag_smem; + + /// Load scale and bias from global memory + + if(PerChannelScale) + scale_iterator.load(tb_frag_scale); + + bias_iterator.load(tb_frag_bias); + + /// Iterate over the accumulator tile and store to shared memory + MCTLASS_PRAGMA_UNROLL + for (int rid = 0; rid < AccumulatorFragmentIterator::TileIterations::kRow; ++rid) { + + MCTLASS_PRAGMA_UNROLL + for (int cid = 0; cid < AccumulatorFragmentIterator::TileIterations::kColumn; ++cid) { + + using AccumulatorAccessType = typename OutputOp::FragmentAccumulator; + using ScaleBiasAccessType = typename OutputOp::FragmentScaleBias; + using FragmentSmemAccessType = typename OutputOp::FragmentOutput; + + + ScaleBiasAccessType const * scale_frag_ptr = + reinterpret_cast(&tb_frag_scale); + ScaleBiasAccessType const * bias_frag_ptr = + reinterpret_cast(&tb_frag_bias); + + FragmentSmemAccessType * smem_frag_ptr = + reinterpret_cast(&tb_frag_smem); + + MCTLASS_PRAGMA_UNROLL + for (int idx = 0; idx < AccumulatorFragmentIterator::kIterationsPerTile; ++idx) { + frag_iterator_accum.load(tb_frag_accum); + ++frag_iterator_accum; + + AccumulatorAccessType const * accumulator_frag_ptr = + reinterpret_cast(&tb_frag_accum); + const int kOutputIterations = FragmentAccumulator::kElements / OutputOp::kCount; + + MCTLASS_PRAGMA_UNROLL + for (int it = 0; it < kOutputIterations; it++) { + smem_frag_ptr[idx * kOutputIterations + it] = output_op(accumulator_frag_ptr[it], + scale_frag_ptr[cid * kOutputIterations + it], bias_frag_ptr[cid * kOutputIterations + it]); + } + } + + smem_iterator.store(tb_frag_smem); + ++smem_iterator; + + } + } + } + + /// Streams the result to shared memory + MCTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + SmemTileIterator smem_iterator, ///< Tile iterator for destination in shared memory + AccumulatorTile const &accumulator) { ///< Complete warp-level accumulator tile + + /// Fragment Iterator to load slice of accumulator tile + AccumulatorFragmentIterator frag_iterator_accum(accumulator); + FragmentAccumulator tb_frag_accum; + + /// Epilogue output fragment + typename SmemTileIterator::Fragment tb_frag_smem; + + /// Iterate over the accumulator tile and store to shared memory + MCTLASS_PRAGMA_UNROLL + for (int rid = 0; rid < AccumulatorFragmentIterator::TileIterations::kRow; ++rid) { + + MCTLASS_PRAGMA_UNROLL + for (int cid = 0; cid < AccumulatorFragmentIterator::TileIterations::kColumn; ++cid) { + + using AccumulatorAccessType = typename OutputOp::FragmentAccumulator; + using FragmentSmemAccessType = typename OutputOp::FragmentOutput; + + FragmentSmemAccessType * smem_frag_ptr = + reinterpret_cast(&tb_frag_smem); + + MCTLASS_PRAGMA_UNROLL + for (int idx = 0; idx < AccumulatorFragmentIterator::kIterationsPerTile; ++idx) { + frag_iterator_accum.load(tb_frag_accum); + ++frag_iterator_accum; + + AccumulatorAccessType const * accumulator_frag_ptr = + reinterpret_cast(&tb_frag_accum); + const int kOutputIterations = FragmentAccumulator::kElements / OutputOp::kCount; + + MCTLASS_PRAGMA_UNROLL + for (int it = 0; it < kOutputIterations; it++) { + smem_frag_ptr[idx * kOutputIterations + it] = output_op(accumulator_frag_ptr[it]); + } + } + + smem_iterator.store(tb_frag_smem); + ++smem_iterator; + + } + } + } + +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_streamk_with_broadcast.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_streamk_with_broadcast.h new file mode 100644 index 0000000..05e69d3 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_streamk_with_broadcast.h @@ -0,0 +1,443 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#include +#else +#include +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/tensor_coord.h" +#include "mctlass/aligned_buffer.h" +#include "mctlass/functional.h" +#include "mctlass/fast_math.h" +#include "mctlass/layout/vector.h" +#include "mctlass/layout/tensor.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_iterator.h" + +#include "mctlass/epilogue/threadblock/epilogue_base.h" +#include "mctlass/epilogue/threadblock/epilogue_base_streamk.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator.h" + +#include "mctlass/numeric_types.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// This base class is meant to define the concept required of the +/// EpilogueStreamkWithBroadcast::OutputOp +template < + typename ElementC_, + typename ElementAccumulator_, + typename ElementCompute_, + typename ElementZ_, + typename ElementT_, + int ElementsPerAccess, + bool StoreZ = true, + bool StoreT = true +> +struct EpilogueStreamkWithBroadcastOpBase : EpilogueWithBroadcastOpBase< + ElementC_, + ElementAccumulator_, + ElementCompute_, + ElementZ_, + ElementT_, + ElementsPerAccess, + StoreZ, + StoreT + > +{ + + /// Parameters structure - required + struct Params { }; + + // + // Methods + // + + /// Constructor from Params + EpilogueStreamkWithBroadcastOpBase(Params const ¶ms_) { } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator with bias vector broadcast over columns. +/// +/// Computes the following: +/// +/// +/// Z, T = OutputOp(AB, C, Broadcast) +/// +/// if (ElementwiseOp::kStoreZ) { +/// store(converted_u); +/// } +/// +/// if (ElementwiseOp::kStoreT) { +/// store(v); +/// } +/// +template < + typename Shape_, ///< Shape of threadblock tile (concept: GemmShape) + typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp) + int PartitionsK, ///< Number of partitions of the K dimension + typename OutputTileIterator_, ///< Tile iterator reading and writing output tensors (z) + typename TensorTileIterator_, ///< Additional tile iterator for tensor-valued operands (t) + typename ElementVector_, ///< Pointer to broadcast vector + typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators + typename WarpTileIterator_, ///< Warp-scoped tile iterator writing accumulators to SMEM + typename SharedLoadIterator_, ///< Threadblock-scoped tile iterator loading from SMEM + typename OutputOp_, ///< Output operator - concept is EpilogueWithBroadcastOp + typename Padding_, ///< Padding added to SMEM allocation to avoid bank conflicts (concept: MatrixShape) + int FragmentsPerPartition = 1, ///< Used to coarsten the epilogue granularity + int IterationsUnroll = ///< Used to reduce binary size when epilogue op is large + (!IsEpilogueFunctorHeavy::value), + bool IsSingleSource = OutputOp_::kIsSingleSource +> +class EpilogueStreamkWithBroadcast; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// EpilogueStreamkWithBroadcast: Two sources + +template < + typename Shape_, + typename WarpMmaOperator_, + int PartitionsK, + typename OutputTileIterator_, + typename TensorTileIterator_, + typename ElementVector_, + typename AccumulatorFragmentIterator_, + typename WarpTileIterator_, + typename SharedLoadIterator_, + typename OutputOp_, + typename Padding_, + int FragmentsPerPartition, + int IterationsUnroll +> +class EpilogueStreamkWithBroadcast< + Shape_, + WarpMmaOperator_, + PartitionsK, + OutputTileIterator_, + TensorTileIterator_, + ElementVector_, + AccumulatorFragmentIterator_, + WarpTileIterator_, + SharedLoadIterator_, + OutputOp_, + Padding_, + FragmentsPerPartition, + IterationsUnroll, + false +> : + public EpilogueWithBroadcast< + Shape_, + WarpMmaOperator_, + PartitionsK, + OutputTileIterator_, + TensorTileIterator_, + ElementVector_, + AccumulatorFragmentIterator_, + WarpTileIterator_, + SharedLoadIterator_, + OutputOp_, + Padding_, + FragmentsPerPartition, + IterationsUnroll, + false>, + public EpilogueBaseStreamK< + Shape_, + PartitionsK, + WarpMmaOperator_, + AccumulatorFragmentIterator_> +{ + +public: + + using Base = EpilogueWithBroadcast< + Shape_, + WarpMmaOperator_, + PartitionsK, + OutputTileIterator_, + TensorTileIterator_, + ElementVector_, + AccumulatorFragmentIterator_, + WarpTileIterator_, + SharedLoadIterator_, + OutputOp_, + Padding_, + FragmentsPerPartition, + IterationsUnroll, + false>; + + using BaseStreamK = EpilogueBaseStreamK< + Shape_, + PartitionsK, + WarpMmaOperator_, + AccumulatorFragmentIterator_>; + + using Shape = Shape_; + static int const kPartitionsK = PartitionsK; + using OutputTileIterator = OutputTileIterator_; + using TensorTileIterator = TensorTileIterator_; + using ElementVector = ElementVector_; + using SharedLoadIterator = SharedLoadIterator_; + using OutputOp = OutputOp_; + + /// Fragment type used by the accumulator tile's fragment iterator + using AccumulatorFragment = typename Base::AccumulatorFragmentIterator::Fragment; + + /// Shared storage structure (shadows base) with additional SMEM buffer for reduction + using SharedStorage = typename Base::SharedStorage; + +public: + + /// Constructor + MCTLASS_DEVICE + EpilogueStreamkWithBroadcast( + SharedStorage &shared_storage, ///< Shared storage object + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx ///< Id of thread within warp + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + BaseStreamK(thread_idx) + { } + + + /// Aggregates the accumulator sets shared by peer blocks in the global workspace, + /// performing epilogue computations, writing to output + MCTLASS_DEVICE + void reduce( + int peer_idx_begin, + int peer_idx_end, + int reduce_fragment_idx, + void *element_workspace, + OutputOp const &output_op, ///< Output operator + ElementVector const * broadcast_ptr, ///< Broadcast vector + OutputTileIterator destination_iterator, ///< Tile iterator for destination + OutputTileIterator source_iterator1, ///< Tile iterator for first source accumulator matrix + OutputTileIterator source_iterator2, ///< Tile iterator for second source accumulator matrix + TensorTileIterator tensor_iterator, ///< Threadblock tile iterator for additional tensor operand + MatrixCoord const &problem_size = ///< Problem size needed to guard against out-of-bounds accesses + MatrixCoord(Shape::kM, Shape::kN), + MatrixCoord const &threadblock_offset = ///< Threadblock's initial offset within the problem size space + MatrixCoord()) + { + // Reduce peer accumulator fragments into one fragment + AccumulatorFragment accum_fragment; + BaseStreamK::reduce(accum_fragment, peer_idx_begin, peer_idx_end, reduce_fragment_idx, element_workspace); + + // Store fragment to shared memory + this->warp_tile_iterator_.store(accum_fragment); + + __syncthreads(); + + Base::reduce(reduce_fragment_idx, output_op, broadcast_ptr, destination_iterator, source_iterator1, source_iterator2, tensor_iterator, problem_size, threadblock_offset); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// EpilogueStreamkWithBroadcast: Single source + +template < + typename Shape_, + typename WarpMmaOperator_, + int PartitionsK, + typename OutputTileIterator_, + typename TensorTileIterator_, + typename ElementVector_, + typename AccumulatorFragmentIterator_, + typename WarpTileIterator_, + typename SharedLoadIterator_, + typename OutputOp_, + typename Padding_, + int FragmentsPerPartition, + int IterationsUnroll +> +class EpilogueStreamkWithBroadcast< + Shape_, + WarpMmaOperator_, + PartitionsK, + OutputTileIterator_, + TensorTileIterator_, + ElementVector_, + AccumulatorFragmentIterator_, + WarpTileIterator_, + SharedLoadIterator_, + OutputOp_, + Padding_, + FragmentsPerPartition, + IterationsUnroll, + true +> : + public EpilogueWithBroadcast< + Shape_, + WarpMmaOperator_, + PartitionsK, + OutputTileIterator_, + TensorTileIterator_, + ElementVector_, + AccumulatorFragmentIterator_, + WarpTileIterator_, + SharedLoadIterator_, + OutputOp_, + Padding_, + FragmentsPerPartition, + IterationsUnroll, + true>, + public EpilogueBaseStreamK< + Shape_, + PartitionsK, + WarpMmaOperator_, + AccumulatorFragmentIterator_> +{ + +public: + + using Base = EpilogueWithBroadcast< + Shape_, + WarpMmaOperator_, + PartitionsK, + OutputTileIterator_, + TensorTileIterator_, + ElementVector_, + AccumulatorFragmentIterator_, + WarpTileIterator_, + SharedLoadIterator_, + OutputOp_, + Padding_, + FragmentsPerPartition, + IterationsUnroll, + true>; + + using BaseStreamK = EpilogueBaseStreamK< + Shape_, + PartitionsK, + WarpMmaOperator_, + AccumulatorFragmentIterator_>; + + using Shape = Shape_; + static int const kPartitionsK = PartitionsK; + using OutputTileIterator = OutputTileIterator_; + using TensorTileIterator = TensorTileIterator_; + using ElementVector = ElementVector_; + using SharedLoadIterator = SharedLoadIterator_; + using OutputOp = OutputOp_; + + /// Fragment type used by the accumulator tile's fragment iterator + using AccumulatorFragment = typename Base::AccumulatorFragmentIterator::Fragment; + + /// Shared storage structure (shadows base) with additional SMEM buffer for reduction + using SharedStorage = typename Base::SharedStorage; + +public: + + /// Constructor + MCTLASS_DEVICE + EpilogueStreamkWithBroadcast( + SharedStorage &shared_storage, ///< Shared storage object + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx ///< Id of thread within warp + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + BaseStreamK(thread_idx) + { } + + + /// Aggregates the accumulator sets shared by peer blocks in the global workspace, + /// performing epilogue computations, writing to output + MCTLASS_DEVICE + void reduce( + int peer_idx_begin, + int peer_idx_end, + int reduce_fragment_idx, + void *element_workspace, + OutputOp const &output_op, ///< Output operator + ElementVector const * broadcast_ptr, ///< Broadcast vector + OutputTileIterator destination_iterator, ///< Tile iterator for destination + OutputTileIterator source_iterator, ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + TensorTileIterator tensor_iterator, ///< Threadblock tile iterator for additional tensor operand + MatrixCoord const &problem_size = ///< Problem size needed to guard against out-of-bounds accesses + MatrixCoord(Shape::kM, Shape::kN), + MatrixCoord const &threadblock_offset = ///< Threadblock's initial offset within the problem size space + MatrixCoord()) + { + // Reduce peer accumulator fragments into one fragment + AccumulatorFragment accum_fragment; + BaseStreamK::reduce(accum_fragment, peer_idx_begin, peer_idx_end, reduce_fragment_idx, element_workspace); + + // Store fragment to shared memory + this->warp_tile_iterator_.store(accum_fragment); + + __syncthreads(); + + Base::reduce(reduce_fragment_idx, output_op, broadcast_ptr, destination_iterator, source_iterator, tensor_iterator, problem_size, threadblock_offset); + + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_visitor_with_softmax.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_visitor_with_softmax.h new file mode 100644 index 0000000..43b6009 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_visitor_with_softmax.h @@ -0,0 +1,513 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue visitor for threadblock scoped GEMMs that process softmax computations in epilogue. + + The epilogue finds max values in each row of the row-major output matrix and stores them. + The max values are also used for a further round of threadblock scoped reduction operation, where + the partial reduction results are stored in a pre-allocated array and used for further full reduction. + +*/ + +#pragma once + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "mctlass/mctlass.h" +#include "mctlass/arch/memory.h" +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/fast_math.h" + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +template < + typename ThreadblockShape_, + int ThreadCount, + typename OutputTileIterator_, + typename ElementAccumulator_, + typename ElementNorm_, + typename ElementSum_, + typename ElementSoftmaxCompute_, + typename ElementwiseFunctor_, + bool UseMasking_ = false +> +class EpilogueVisitorSoftmax { +public: + + using ThreadblockShape = ThreadblockShape_; + static int const kThreadCount = ThreadCount; + + using OutputTileIterator = OutputTileIterator_; + using ElementwiseFunctor = ElementwiseFunctor_; + + static int const kIterations = OutputTileIterator::kIterations; + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + using ElementOutput = typename OutputTileIterator::Element; + using LayoutOutput = mctlass::layout::RowMajor; + using ElementAccumulator = ElementAccumulator_; + + using ElementNorm = ElementNorm_; + using ElementSum = ElementSum_; + using ElementSoftmaxCompute = ElementSoftmaxCompute_; + + using AccumulatorFragment = Array; + using SoftmaxFragment = Array; + using OutputVector = Array; + using TensorRefD = TensorRef; + + static int const kThreadsPerRow = OutputTileIterator::ThreadMap::Detail::kAccessWidth; + static bool const kHasMultiStepsInRow = (OutputTileIterator::ThreadMap::Iterations::kColumn > 1); + static bool const kUseMasking = UseMasking_; + + /// Argument structure + struct Arguments { + + typename ElementwiseFunctor::Params elementwise; + int64_t batch_stride_C; + int64_t batch_stride_D; + int64_t batch_stride_Max; + int64_t batch_stride_Sum; + + // + // Methods + // + Arguments(): + batch_stride_C(0), + batch_stride_D(0), + batch_stride_Max(0), + batch_stride_Sum(0) + { + + } + + Arguments( + typename ElementwiseFunctor::Params elementwise_ + ): + elementwise(elementwise_), + batch_stride_C(0), + batch_stride_D(0), + batch_stride_Max(0), + batch_stride_Sum(0) + { + + } + + Arguments( + typename ElementwiseFunctor::Params elementwise_, + int64_t batch_stride_C_, + int64_t batch_stride_D_, + int64_t batch_stride_Max_, + int64_t batch_stride_Sum_ + ): + elementwise(elementwise_), + batch_stride_C(batch_stride_C_), + batch_stride_D(batch_stride_D_), + batch_stride_Max(batch_stride_Max_), + batch_stride_Sum(batch_stride_Sum_) + { + + } + + }; + + struct Params { + + typename ElementwiseFunctor::Params elementwise; + int64_t batch_stride_C; + int64_t batch_stride_D; + int64_t batch_stride_Max; + int64_t batch_stride_Sum; + // + // Methods + // + MCTLASS_HOST_DEVICE + Params() + { + + } + + MCTLASS_HOST_DEVICE + Params(Arguments const &args): + elementwise(args.elementwise), + batch_stride_C(args.batch_stride_C), + batch_stride_D(args.batch_stride_D), + batch_stride_Max(args.batch_stride_Max), + batch_stride_Sum(args.batch_stride_Sum) + { + + } + }; + + /// Shared storage + struct SharedStorage { + + }; + +private: + + Params const & params_; + SharedStorage & shared_storage_; + MatrixCoord extent_; + MatrixCoord extent_real_; + ElementwiseFunctor elementwise_; + + OutputTileIterator iterator_C_; + OutputTileIterator iterator_D_; + typename OutputTileIterator::Fragment fragment_C_; + typename OutputTileIterator::Fragment fragment_D_; + + ElementAccumulator alpha_; + ElementAccumulator beta_; + + ElementNorm *ptr_Max_; + ElementSum *ptr_Sum_; + + int column_offset_; + + ElementSoftmaxCompute accum_max_; + ElementSoftmaxCompute accum_sum_; + + MatrixCoord thread_offset_; + + float infinity_; + +public: + + MCTLASS_DEVICE + EpilogueVisitorSoftmax( + Params const ¶ms, + SharedStorage &shared_storage, + mctlass::MatrixCoord const &problem_size, + int thread_idx, + int warp_idx, + int lane_idx, + typename OutputTileIterator::Params params_C, + typename OutputTileIterator::Params params_D, + typename OutputTileIterator::Element *ptr_C, + typename OutputTileIterator::Element *ptr_D, + ElementNorm *ptr_Max = nullptr, + ElementSum *ptr_Sum = nullptr, + mctlass::MatrixCoord const &threadblock_offset = mctlass::MatrixCoord(0, 0), + int column_offset = 0, + mctlass::MatrixCoord const &problem_size_real = mctlass::MatrixCoord(0, 0), + float infinity = 10000.0f + ): + params_(params), + shared_storage_(shared_storage), + extent_(problem_size), + elementwise_(params.elementwise), + iterator_C_(params_C, ptr_C, problem_size, thread_idx, threadblock_offset), + iterator_D_(params_D, ptr_D, problem_size, thread_idx, threadblock_offset), + ptr_Max_(ptr_Max), + ptr_Sum_(ptr_Sum), + column_offset_(column_offset), + extent_real_(problem_size_real), + infinity_(infinity) + { + alpha_ = (params.elementwise.alpha_ptr ? *params.elementwise.alpha_ptr : params.elementwise.alpha); + beta_ = (params.elementwise.beta_ptr ? *params.elementwise.beta_ptr : params.elementwise.beta); + + if (beta_ == ElementAccumulator()) { + iterator_C_.clear_mask(); + } + } + + /// Helper to indicate split-K behavior + MCTLASS_DEVICE + void set_k_partition( + int split_k_index, ///< Index of this threadblock within split-K partitioned scheme + int split_k_slices) { ///< Total number of split-K slices + + } + + /// Called to set the batch index + MCTLASS_DEVICE + void set_batch_index(int batch_idx) { + iterator_C_.add_pointer_offset(batch_idx * params_.batch_stride_C); + iterator_D_.add_pointer_offset(batch_idx * params_.batch_stride_D); + } + + /// Called at the start of the epilogue just before iterating over accumulator slices + MCTLASS_DEVICE + void begin_epilogue() { + + } + + /// Called at the start of one step before starting accumulator exchange + MCTLASS_DEVICE + void begin_step(int step_idx) { + fragment_D_.clear(); + fragment_C_.clear(); + + if (elementwise_.kScale != mctlass::epilogue::thread::ScaleType::OnlyAlphaScaling) { + iterator_C_.load(fragment_C_); + ++iterator_C_; + } + + } + + /// Called at the start of a row + MCTLASS_DEVICE + void begin_row(int row_idx) { + // Clear accumulators for max and sum when starting a whole row + clear_accum_(); + + } + + /// Called after accumulators have been exchanged for each accumulator vector + MCTLASS_DEVICE + void visit( + int iter_idx, + int row_idx, + int column_idx, + int frag_idx, + AccumulatorFragment const &accum) { + + using Mul = mctlass::multiplies; + using Minus = mctlass::minus; + using Exp = mctlass::fast_exp_op; + + Minus minus; + Exp exponential; + + SoftmaxFragment result; + + NumericArrayConverter source_converter; + OutputVector &source_vector = reinterpret_cast(&fragment_C_)[frag_idx]; + + if (elementwise_.kScale == mctlass::epilogue::thread::ScaleType::OnlyAlphaScaling) { + result = source_converter(elementwise_(accum)); + }else{ + result = source_converter(elementwise_(accum, source_vector)); + } + + thread_offset_ = + iterator_D_.thread_start() + + OutputTileIterator::ThreadMap::iteration_offset(frag_idx); + + bool column_guard = (thread_offset_.column() < extent_.column()); + + if (kUseMasking) { + int elements_in_boundary = extent_real_.column() - thread_offset_.column(); + elements_in_boundary = (elements_in_boundary > kElementsPerAccess) ? kElementsPerAccess : elements_in_boundary; + elementwise_padding_(result, elements_in_boundary); + } + + ElementSoftmaxCompute accum_max_prev = accum_max_; + + // Compute the maximum within one row + if (!column_idx) { + // This is the first fragment in a new row + if (column_guard) { + accum_max_ = maximum_accumulator_(result); + } + } + else { + // This is an additional fragment in the same row + if (column_guard) { + accum_max_ = maximum_accumulator_(result, accum_max_); + } + } + + // proactively compute max in warps + accum_max_ = warp_reduce_max_(accum_max_); + + ElementSoftmaxCompute updater = fast_exp(accum_max_prev - accum_max_); + + SoftmaxFragment intermediate = exponential(minus(result, accum_max_)); + + if (kHasMultiStepsInRow) { + if (!column_idx) { + accum_sum_ = (column_guard) ? \ + sum_accumulator_(intermediate) : ElementSoftmaxCompute(0); + } else { + // Algorithm in $3.1, https://arxiv.org/pdf/2205.14135v1.pdf + // S* = S* x updater + sum_row(P'), where updater = exp(M* - M_row) + accum_sum_ = (column_guard) ? \ + sum_accumulator_(intermediate, accum_sum_ * updater) : accum_sum_ * updater; + } + } else { + accum_sum_ = (column_guard) ? sum_accumulator_(intermediate, accum_sum_) : ElementSoftmaxCompute(0); + } + + // Convert to the output + NumericArrayConverter output_converter; + OutputVector &output = reinterpret_cast(&fragment_D_)[frag_idx]; + output = output_converter(result); + } + + /// Called at the end of a row + MCTLASS_DEVICE + void end_row(int row_idx) { + + using ConvertSumOutput = mctlass::NumericConverter; + using ConvertNormOutput = mctlass::NumericConverter; + + ConvertSumOutput convert_sum_output; + ConvertNormOutput convert_norm_output; + + // Compute accumulate sum only in the last step + accum_sum_ = warp_reduce_sum_(accum_sum_); + + bool is_first_thread_in_tile = ((threadIdx.x % kThreadsPerRow) == 0); + bool row_guard = thread_offset_.row() < extent_.row(); + bool is_write_thread = row_guard && is_first_thread_in_tile; + + int block_batch = blockIdx.z; + + ElementNorm *curr_ptr_max = ptr_Max_ + thread_offset_.row() + column_offset_ + block_batch * params_.batch_stride_Max; + ElementSum *curr_ptr_sum = ptr_Sum_ + thread_offset_.row() + column_offset_ + block_batch * params_.batch_stride_Sum; + + arch::global_store( + convert_norm_output(accum_max_), + (void *)curr_ptr_max, + is_write_thread); + + arch::global_store( + convert_sum_output(accum_sum_), + (void *)curr_ptr_sum, + is_write_thread); + + // Clear accumulators for max and sum when finishing a whole row + clear_accum_(); + + } + + /// Called after all accumulator elements have been visited + MCTLASS_DEVICE + void end_step(int step_idx) { + + iterator_D_.store(fragment_D_); + ++iterator_D_; + } + + /// Called after all steps have been completed + MCTLASS_DEVICE + void end_epilogue() { + + } + +private: + + MCTLASS_DEVICE + void elementwise_padding_(SoftmaxFragment &result, int elements_in_boundary) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < SoftmaxFragment::kElements; ++i) { + result[i] = (i < elements_in_boundary) ? result[i] : ElementSoftmaxCompute(-infinity_); + } + } + + MCTLASS_DEVICE + ElementSoftmaxCompute warp_reduce_sum_(ElementSoftmaxCompute sum_) { + int half_thread_in_row = (kThreadsPerRow >> 1); + MCTLASS_PRAGMA_UNROLL + for (int i = half_thread_in_row; i > 0; i >>= 1) { + ElementSoftmaxCompute tmp = __shfl_xor_sync(0xFFFFFFFF, sum_, i); + sum_ += tmp; + } + return sum_; + } + + MCTLASS_DEVICE + ElementSoftmaxCompute warp_reduce_max_(ElementSoftmaxCompute max_) { + int half_thread_in_row = (kThreadsPerRow >> 1); + MCTLASS_PRAGMA_UNROLL + for (int i = half_thread_in_row; i > 0; i >>= 1) { + ElementSoftmaxCompute tmp = __shfl_xor_sync(0xFFFFFFFF, max_, i); + max_ = fast_max(max_, tmp); + } + return max_; + } + + MCTLASS_DEVICE + void clear_accum_() { + + uint32_t float_max_bits = 0xff7fffff; // -FLT_MAX + float min_float = reinterpret_cast(float_max_bits); + accum_max_ = ElementSoftmaxCompute(min_float); + accum_sum_ = ElementSoftmaxCompute(0); + } + + MCTLASS_DEVICE + ElementSoftmaxCompute sum_accumulator_(SoftmaxFragment const &accum) { + ElementSoftmaxCompute sum_ = ElementSoftmaxCompute(0); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < SoftmaxFragment::kElements; ++i) { + sum_ += ElementSoftmaxCompute(accum[i]); + } + + return sum_; + } + + MCTLASS_DEVICE + ElementSoftmaxCompute sum_accumulator_(SoftmaxFragment const &accum, ElementSoftmaxCompute sum_) { + // ElementSoftmaxCompute sum_ = ElementSoftmaxCompute(0); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < SoftmaxFragment::kElements; ++i) { + sum_ += ElementSoftmaxCompute(accum[i]); + } + + return sum_; + } + + MCTLASS_DEVICE + ElementSoftmaxCompute maximum_accumulator_(SoftmaxFragment const &accum) { + ElementSoftmaxCompute max_ = accum[0]; + + MCTLASS_PRAGMA_UNROLL + for (int i = 1; i < SoftmaxFragment::kElements; ++i) { + max_ = fast_max(max_, ElementSoftmaxCompute(accum[i])); + } + + return max_; + } + + MCTLASS_DEVICE + ElementSoftmaxCompute maximum_accumulator_(SoftmaxFragment const &accum, ElementSoftmaxCompute max_) { + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < SoftmaxFragment::kElements; ++i) { + max_ = fast_max(max_, ElementSoftmaxCompute(accum[i])); + } + + return max_; + } +}; + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_with_broadcast.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_with_broadcast.h new file mode 100644 index 0000000..b0aa588 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_with_broadcast.h @@ -0,0 +1,1778 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#include +#else +#include +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/tensor_coord.h" +#include "mctlass/aligned_buffer.h" +#include "mctlass/functional.h" +#include "mctlass/fast_math.h" +#include "mctlass/layout/vector.h" +#include "mctlass/layout/tensor.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_iterator.h" + +#include "mctlass/epilogue/threadblock/epilogue_base.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator.h" + +#include "mctlass/numeric_types.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// This base class is meant to define the concept required of the +/// EpilogueWithBroadcast::OutputOp +template < + typename ElementC_, + typename ElementAccumulator_, + typename ElementCompute_, + typename ElementZ_, + typename ElementT_, + int ElementsPerAccess, + bool StoreZ = true, + bool StoreT = true +> +struct EpilogueWithBroadcastOpBase { + + using ElementOutput = ElementC_; + using ElementAccumulator = ElementAccumulator_; + using ElementCompute = ElementCompute_; + using ElementZ = ElementZ_; + using ElementT = ElementT_; + static int const kElementsPerAccess = ElementsPerAccess; + + using FragmentAccumulator = Array; + using FragmentCompute = Array; + using FragmentC = Array; + using FragmentZ = Array; + using FragmentT = Array; + + /// If true, the 'Z' tensor is stored + static bool const kStoreZ = StoreZ; + + /// If true, the 'T' tensor is stored + static bool const kStoreT = StoreT; + + /// Parameters structure - required + struct Params { }; + + // + // Methods + // + + /// Constructor from Params + EpilogueWithBroadcastOpBase(Params const ¶ms_) { } + + /// Determine if the source is needed. May return false if + bool is_source_needed() const { + return true; + } + + MCTLASS_HOST_DEVICE + void set_k_partition(int k_partition, int k_partition_count) { } + + /// Applies the operation when is_source_needed() is true + MCTLASS_HOST_DEVICE + void operator()( + FragmentZ &frag_Z, + FragmentT &frag_T, + FragmentAccumulator const &AB, + FragmentC const &frag_C1, + FragmentC const &frag_C2, + FragmentCompute const &V) const { + + } + + /// Applies the operation when is_source_needed() is false + MCTLASS_HOST_DEVICE + void operator()( + FragmentZ &frag_Z, + FragmentT &frag_T, + FragmentAccumulator const &AB, + FragmentCompute const &V) const { + + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator with bias vector broadcast over columns. +/// +/// Computes the following: +/// +/// +/// Z, T = OutputOp(AB, C, Broadcast) +/// +/// if (ElementwiseOp::kStoreZ) { +/// store(converted_u); +/// } +/// +/// if (ElementwiseOp::kStoreT) { +/// store(v); +/// } +/// +template < + typename Shape_, ///< Shape of threadblock tile (concept: GemmShape) + typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp) + int PartitionsK, ///< Number of partitions of the K dimension + typename OutputTileIterator_, ///< Tile iterator reading and writing output tensors (z) + typename TensorTileIterator_, ///< Additional tile iterator for tensor-valued operands (t) + typename ElementVector_, ///< Pointer to broadcast vector + typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators + typename WarpTileIterator_, ///< Warp-scoped tile iterator writing accumulators to SMEM + typename SharedLoadIterator_, ///< Threadblock-scoped tile iterator loading from SMEM + typename OutputOp_, ///< Output operator - concept is EpilogueWithBroadcastOp + typename Padding_, ///< Padding added to SMEM allocation to avoid bank conflicts (concept: MatrixShape) + int FragmentsPerPartition = 1, ///< Used to coarsten the epilogue granularity + int IterationsUnroll = ///< Used to reduce binary size when epilogue op is large + (!IsEpilogueFunctorHeavy::value), + bool IsSingleSource = OutputOp_::kIsSingleSource +> +class EpilogueWithBroadcast; + +template < + typename Shape_, + typename WarpMmaOperator_, + int PartitionsK, + typename OutputTileIterator_, + typename TensorTileIterator_, + typename ElementVector_, + typename AccumulatorFragmentIterator_, + typename WarpTileIterator_, + typename SharedLoadIterator_, + typename OutputOp_, + typename Padding_, + int FragmentsPerPartition, + int IterationsUnroll +> +class EpilogueWithBroadcast< + Shape_, + WarpMmaOperator_, + PartitionsK, + OutputTileIterator_, + TensorTileIterator_, + ElementVector_, + AccumulatorFragmentIterator_, + WarpTileIterator_, + SharedLoadIterator_, + OutputOp_, + Padding_, + FragmentsPerPartition, + IterationsUnroll, + false +> : + public EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_, + FragmentsPerPartition> { + +public: + + using Base = EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_, + FragmentsPerPartition>; + + static bool const kIsSingleSource = false; + using Shape = Shape_; + using WarpMmaOperator = WarpMmaOperator_; + static int const kPartitionsK = PartitionsK; + using OutputTileIterator = OutputTileIterator_; + using TensorTileIterator = TensorTileIterator_; + using ElementVector = ElementVector_; + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using WarpTileIterator = WarpTileIterator_; + using SharedLoadIterator = SharedLoadIterator_; + using OutputOp = OutputOp_; + using Padding = Padding_; + + using Layout = layout::RowMajor; + using LongIndex = typename Layout::LongIndex; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename Base::AccumulatorTile; + + /// Accumulator element + using ElementAccumulator = typename WarpTileIterator::Element; + + /// Compute data type produced by the output op + using ElementCompute = typename OutputOp::ElementCompute; + + /// Compute fragment + using FragmentCompute = Array; + + /// Thread map used by output tile iterators + using ThreadMap = typename OutputTileIterator::ThreadMap; + + /// Fragment object used to store the broadcast values + using BroadcastFragment = Array< + ElementCompute, + ThreadMap::Iterations::kColumn * ThreadMap::kElementsPerAccess>; + + /// Output element + using ElementOutput = typename OutputTileIterator::Element; + + /// Data type of additional tensor + using ElementTensor = typename TensorTileIterator::Element; + + /// Output access size + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + /// Tensor reference to destination tensor + using TensorRef = typename OutputTileIterator::TensorRef; + + /// Tensor reference to sync tensor + using SyncTensorRef = typename mctlass::TensorRef; + + /// Const tensor reference to source tensor + using ConstTensorRef = typename OutputTileIterator::ConstTensorRef; + + /// Array type used to output + using OutputAccessType = Array< + typename OutputTileIterator::Element, OutputTileIterator::kElementsPerAccess>; + + /// Array type used by output functor + using AccumulatorAccessType = Array; + + /// Array type used by output functor + using ComputeAccessType = Array; + + /// Tensor access type + using TensorAccessType = Array; + + /// Number of warps + using WarpCount = typename Base::WarpCount; + + /// Shared memory allocation from epilogue base class + using BaseSharedStorage = typename Base::SharedStorage; + + static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1 ? Base::kFragmentsPerIteration : kPartitionsK; + static int constexpr kSmemPointerOffset = Base::SharedStorage::StorageShape::kCount / kSmemTiles; + + /// Used for the broadcast + struct BroadcastDetail { + + /// Number of threads per warp + static int const kWarpSize = 32; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + + /// Number of distinct scalar column indices handled by each thread + static int const kColumnsPerThread = ThreadMap::Iterations::kColumn * ThreadMap::kElementsPerAccess; + + /// Number of distinct scalar row indices handled by each thread + static int const kRowsPerThread = ThreadMap::Iterations::kCount / ThreadMap::Iterations::kColumn; + + /// Number of threads per threadblock + static int const kThreadCount = kWarpSize * WarpCount::kCount; + + /// Number of distinct threads per row of output tile + static int const kThreadsPerRow = (Shape::kN / kColumnsPerThread); + + /// Number of distinct threads which must be reduced during the final reduction phase within the threadblock. + static int const kThreadRows = kThreadCount / kThreadsPerRow; + + /// I'm not sure what I meant here. + static int const kThreadAccessesPerRow = const_max(1, (Shape::kN + kThreadCount - 1) / kThreadCount); + + /// Shape of the shared memory allocation for the epilogue + using StorageShape = MatrixShape< + kThreadRows, + Shape::kN + >; + + /// Debug printing + MCTLASS_DEVICE + static void print() { +#if 0 + printf("BroadcastDetail {\n"); + printf( + " kColumnsPerThread: %d\nkRowsPerThread: %d\n,kThreadCount: %d\nkThreadsPerRow: %d\n" + "kThreadRows: %d\nThreadAccessesPerRow: %d\nStorageShape: %d x %d (count: %d)\n", + kColumnsPerThread, + kRowsPerThread, + kThreadCount, + kThreadsPerRow, + kThreadRows, + kThreadAccessesPerRow, + StorageShape::kRow, + StorageShape::kColumn, + StorageShape::kCount + ); + printf("};\n"); +#endif + } + }; + + /// Shared storage structure (shadows base) with additional SMEM buffer for reduction + struct SharedStorage { + union { + BaseSharedStorage base; + }; + + MCTLASS_HOST_DEVICE + SharedStorage() { } + }; + +public: + + + static_assert(SharedLoadIterator::Fragment::kElements == OutputTileIterator::Fragment::kElements, + "Mismatch between shared load iterator and output tile iterator."); + + static_assert(OutputTileIterator::kElementsPerAccess, "OutputTileIterator::kElementsPerAccess must not be zero."); + + static_assert(!(OutputTileIterator::Fragment::kElements % OutputTileIterator::kElementsPerAccess), + "Divisibility"); + +private: + + /// Loads fragment from shared memory aligned with output tensor + SharedLoadIterator shared_load_iterator_; + + /// Thread index within the threadblock + int thread_idx_; + +public: + + /// Constructor + MCTLASS_DEVICE + EpilogueWithBroadcast( + SharedStorage &shared_storage, ///< Shared storage object + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx ///< Id of thread within warp + ): + Base(shared_storage.base, thread_idx, warp_idx, lane_idx), + shared_load_iterator_(shared_storage.base.reference(), thread_idx), + thread_idx_(thread_idx) + { + + } + + /// Streams the result to global memory + MCTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + ElementVector const * broadcast_ptr, ///< Broadcast vector + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator1, ///< Tile iterator for first source accumulator matrix + OutputTileIterator source_iterator2, ///< Tile iterator for second source accumulator matrix + TensorTileIterator tensor_iterator, ///< Threadblock tile iterator for additional tensor operand + MatrixCoord const &problem_size = ///< Problem size needed to guard against out-of-bounds accesses + MatrixCoord(Shape::kM, Shape::kN), + MatrixCoord const &threadblock_offset = ///< Threadblock's initial offset within the problem size space + MatrixCoord()) { + + BroadcastFragment broadcast_fragment; + + load_broadcast_fragment_(broadcast_fragment, broadcast_ptr, problem_size, threadblock_offset); + + using Element = typename AccumulatorTile::Element; + const int lane_id = __lane_id(); + const int lane_delta = 2 * (lane_id % 4) + lane_id / 16 * 16 - lane_id; + for (int i = 0; i < accumulators.size() / 4; i++) { + + Element x = accumulators[4 * i + 0]; + Element y = accumulators[4 * i + 1]; + Element z = accumulators[4 * i + 2]; + Element w = accumulators[4 * i + 3]; + + Element x0 = shfl_down_sync(UINT64_MAX, x, lane_delta); + Element y0 = shfl_down_sync(UINT64_MAX, y, lane_delta); + Element z0 = shfl_down_sync(UINT64_MAX, z, lane_delta); + Element w0 = shfl_down_sync(UINT64_MAX, w, lane_delta); + + Element x1 = shfl_down_sync(UINT64_MAX, x, lane_delta + 1); + Element y1 = shfl_down_sync(UINT64_MAX, y, lane_delta + 1); + Element z1 = shfl_down_sync(UINT64_MAX, z, lane_delta + 1); + Element w1 = shfl_down_sync(UINT64_MAX, w, lane_delta + 1); + + Element x2 = shfl_down_sync(UINT64_MAX, x, lane_delta + 32); + Element y2 = shfl_down_sync(UINT64_MAX, y, lane_delta + 32); + Element z2 = shfl_down_sync(UINT64_MAX, z, lane_delta + 32); + Element w2 = shfl_down_sync(UINT64_MAX, w, lane_delta + 32); + + Element x3 = shfl_down_sync(UINT64_MAX, x, lane_delta + 33); + Element y3 = shfl_down_sync(UINT64_MAX, y, lane_delta + 33); + Element z3 = shfl_down_sync(UINT64_MAX, z, lane_delta + 33); + Element w3 = shfl_down_sync(UINT64_MAX, w, lane_delta + 33); + + Element dst_x = 0, dst_y = 0, dst_z = 0, dst_w = 0; + switch ((lane_id % 16) / 4) { + case 0: + dst_x = x0, dst_y = x1, dst_z = x2, dst_w = x3; + break; + case 1: + dst_x = y0, dst_y = y1, dst_z = y2, dst_w = y3; + break; + case 2: + dst_x = z0, dst_y = z1, dst_z = z2, dst_w = z3; + break; + case 3: + dst_x = w0, dst_y = w1, dst_z = w2, dst_w = w3; + break; + default: + break; + } + + accumulators[4 * i + 0] = dst_x; + accumulators[4 * i + 1] = dst_y; + accumulators[4 * i + 2] = dst_z; + accumulators[4 * i + 3] = dst_w; + + } + + if (!output_op.is_source_needed()) { + compute_source_not_needed_( + output_op, + broadcast_fragment, + destination_iterator, + accumulators, + tensor_iterator); + } + else { + compute_source_needed_( + output_op, + broadcast_fragment, + destination_iterator, + accumulators, + source_iterator1, + source_iterator2, + tensor_iterator); + } + } + +private: + + MCTLASS_DEVICE + void load_broadcast_fragment_( + BroadcastFragment & broadcast_fragment, ///< Fragment containing the accumulated partial reduction over columns + ElementVector const * broadcast_ptr, ///< Broadcast vector + MatrixCoord const &problem_size, ///< Problem size needed to guard against out-of-bounds accesses + MatrixCoord const &threadblock_offset ///< Threadblock's initial offset within the problem size space + ) { + + broadcast_fragment.clear(); + + // If no pointer is supplied, set with all zeros and avoid memory accesses + if (!broadcast_ptr) { + return; + } + + int thread_initial_column = ThreadMap::initial_offset(thread_idx_).column(); + + int thread_column_idx = threadblock_offset.column() + thread_initial_column; + broadcast_ptr += thread_initial_column; + + NumericArrayConverter converter; + using AccessType = AlignedArray; + using ComputeFragmentType = Array; + + ComputeFragmentType *frag_ptr = reinterpret_cast(&broadcast_fragment); + + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < ThreadMap::Iterations::kColumn; ++j) { + + AccessType loaded; + + loaded.clear(); + + if (thread_column_idx < problem_size.column()) { + loaded = *reinterpret_cast(broadcast_ptr); + } + + ComputeFragmentType cvt = converter(loaded); + frag_ptr[j] = cvt; + + thread_column_idx += ThreadMap::Delta::kColumn; + broadcast_ptr += ThreadMap::Delta::kColumn; + } + } + + template + struct acc2smem_source_not_needed; + + template + struct acc2smem_source_not_needed> { + template + MCTLASS_DEVICE static void helper(AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator &warp_tile_iterator) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { + ++accum_fragment_iterator; + } + + MCTLASS_PRAGMA_UNROLL + for (int p = 0; p < Base::kFragmentsPerIteration; ++p) { + typename AccumulatorFragmentIterator::Fragment accum_fragment; + + accum_fragment_iterator.load(accum_fragment); + ++accum_fragment_iterator; + + warp_tile_iterator.store(accum_fragment); + if (p < Base::kFragmentsPerIteration - 1) { + warp_tile_iterator.add_pointer_offset(kSmemPointerOffset); + } + } + + if (Base::kFragmentsPerIteration > 1) { + warp_tile_iterator.add_pointer_offset(kSmemPointerOffset * + (1 - Base::kFragmentsPerIteration)); + } + } + + MCTLASS_DEVICE + static void push(size_t pos, + AccumulatorFragmentIterator const &iterator_begin, + WarpTileIterator &warp_tile_iterator) { + int dummy[] = { + (pos == (Seq * Base::kFragmentsPerIteration)) && + (helper(iterator_begin, warp_tile_iterator), 0)...}; + + MCTLASS_UNUSED(dummy[0]); + } + }; + + /// Streams the result to global memory + MCTLASS_DEVICE + void compute_source_not_needed_( + OutputOp const &output_op, ///< Output operator + BroadcastFragment const &broadcast_fragment, ///< Fragment containing the accumulated partial reduction over columns + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + TensorTileIterator tensor_iterator ///< Threadblock tile iterator for additioanl tensor operand + ) { + + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + + // MCTLASS_PRAGMA_UNROLL + //#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations / Base::kFragmentsPerIteration : 1) + //mxcc not support using unroll with parentheses, + //and we would support with unroll with OutputTileIterator::kIterations / Base::kFragmentsPerIteration someday + #pragma unroll 1 + for (int iter = 0; iter < OutputTileIterator::kIterations; iter += Base::kFragmentsPerIteration) { + + // + // Convert and store fragment + // + + + __syncthreads(); + + acc2smem_source_not_needed< + mctlass::make_index_sequence>::push(iter, + accum_fragment_iterator, + this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + MCTLASS_PRAGMA_UNROLL + for (int p = 0; p < Base::kFragmentsPerIteration; ++p) { + + + typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment[0]); + + if (p < Base::kFragmentsPerIteration - 1) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + } + else if (kPartitionsK > 1) { + + plus add_fragments; + + MCTLASS_PRAGMA_UNROLL + for ( int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + + shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * kSmemPointerOffset); + } + + // + // Apply output operation + // + + typename OutputTileIterator::Fragment frag_Z; + typename TensorTileIterator::Fragment frag_T; + + apply_output_operator_source_not_needed_( + frag_Z, + frag_T, + output_op, + aligned_accum_fragment[0], + broadcast_fragment); + + // + // Conditionally store fragments + // + + if (OutputOp::kStoreZ) { + destination_iterator.store(frag_Z); + ++destination_iterator; + } + + if (OutputOp::kStoreT) { + tensor_iterator.store(frag_T); + ++tensor_iterator; + } + } + + if (Base::kFragmentsPerIteration > 1) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset * (1 - Base::kFragmentsPerIteration)); + } + } + } + + + template + struct acc2smem_source_needed; + + template + struct acc2smem_source_needed> { + template + MCTLASS_DEVICE + static void helper(AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator &warp_tile_iterator) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { + ++accum_fragment_iterator; + } + + typename AccumulatorFragmentIterator::Fragment accum_fragment; + accum_fragment_iterator.load(accum_fragment); + warp_tile_iterator.store(accum_fragment); + } + + MCTLASS_DEVICE + static void push(size_t pos, + AccumulatorFragmentIterator const &iterator_begin, + WarpTileIterator &warp_tile_iterator) { + int dummy[] = {(pos == Seq) && (helper(iterator_begin, warp_tile_iterator), 0)...}; + } + }; + + + /// Streams the result to global memory + MCTLASS_DEVICE + void compute_source_needed_( + OutputOp const &output_op, ///< Output operator + BroadcastFragment const &broadcast_fragment, ///< Fragment containing the accumulated partial reduction over columns + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator1, ///< Tile iterator for first source accumulator matrix + OutputTileIterator source_iterator2, ///< Tile iterator for second source accumulator matrix + TensorTileIterator tensor_iterator ///< Threadblock tile iterator for additioanl tensor operand + ) { + + typename OutputTileIterator::Fragment source_fragment1; + source_fragment1.clear(); + typename OutputTileIterator::Fragment source_fragment2; + source_fragment2.clear(); + + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + + #pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1) + for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) { + + // + // Load the source + // + + source_iterator1.load(source_fragment1); + ++source_iterator1; + + source_iterator2.load(source_fragment2); + ++source_iterator2; + + // + // Convert and store fragment + // + + __syncthreads(); + + acc2smem_source_needed>::push( + iter, accum_fragment_iterator, this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment[0]); + + // If the number of k-slices is > 1 - perform a reduction amongst the k-slices + if (kPartitionsK > 1) + { + plus add_fragments; + const int tile_row_offset = Base::SharedStorage::StorageShape::kRow / PartitionsK; + + MCTLASS_PRAGMA_UNROLL + for ( int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_tile_offset({tile_row_offset , 0}); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + + shared_load_iterator_.add_tile_offset({-1 * (kPartitionsK-1) * tile_row_offset, 0}); + } + + // + // Apply output operation + // + + typename OutputTileIterator::Fragment frag_Z; + typename TensorTileIterator::Fragment frag_T; + + apply_output_operator_( + frag_Z, + frag_T, + output_op, + aligned_accum_fragment[0], + source_fragment1, + source_fragment2, + broadcast_fragment); + + // + // Conditionally store fragments + // + + if (OutputOp::kStoreZ) { + destination_iterator.store(frag_Z); + ++destination_iterator; + } + + if (OutputOp::kStoreT) { + tensor_iterator.store(frag_T); + ++tensor_iterator; + } + } + } + + /// Helper to invoke the output functor over each vector of output + MCTLASS_DEVICE + void apply_output_operator_( + typename OutputTileIterator::Fragment &frag_Z, + typename TensorTileIterator::Fragment &frag_T, + OutputOp const &output_op, + typename SharedLoadIterator::Fragment const &frag_AB, + typename OutputTileIterator::Fragment const &frag_C1, + typename OutputTileIterator::Fragment const &frag_C2, + BroadcastFragment const &frag_Broadcast) { + + using AccessTypeZ = Array; + using AccessTypeT = Array; + using AccessTypeBroadcast = Array; + + AccessTypeZ *frag_Z_ptr = reinterpret_cast(&frag_Z); + AccessTypeT *frag_T_ptr = reinterpret_cast(&frag_T); + + AccumulatorAccessType const *frag_AB_ptr = + reinterpret_cast(&frag_AB); + + OutputAccessType const *frag_C1_ptr = + reinterpret_cast(&frag_C1); + + OutputAccessType const *frag_C2_ptr = + reinterpret_cast(&frag_C2); + + AccessTypeBroadcast const *frag_Broadcast_ptr = + reinterpret_cast(&frag_Broadcast); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + output_op( + frag_Z_ptr[i], + frag_T_ptr[i], + frag_AB_ptr[i], + frag_C1_ptr[i], + frag_C2_ptr[i], + frag_Broadcast_ptr[i % ThreadMap::Iterations::kColumn]); + } + } + + /// Helper to invoke the output functor over each vector of output + MCTLASS_DEVICE + void apply_output_operator_source_not_needed_( + typename OutputTileIterator::Fragment &frag_Z, + typename TensorTileIterator::Fragment &frag_T, + OutputOp const &output_op, + typename SharedLoadIterator::Fragment const &frag_AB, + BroadcastFragment const &frag_Broadcast) { + + using AccessTypeZ = Array; + using AccessTypeT = Array; + using AccessTypeBroadcast = Array; + + AccessTypeZ *frag_Z_ptr = reinterpret_cast(&frag_Z); + AccessTypeT *frag_T_ptr = reinterpret_cast(&frag_T); + + AccumulatorAccessType const *frag_AB_ptr = + reinterpret_cast(&frag_AB); + + AccessTypeBroadcast const *frag_Broadcast_ptr = + reinterpret_cast(&frag_Broadcast); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + + output_op( + frag_Z_ptr[i], + frag_T_ptr[i], + frag_AB_ptr[i], + frag_Broadcast_ptr[i % ThreadMap::Iterations::kColumn]); + } + } + + public: + /// Stream-K reduce helper + MCTLASS_DEVICE + void reduce( + int reduce_fragment_idx, ///< Reduce fragment index + OutputOp const &output_op, ///< Output operator + ElementVector const * broadcast_ptr, ///< Broadcast vector + OutputTileIterator destination_iterator, ///< Tile iterator for destination + OutputTileIterator source_iterator1, ///< Tile iterator for first source accumulator matrix + OutputTileIterator source_iterator2, ///< Tile iterator for second source accumulator matrix + TensorTileIterator tensor_iterator, ///< Threadblock tile iterator for additional tensor operand + MatrixCoord const &problem_size = ///< Problem size needed to guard against out-of-bounds accesses + MatrixCoord(Shape::kM, Shape::kN), + MatrixCoord const &threadblock_offset = ///< Threadblock's initial offset within the problem size space + MatrixCoord()) + { + + BroadcastFragment broadcast_fragment; + load_broadcast_fragment_(broadcast_fragment, broadcast_ptr, problem_size, threadblock_offset); + + // Initialize/load source-fragment data + typename OutputTileIterator::Fragment source_fragment1; + source_fragment1.clear(); + typename OutputTileIterator::Fragment source_fragment2; + source_fragment2.clear(); + + if (output_op.is_source_needed()) + { + source_iterator1 += reduce_fragment_idx; + source_iterator1.load(source_fragment1); + + source_iterator2 += reduce_fragment_idx; + source_iterator2.load(source_fragment2); + } + + // Load fragment from shared memory + typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK]; + shared_load_iterator_.load(aligned_accum_fragment[0]); + + // Add fragments shared by other k partitions + if (kPartitionsK > 1) + { + plus add_fragments; + + MCTLASS_PRAGMA_UNROLL + for ( int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + } + + // + // Apply output operation + // + + typename OutputTileIterator::Fragment frag_Z; + typename TensorTileIterator::Fragment frag_T; + + if (!output_op.is_source_needed()) { + apply_output_operator_source_not_needed_( + frag_Z, + frag_T, + output_op, + aligned_accum_fragment[0], + broadcast_fragment); + } else { + apply_output_operator_( + frag_Z, + frag_T, + output_op, + aligned_accum_fragment[0], + source_fragment1, + source_fragment2, + broadcast_fragment); + } + + // + // Conditionally store fragments + // + + if (OutputOp::kStoreZ) { + destination_iterator.store(frag_Z); + ++destination_iterator; + } + + if (OutputOp::kStoreT) { + tensor_iterator.store(frag_T); + ++tensor_iterator; + } + } +}; + + +template < + typename Shape_, + typename WarpMmaOperator_, + int PartitionsK, + typename OutputTileIterator_, + typename TensorTileIterator_, + typename ElementVector_, + typename AccumulatorFragmentIterator_, + typename WarpTileIterator_, + typename SharedLoadIterator_, + typename OutputOp_, + typename Padding_, + int FragmentsPerPartition, + int IterationsUnroll +> +class EpilogueWithBroadcast< + Shape_, + WarpMmaOperator_, + PartitionsK, + OutputTileIterator_, + TensorTileIterator_, + ElementVector_, + AccumulatorFragmentIterator_, + WarpTileIterator_, + SharedLoadIterator_, + OutputOp_, + Padding_, + FragmentsPerPartition, + IterationsUnroll, + true +> : + public EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_, + FragmentsPerPartition> { + +public: + + using Base = EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_, + FragmentsPerPartition>; + + static bool const kIsSingleSource = true; + using Shape = Shape_; + using WarpMmaOperator = WarpMmaOperator_; + static int const kPartitionsK = PartitionsK; + using OutputTileIterator = OutputTileIterator_; + using TensorTileIterator = TensorTileIterator_; + using ElementVector = ElementVector_; + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using WarpTileIterator = WarpTileIterator_; + using SharedLoadIterator = SharedLoadIterator_; + using OutputOp = OutputOp_; + using Padding = Padding_; + + using Layout = layout::RowMajor; + using LongIndex = typename Layout::LongIndex; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename Base::AccumulatorTile; + + /// Accumulator element + using ElementAccumulator = typename WarpTileIterator::Element; + + /// Compute data type produced by the output op + using ElementCompute = typename OutputOp::ElementCompute; + + /// Compute fragment + using FragmentCompute = Array; + + /// Thread map used by output tile iterators + using ThreadMap = typename OutputTileIterator::ThreadMap; + + /// Fragment object used to store the broadcast values + using BroadcastFragment = Array< + ElementCompute, + ThreadMap::Iterations::kColumn * ThreadMap::kElementsPerAccess>; + + /// Output element + using ElementOutput = typename OutputTileIterator::Element; + + /// Data type of additional tensor + using ElementTensor = typename TensorTileIterator::Element; + + /// Output access size + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + /// Tensor reference to destination tensor + using TensorRef = typename OutputTileIterator::TensorRef; + + /// Tensor reference to sync tensor + using SyncTensorRef = typename mctlass::TensorRef; + + /// Const tensor reference to source tensor + using ConstTensorRef = typename OutputTileIterator::ConstTensorRef; + + /// Array type used to output + using OutputAccessType = Array< + typename OutputTileIterator::Element, OutputTileIterator::kElementsPerAccess>; + + /// Array type used by output functor + using AccumulatorAccessType = Array; + + /// Array type used by output functor + using ComputeAccessType = Array; + + /// Tensor access type + using TensorAccessType = Array; + + /// Number of warps + using WarpCount = typename Base::WarpCount; + + /// Shared memory allocation from epilogue base class + using BaseSharedStorage = typename Base::SharedStorage; + + static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1 ? Base::kFragmentsPerIteration : kPartitionsK; + static int constexpr kSmemPointerOffset = Base::SharedStorage::StorageShape::kCount / kSmemTiles; + + /// Used for the broadcast + struct BroadcastDetail { + + /// Number of threads per warp + static int const kWarpSize = 32; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + + /// Number of distinct scalar column indices handled by each thread + static int const kColumnsPerThread = ThreadMap::Iterations::kColumn * ThreadMap::kElementsPerAccess; + + /// Number of distinct scalar row indices handled by each thread + static int const kRowsPerThread = ThreadMap::Iterations::kCount / ThreadMap::Iterations::kColumn; + + /// Number of threads per threadblock + static int const kThreadCount = kWarpSize * WarpCount::kCount; + + /// Number of distinct threads per row of output tile + static int const kThreadsPerRow = (Shape::kN / kColumnsPerThread); + + /// Number of distinct threads which must be reduced during the final reduction phase within the threadblock. + static int const kThreadRows = kThreadCount / kThreadsPerRow; + + /// I'm not sure what I meant here. + static int const kThreadAccessesPerRow = const_max(1, (Shape::kN + kThreadCount - 1) / kThreadCount); + + /// Shape of the shared memory allocation for the epilogue + using StorageShape = MatrixShape< + kThreadRows, + Shape::kN + >; + + /// Debug printing + MCTLASS_DEVICE + static void print() { +#if 0 + printf("BroadcastDetail {\n"); + printf( + " kColumnsPerThread: %d\nkRowsPerThread: %d\n,kThreadCount: %d\nkThreadsPerRow: %d\n" + "kThreadRows: %d\nThreadAccessesPerRow: %d\nStorageShape: %d x %d (count: %d)\n", + kColumnsPerThread, + kRowsPerThread, + kThreadCount, + kThreadsPerRow, + kThreadRows, + kThreadAccessesPerRow, + StorageShape::kRow, + StorageShape::kColumn, + StorageShape::kCount + ); + printf("};\n"); +#endif + } + }; + + /// Shared storage structure (shadows base) with additional SMEM buffer for reduction + struct SharedStorage { + union { + BaseSharedStorage base; + }; + + MCTLASS_HOST_DEVICE + SharedStorage() { } + }; + +public: + + + static_assert(SharedLoadIterator::Fragment::kElements == OutputTileIterator::Fragment::kElements, + "Mismatch between shared load iterator and output tile iterator."); + + static_assert(OutputTileIterator::kElementsPerAccess, "OutputTileIterator::kElementsPerAccess must not be zero."); + + static_assert(!(OutputTileIterator::Fragment::kElements % OutputTileIterator::kElementsPerAccess), + "Divisibility"); + +private: + + /// Loads fragment from shared memory aligned with output tensor + SharedLoadIterator shared_load_iterator_; + + /// Thread index within the threadblock + int thread_idx_; + +public: + + /// Constructor + MCTLASS_DEVICE + EpilogueWithBroadcast( + SharedStorage &shared_storage, ///< Shared storage object + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx ///< Id of thread within warp + ): + Base(shared_storage.base, thread_idx, warp_idx, lane_idx), + shared_load_iterator_(shared_storage.base.reference(), thread_idx), + thread_idx_(thread_idx) + { + + } + + /// Streams the result to global memory + MCTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + ElementVector const * broadcast_ptr, ///< Broadcast vector + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator, ///< Tile iterator for source accumulator matrix + TensorTileIterator tensor_iterator, ///< Threadblock tile iterator for additional tensor operand + MatrixCoord const &problem_size = ///< Problem size needed to guard against out-of-bounds accesses + MatrixCoord(Shape::kM, Shape::kN), + MatrixCoord const &threadblock_offset = ///< Threadblock's initial offset within the problem size space + MatrixCoord()) { + + BroadcastFragment broadcast_fragment; + + load_broadcast_fragment_(broadcast_fragment, broadcast_ptr, problem_size, threadblock_offset); + + if (!output_op.is_source_needed()) { + compute_source_not_needed_( + output_op, + broadcast_fragment, + destination_iterator, + accumulators, + tensor_iterator); + } + else { + compute_source_needed_( + output_op, + broadcast_fragment, + destination_iterator, + accumulators, + source_iterator, + tensor_iterator); + } + } + +private: + + MCTLASS_DEVICE + void load_broadcast_fragment_( + BroadcastFragment & broadcast_fragment, ///< Fragment containing the accumulated partial reduction over columns + ElementVector const * broadcast_ptr, ///< Broadcast vector + MatrixCoord const &problem_size, ///< Problem size needed to guard against out-of-bounds accesses + MatrixCoord const &threadblock_offset ///< Threadblock's initial offset within the problem size space + ) { + + broadcast_fragment.clear(); + + // If no pointer is supplied, set with all zeros and avoid memory accesses + if (!broadcast_ptr) { + return; + } + + int thread_initial_column = ThreadMap::initial_offset(thread_idx_).column(); + + int thread_column_idx = threadblock_offset.column() + thread_initial_column; + broadcast_ptr += thread_initial_column; + + NumericArrayConverter converter; + using AccessType = AlignedArray; + using ComputeFragmentType = Array; + + ComputeFragmentType *frag_ptr = reinterpret_cast(&broadcast_fragment); + + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < ThreadMap::Iterations::kColumn; ++j) { + + AccessType loaded; + + loaded.clear(); + + if (thread_column_idx < problem_size.column()) { + loaded = *reinterpret_cast(broadcast_ptr); + } + + ComputeFragmentType cvt = converter(loaded); + frag_ptr[j] = cvt; + + thread_column_idx += ThreadMap::Delta::kColumn; + broadcast_ptr += ThreadMap::Delta::kColumn; + } + } + + template + struct acc2smem_source_not_needed; + + template + struct acc2smem_source_not_needed> { + template + MCTLASS_DEVICE static void helper(AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator &warp_tile_iterator) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { + ++accum_fragment_iterator; + } + + MCTLASS_PRAGMA_UNROLL + for (int p = 0; p < Base::kFragmentsPerIteration; ++p) { + typename AccumulatorFragmentIterator::Fragment accum_fragment; + + accum_fragment_iterator.load(accum_fragment); + ++accum_fragment_iterator; + + warp_tile_iterator.store(accum_fragment); + if (p < Base::kFragmentsPerIteration - 1) { + warp_tile_iterator.add_pointer_offset(kSmemPointerOffset); + } + } + + if (Base::kFragmentsPerIteration > 1) { + warp_tile_iterator.add_pointer_offset(kSmemPointerOffset * + (1 - Base::kFragmentsPerIteration)); + } + } + + MCTLASS_DEVICE + static void push(size_t pos, + AccumulatorFragmentIterator const &iterator_begin, + WarpTileIterator &warp_tile_iterator) { + int dummy[] = { + (pos == (Seq * Base::kFragmentsPerIteration)) && + (helper(iterator_begin, warp_tile_iterator), 0)...}; + + MCTLASS_UNUSED(dummy[0]); + } + }; + + /// Streams the result to global memory + MCTLASS_DEVICE + void compute_source_not_needed_( + OutputOp const &output_op, ///< Output operator + BroadcastFragment const &broadcast_fragment, ///< Fragment containing the accumulated partial reduction over columns + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + TensorTileIterator tensor_iterator ///< Threadblock tile iterator for additioanl tensor operand + ) { + + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + + // MCTLASS_PRAGMA_UNROLL + #pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations / Base::kFragmentsPerIteration : 1) + for (int iter = 0; iter < OutputTileIterator::kIterations; iter += Base::kFragmentsPerIteration) { + + // + // Convert and store fragment + // + + + __syncthreads(); + + acc2smem_source_not_needed< + mctlass::make_index_sequence>::push(iter, + accum_fragment_iterator, + this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + MCTLASS_PRAGMA_UNROLL + for (int p = 0; p < Base::kFragmentsPerIteration; ++p) { + + + typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment[0]); + + if (p < Base::kFragmentsPerIteration - 1) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + } + else if (kPartitionsK > 1) { + + plus add_fragments; + + MCTLASS_PRAGMA_UNROLL + for ( int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + + shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * kSmemPointerOffset); + } + + // + // Apply output operation + // + + typename OutputTileIterator::Fragment frag_Z; + typename TensorTileIterator::Fragment frag_T; + + apply_output_operator_source_not_needed_( + frag_Z, + frag_T, + output_op, + aligned_accum_fragment[0], + broadcast_fragment); + + // + // Conditionally store fragments + // + + if (OutputOp::kStoreZ) { + destination_iterator.store(frag_Z); + ++destination_iterator; + } + + if (OutputOp::kStoreT) { + tensor_iterator.store(frag_T); + ++tensor_iterator; + } + } + + if (Base::kFragmentsPerIteration > 1) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset * (1 - Base::kFragmentsPerIteration)); + } + } + } + + + template + struct acc2smem_source_needed; + + template + struct acc2smem_source_needed> { + template + MCTLASS_DEVICE + static void helper(AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator &warp_tile_iterator) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { + ++accum_fragment_iterator; + } + + typename AccumulatorFragmentIterator::Fragment accum_fragment; + accum_fragment_iterator.load(accum_fragment); + warp_tile_iterator.store(accum_fragment); + } + + MCTLASS_DEVICE + static void push(size_t pos, + AccumulatorFragmentIterator const &iterator_begin, + WarpTileIterator &warp_tile_iterator) { + int dummy[] = {(pos == Seq) && (helper(iterator_begin, warp_tile_iterator), 0)...}; + } + }; + + + /// Streams the result to global memory + MCTLASS_DEVICE + void compute_source_needed_( + OutputOp const &output_op, ///< Output operator + BroadcastFragment const &broadcast_fragment, ///< Fragment containing the accumulated partial reduction over columns + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator, ///< Tile iterator for source accumulator matrix + TensorTileIterator tensor_iterator ///< Threadblock tile iterator for additioanl tensor operand + ) { + + typename OutputTileIterator::Fragment source_fragment; + source_fragment.clear(); + + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + + //#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1) + //mxcc not support using unroll with parentheses, and we would support with unroll with OutputTileIterator::kIterations someday + #pragma unroll 1 + for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) { + + // + // Load the source + // + + source_iterator.load(source_fragment); + ++source_iterator; + + // + // Convert and store fragment + // + + __syncthreads(); + + acc2smem_source_needed>::push( + iter, accum_fragment_iterator, this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment[0]); + + // If the number of k-slices is > 1 - perform a reduction amongst the k-slices + if (kPartitionsK > 1) + { + plus add_fragments; + const int tile_row_offset = Base::SharedStorage::StorageShape::kRow / PartitionsK; + + MCTLASS_PRAGMA_UNROLL + for ( int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_tile_offset({tile_row_offset , 0}); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + + shared_load_iterator_.add_tile_offset({-1 * (kPartitionsK-1) * tile_row_offset, 0}); + } + + // + // Apply output operation + // + + typename OutputTileIterator::Fragment frag_Z; + typename TensorTileIterator::Fragment frag_T; + + apply_output_operator_( + frag_Z, + frag_T, + output_op, + aligned_accum_fragment[0], + source_fragment, + broadcast_fragment); + + // + // Conditionally store fragments + // + + if (OutputOp::kStoreZ) { + destination_iterator.store(frag_Z); + ++destination_iterator; + } + + if (OutputOp::kStoreT) { + tensor_iterator.store(frag_T); + ++tensor_iterator; + } + } + } + + /// Helper to invoke the output functor over each vector of output + MCTLASS_DEVICE + void apply_output_operator_( + typename OutputTileIterator::Fragment &frag_Z, + typename TensorTileIterator::Fragment &frag_T, + OutputOp const &output_op, + typename SharedLoadIterator::Fragment const &frag_AB, + typename OutputTileIterator::Fragment const &frag_C, + BroadcastFragment const &frag_Broadcast) { + + using AccessTypeZ = Array; + using AccessTypeT = Array; + using AccessTypeBroadcast = Array; + + AccessTypeZ *frag_Z_ptr = reinterpret_cast(&frag_Z); + AccessTypeT *frag_T_ptr = reinterpret_cast(&frag_T); + + AccumulatorAccessType const *frag_AB_ptr = + reinterpret_cast(&frag_AB); + + OutputAccessType const *frag_C_ptr = + reinterpret_cast(&frag_C); + + AccessTypeBroadcast const *frag_Broadcast_ptr = + reinterpret_cast(&frag_Broadcast); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + output_op( + frag_Z_ptr[i], + frag_T_ptr[i], + frag_AB_ptr[i], + frag_C_ptr[i], + frag_Broadcast_ptr[i % ThreadMap::Iterations::kColumn]); + } + } + + /// Helper to invoke the output functor over each vector of output + MCTLASS_DEVICE + void apply_output_operator_source_not_needed_( + typename OutputTileIterator::Fragment &frag_Z, + typename TensorTileIterator::Fragment &frag_T, + OutputOp const &output_op, + typename SharedLoadIterator::Fragment const &frag_AB, + BroadcastFragment const &frag_Broadcast) { + + using AccessTypeZ = Array; + using AccessTypeT = Array; + using AccessTypeBroadcast = Array; + + AccessTypeZ *frag_Z_ptr = reinterpret_cast(&frag_Z); + AccessTypeT *frag_T_ptr = reinterpret_cast(&frag_T); + + AccumulatorAccessType const *frag_AB_ptr = + reinterpret_cast(&frag_AB); + + AccessTypeBroadcast const *frag_Broadcast_ptr = + reinterpret_cast(&frag_Broadcast); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + + output_op( + frag_Z_ptr[i], + frag_T_ptr[i], + frag_AB_ptr[i], + frag_Broadcast_ptr[i % ThreadMap::Iterations::kColumn]); + } + } + + + public: + /// Stream-K reduce helper + MCTLASS_DEVICE + void reduce( + int reduce_fragment_idx, ///< Reduce fragment index + OutputOp const &output_op, ///< Output operator + ElementVector const * broadcast_ptr, ///< Broadcast vector + OutputTileIterator destination_iterator, ///< Tile iterator for destination + OutputTileIterator source_iterator, ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + TensorTileIterator tensor_iterator, ///< Threadblock tile iterator for additional tensor operand + MatrixCoord const &problem_size = ///< Problem size needed to guard against out-of-bounds accesses + MatrixCoord(Shape::kM, Shape::kN), + MatrixCoord const &threadblock_offset = ///< Threadblock's initial offset within the problem size space + MatrixCoord()) + { + + BroadcastFragment broadcast_fragment; + load_broadcast_fragment_(broadcast_fragment, broadcast_ptr, problem_size, threadblock_offset); + + // Initialize/load source-fragment data + typename OutputTileIterator::Fragment source_fragment; + source_fragment.clear(); + + if (output_op.is_source_needed()) + { + source_iterator += reduce_fragment_idx; + source_iterator.load(source_fragment); + } + + // Load fragment from shared memory + typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK]; + shared_load_iterator_.load(aligned_accum_fragment[0]); + + // Add fragments shared by other k partitions + if (kPartitionsK > 1) + { + plus add_fragments; + + MCTLASS_PRAGMA_UNROLL + for ( int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + } + + // + // Apply output operation + // + + typename OutputTileIterator::Fragment frag_Z; + typename TensorTileIterator::Fragment frag_T; + + if (!output_op.is_source_needed()) { + apply_output_operator_source_not_needed_( + frag_Z, + frag_T, + output_op, + aligned_accum_fragment[0], + broadcast_fragment); + } else { + apply_output_operator_( + frag_Z, + frag_T, + output_op, + aligned_accum_fragment[0], + source_fragment, + broadcast_fragment); + } + + // + // Conditionally store fragments + // + + if (OutputOp::kStoreZ) { + destination_iterator.store(frag_Z); + ++destination_iterator; + } + + if (OutputOp::kStoreT) { + tensor_iterator.store(frag_T); + ++tensor_iterator; + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_with_reduction.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_with_reduction.h new file mode 100644 index 0000000..3d2875b --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_with_reduction.h @@ -0,0 +1,827 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/tensor_coord.h" +#include "mctlass/aligned_buffer.h" +#include "mctlass/functional.h" +#include "mctlass/fast_math.h" +#include "mctlass/layout/vector.h" +#include "mctlass/layout/tensor.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_iterator.h" + +#include "mctlass/epilogue/threadblock/epilogue_base.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator with reduction over each column +template < + typename Shape_, ///< Shape of threadblock tile (concept: GemmShape) + typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp) + int PartitionsK, ///< Number of partitions of the K dimension + typename OutputTileIterator_, ///< Tile iterator reading and writing output tensors + typename TensorTileIterator_, ///< Additional tile iterator for tensor-valued operands + typename ElementVector_, ///< Pointer to reduction vector + typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators + typename WarpTileIterator_, ///< Warp-scoped tile iterator writing accumulators to SMEM + typename SharedLoadIterator_, ///< Threadblock-scoped tile iterator loading from SMEM + typename OutputOp_, ///< Output operator + typename ReductionOp_, ///< Reduction operator + typename Padding_, ///< Padding added to SMEM allocation to avoid bank conflicts (concept: MatrixShape) + int IterationsUnroll = ///< Used to reduce binary size when epilogue op is large + (!IsEpilogueFunctorHeavy::value) +> +class EpilogueWithReduction : + public EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_> { + +public: + + using Base = EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_>; + + using Shape = Shape_; + using WarpMmaOperator = WarpMmaOperator_; + static int const kPartitionsK = PartitionsK; + using OutputTileIterator = OutputTileIterator_; + using TensorTileIterator = TensorTileIterator_; + using ElementVector = ElementVector_; + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using WarpTileIterator = WarpTileIterator_; + using SharedLoadIterator = SharedLoadIterator_; + using OutputOp = OutputOp_; + using ReductionOp = ReductionOp_; + using Padding = Padding_; + + using Layout = layout::RowMajor; + using LongIndex = typename Layout::LongIndex; + + static bool const kIsSingleSource = true; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename Base::AccumulatorTile; + + /// Accumulator element + using ElementAccumulator = typename WarpTileIterator::Element; + + /// Compute data type produced by the output op + using ElementCompute = typename OutputOp::ElementCompute; + + /// Compute fragment + using FragmentCompute = Array; + + /// Thread map used by output tile iterators + using ThreadMap = typename OutputTileIterator::ThreadMap; + + /// Fragment object used in reduction + using ReductionFragment = Array< + ElementAccumulator, + ThreadMap::Iterations::kColumn * ThreadMap::kElementsPerAccess>; + + /// Output element + using ElementOutput = typename OutputTileIterator::Element; + + /// Data type of additional tensor + using ElementTensor = typename TensorTileIterator::Element; + + /// Output access size + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + /// Tensor reference to destination tensor + using TensorRef = typename OutputTileIterator::TensorRef; + + /// Tensor reference to sync tensor + using SyncTensorRef = typename mctlass::TensorRef; + + /// Const tensor reference to source tensor + using ConstTensorRef = typename OutputTileIterator::ConstTensorRef; + + /// Array type used to output + using OutputAccessType = Array< + typename OutputTileIterator::Element, OutputTileIterator::kElementsPerAccess>; + + /// Array type used by output functor + using AccumulatorAccessType = Array; + + /// Array type used by output functor + using ComputeAccessType = Array; + + /// Tensor access type + using TensorAccessType = Array; + + /// Number of warps + using WarpCount = typename Base::WarpCount; + + /// Shared memory allocation from epilogue base class + using BaseSharedStorage = typename Base::SharedStorage; + + /// Used for the reduction + struct ReductionDetail { + + /// If true, accumulator coordinates are computed and out-of-bounds checks are enabled when + /// performing the reduction. + static bool const kOobCheck = false; + + /// Number of threads per warp + static int const kWarpSize = 32; + + /// Number of distinct scalar column indices handled by each thread + static int const kColumnsPerThread = ThreadMap::Iterations::kColumn * ThreadMap::kElementsPerAccess; + + /// Number of distinct scalar row indices handled by each thread + static int const kRowsPerThread = ThreadMap::Iterations::kCount / ThreadMap::Iterations::kColumn; + + /// Number of threads per threadblock + static int const kThreadCount = kWarpSize * WarpCount::kCount; + + /// Number of distinct threads per row of output tile + static int const kThreadsPerRow = (Shape::kN / kColumnsPerThread); + + /// Number of distinct threads which must be reduced during the final reduction phase within the threadblock. + static int const kThreadRows = kThreadCount / kThreadsPerRow; + + /// I'm not sure what I meant here. + static int const kThreadAccessesPerRow = const_max(1, (Shape::kN + kThreadCount - 1) / kThreadCount); + + /// Shape of the shared memory allocation for the epilogue + using StorageShape = MatrixShape< + kThreadRows, + Shape::kN + >; + + /// Debug printing + MCTLASS_DEVICE + static void print() { +#if 0 + printf("ReductionDetail {\n"); + printf( + " kElementsPerAccess:%d\nkColumnsPerThread: %d\nkRowsPerThread: %d\n,kThreadCount: %d\nkThreadsPerRow: %d\n" + "kThreadRows: %d\nThreadAccessesPerRow: %d\nStorageShape: %d x %d (count: %d)\n", + kElementsPerAccess, + kColumnsPerThread, + kRowsPerThread, + kThreadCount, + kThreadsPerRow, + kThreadRows, + kThreadAccessesPerRow, + StorageShape::kRow, + StorageShape::kColumn, + StorageShape::kCount + ); + printf("};\n"); +#endif + } + }; + + /// Shared storage structure (shadows base) with additional SMEM buffer for reduction + struct SharedStorage { + union { + BaseSharedStorage base; + AlignedArray reduction; ///< Shared storage for reduction + }; + + MCTLASS_HOST_DEVICE + SharedStorage() { } + }; + +public: + + + static_assert(SharedLoadIterator::Fragment::kElements == OutputTileIterator::Fragment::kElements, + "Mismatch between shared load iterator and output tile iterator."); + + static_assert(OutputTileIterator::kElementsPerAccess, "OutputTileIterator::kElementsPerAccess must not be zero."); + + static_assert(!(OutputTileIterator::Fragment::kElements % OutputTileIterator::kElementsPerAccess), + "Divisibility"); + +private: + + /// Loads fragment from shared memory aligned with output tensor + SharedLoadIterator shared_load_iterator_; + + /// Shared memory pointer fo rreduction + ElementAccumulator *reduction_ptr_; + + /// Thread index within the threadblock + int thread_idx_; + +public: + + /// Constructor + MCTLASS_DEVICE + EpilogueWithReduction( + SharedStorage &shared_storage, ///< Shared storage object + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx ///< Id of thread within warp + ): + Base(shared_storage.base, thread_idx, warp_idx, lane_idx), + shared_load_iterator_(shared_storage.base.reference(), thread_idx), + reduction_ptr_(shared_storage.reduction.data()), + thread_idx_(thread_idx) + { + + } + + /// Streams the result to global memory + MCTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + ElementVector * reduction_output_ptr, ///< Reduction output vector + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator, ///< Tile iterator for source accumulator matrix + TensorTileIterator tensor_iterator, ///< Threadblock tile iterator for additional tensor operand + MatrixCoord const &problem_size = ///< Problem size needed to guard against out-of-bounds accesses + MatrixCoord(Shape::kM, Shape::kN), + MatrixCoord const &threadblock_offset = ///< Threadblock's initial offset within the problem size space + MatrixCoord()) { + + ReductionFragment reduction_fragment; + reduction_fragment.clear(); + + if (!output_op.is_source_needed()) { + compute_source_not_needed_( + output_op, + reduction_fragment, + destination_iterator, + accumulators, + tensor_iterator, + problem_size, + threadblock_offset); + } + else { + compute_source_needed_( + output_op, + reduction_fragment, + destination_iterator, + accumulators, + source_iterator, + tensor_iterator, + problem_size, + threadblock_offset); + } + + if (output_op.participates_in_reduction()) { + reduction_(problem_size, threadblock_offset, reduction_output_ptr, reduction_fragment); + } + } + +private: + + /// Perform the reduction + MCTLASS_DEVICE + void reduction_( + MatrixCoord const &problem_size, ///< Problem size needed to guard against out-of-bounds accesses + MatrixCoord const &threadblock_offset, ///< Problem size needed to guard against out-of-bounds accesses + ElementVector * reduction_output_ptr, ///< Reduction output vector + ReductionFragment const & reduction_fragment) { + + // + // Store the partially reduced value to SMEM + // + + // Guard against uses of the existing SMEM tile + __syncthreads(); + + using AccessType = AlignedArray; + + // + // Determine a compacted thread arrangement to store to SMEM. + // + int const kThreadsPerRow = Shape::kN / (ThreadMap::Iterations::kColumn * ThreadMap::kElementsPerAccess); + + MatrixCoord thread_offset( + thread_idx_ / kThreadsPerRow, + (thread_idx_ % kThreadsPerRow) * ThreadMap::kElementsPerAccess); + + // + // Each thread store its fragment to a SMEM + // + + AccessType *aligned_reduction_ptr = reinterpret_cast( + &reduction_ptr_[thread_offset.row() * Shape::kN + thread_offset.column()]); + + AccessType const *frag_ptr = reinterpret_cast(&reduction_fragment); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + int col_idx = column * ThreadMap::Delta::kColumn / ThreadMap::kElementsPerAccess; + + aligned_reduction_ptr[col_idx] = frag_ptr[column]; + } + + __syncthreads(); + + // + // Now, threads are assigned several columns of the output. They fetch over all rows from + // the compacted SMEM tile and perform a reduction. + // + + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < ReductionDetail::kThreadAccessesPerRow; ++j) { + int column_idx = thread_idx_ + j * ReductionDetail::kThreadCount; + + ReductionOp reduction_op; + ElementAccumulator reduction_element = ElementAccumulator(); + + int output_column_idx = threadblock_offset.column() + column_idx; + + if (column_idx < Shape::kN && output_column_idx < problem_size.column()) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ReductionDetail::kThreadRows; ++row) { + if (row) { + auto frag = reduction_ptr_[row * Shape::kN + column_idx]; + + reduction_element = reduction_op(reduction_element, frag); + } + else { + + reduction_element = reduction_ptr_[column_idx]; + } + } + + // Store + reduction_output_ptr[column_idx] = ElementVector(reduction_element); + } + } + } + + template + struct acc2smem; + + template + struct acc2smem> { + template + MCTLASS_DEVICE + static void helper(AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator &warp_tile_iterator) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { + ++accum_fragment_iterator; + } + + typename AccumulatorFragmentIterator::Fragment accum_fragment; + accum_fragment_iterator.load(accum_fragment); + warp_tile_iterator.store(accum_fragment); + } + + MCTLASS_DEVICE + static void push(size_t pos, + AccumulatorFragmentIterator const &iterator_begin, + WarpTileIterator &warp_tile_iterator) { + int dummy[] = {(pos == Seq) && (helper(iterator_begin, warp_tile_iterator), 0)...}; + } + }; + + /// Streams the result to global memory + MCTLASS_DEVICE + void compute_source_not_needed_( + OutputOp const &output_op, ///< Output operator + ReductionFragment &reduction_fragment, ///< Fragment containing the accumulated partial reduction over columns + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + TensorTileIterator tensor_iterator, ///< Threadblock tile iterator for additioanl tensor operand + MatrixCoord const &problem_size, ///< Problem size needed to guard against out-of-bounds accesses + MatrixCoord const &threadblock_offset ///< Threadblock's initial offset within the problem size space + ) { + + // + // Iterator over warp-level accumulator fragment + // + + typename TensorTileIterator::Fragment tensor_fragment; + tensor_fragment.clear(); + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + + //#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1) + //mxcc not support using unroll with parentheses, and we would support with unroll with OutputTileIterator::kIterations someday + #pragma unroll 1 + for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) { + + // + // Convert and store fragment + // + + tensor_iterator.load(tensor_fragment); + ++tensor_iterator; + + __syncthreads(); + + acc2smem>::push( + iter, accum_fragment_iterator, this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment[0]); + + // + // If the number of k-slices is > 1 - perform a reduction amongst the k-slices + // + if (kPartitionsK > 1) + { + plus add_fragments; + const int tile_row_offset = Base::SharedStorage::StorageShape::kRow / PartitionsK; + + MCTLASS_PRAGMA_UNROLL + for ( int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_tile_offset({tile_row_offset , 0}); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + + shared_load_iterator_.add_tile_offset({-1 * (kPartitionsK-1) * tile_row_offset, 0}); + } + + // + // Compute the output result + // + + FragmentCompute compute_fragment; + + apply_output_operator_source_not_needed_( + reduction_fragment, + compute_fragment, + output_op, + aligned_accum_fragment[0], + tensor_fragment, + destination_iterator); + + // + // Store the final result + // + + NumericArrayConverter converter; + + typename OutputTileIterator::Fragment output_fragment = converter(compute_fragment); + + destination_iterator.store(output_fragment); + ++destination_iterator; + } + } + + + /// Streams the result to global memory + MCTLASS_DEVICE + void compute_source_needed_( + OutputOp const &output_op, ///< Output operator + ReductionFragment &reduction_fragment, ///< Fragment containing the accumulated partial reduction over columns + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator, ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + TensorTileIterator tensor_iterator, ///< Threadblock tile iterator for additioanl tensor operand + MatrixCoord const &problem_size, ///< Problem size needed to guard against out-of-bounds accesses + MatrixCoord const &threadblock_offset ///< Threadblock's initial offset within the problem size space + ) { + + typename OutputTileIterator::Fragment source_fragment; + source_fragment.clear(); + + typename TensorTileIterator::Fragment tensor_fragment; + tensor_fragment.clear(); + + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + + //#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1) + //mxcc not support using unroll with parentheses, and we would support with unroll with OutputTileIterator::kIterations someday + #pragma unroll 1 + for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) { + + // + // Load the source + // + + source_fragment.clear(); + source_iterator.load(source_fragment); + ++source_iterator; + + tensor_iterator.load(tensor_fragment); + ++tensor_iterator; + + // + // Convert and store fragment + // + + __syncthreads(); + + acc2smem>::push( + iter, accum_fragment_iterator, this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment[0]); + + // If the number of k-slices is > 1 - perform a reduction amongst the k-slices + if (kPartitionsK > 1) + { + plus add_fragments; + const int tile_row_offset = Base::SharedStorage::StorageShape::kRow / PartitionsK; + + MCTLASS_PRAGMA_UNROLL + for ( int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_tile_offset({tile_row_offset , 0}); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + + shared_load_iterator_.add_tile_offset({-1 * (kPartitionsK-1) * tile_row_offset, 0}); + } + + // + // Compute the output result + // + + FragmentCompute compute_fragment; + + apply_output_operator_( + reduction_fragment, + compute_fragment, + output_op, + aligned_accum_fragment[0], + source_fragment, + tensor_fragment, + destination_iterator); + + // + // Convert and store the final result + // + + NumericArrayConverter converter; + + typename OutputTileIterator::Fragment output_fragment = converter(compute_fragment); + + destination_iterator.store(output_fragment); + ++destination_iterator; + } + } + + /// Helper to invoke the output functor over each vector of output + MCTLASS_DEVICE + void apply_output_operator_( + ReductionFragment &reduction_fragment, + FragmentCompute &compute_fragment, + OutputOp const &output_op, ///< Output operator + typename SharedLoadIterator::Fragment const &aligned_accum_fragment, + typename OutputTileIterator::Fragment const &source_fragment, + typename TensorTileIterator::Fragment const &tensor_fragment, + OutputTileIterator const & destination_iterator) { + + ComputeAccessType *compute_frag_ptr = + reinterpret_cast(&compute_fragment); + + AccumulatorAccessType const *accum_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + OutputAccessType const *source_frag_ptr = + reinterpret_cast(&source_fragment); + + TensorAccessType const *tensor_frag_ptr = + reinterpret_cast(&tensor_fragment); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + + // Call the output operator + compute_frag_ptr[i] = output_op(accum_frag_ptr[i], source_frag_ptr[i], tensor_frag_ptr[i]); + } + + // + // Partial reduction over each column + // + + ReductionOp reduction_op; + + typename OutputTileIterator::Mask mask; + destination_iterator.get_mask(mask); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ReductionDetail::kColumnsPerThread; ++column) { + + int column_vector_idx = column / ThreadMap::kElementsPerAccess; + bool column_guard = mask.predicates[column_vector_idx]; + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ReductionDetail::kRowsPerThread; ++row) { + + bool fetch; + if (ReductionDetail::kOobCheck) { + int row_idx = (row % ThreadMap::Iterations::kRow); + int residual = (row / ThreadMap::Iterations::kRow); + + int group_idx = (residual % ThreadMap::Iterations::kGroup); + residual = (residual / ThreadMap::Iterations::kGroup); + + int cluster_idx = (residual % ThreadMap::Iterations::kCluster); + + int row_offset = row_idx * ThreadMap::Delta::kRow + + group_idx * ThreadMap::Delta::kGroup + + cluster_idx * ThreadMap::Delta::kCluster; + + int output_row = destination_iterator.thread_start_row() + row_offset; + + fetch = (output_row < destination_iterator.extent_row() && column_guard); + } + else { + fetch = true; + } + + ElementCompute value = ElementCompute(); + if (fetch) { + value = compute_fragment[row * ReductionDetail::kColumnsPerThread + column]; + } + + reduction_fragment[column] = reduction_op( + reduction_fragment[column], + value); + } + } + } + + /// Helper to invoke the output functor over each vector of output + MCTLASS_DEVICE + void apply_output_operator_source_not_needed_( + ReductionFragment &reduction_fragment, + FragmentCompute &compute_fragment, + OutputOp const &output_op, ///< Output operator + typename SharedLoadIterator::Fragment const &aligned_accum_fragment, + typename TensorTileIterator::Fragment const &tensor_fragment, + OutputTileIterator const & destination_iterator + ) { + + ComputeAccessType *compute_frag_ptr = + reinterpret_cast(&compute_fragment); + + AccumulatorAccessType const *accum_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + TensorAccessType const *tensor_frag_ptr = + reinterpret_cast(&tensor_fragment); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) { + + // Call the output operator + compute_frag_ptr[i] = output_op(accum_frag_ptr[i], tensor_frag_ptr[i]); + } + + // + // Partial reduction over each column + // + + ReductionOp reduction_op; + + typename OutputTileIterator::Mask mask; + destination_iterator.get_mask(mask); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ReductionDetail::kColumnsPerThread; ++column) { + + int column_vector_idx = column / ThreadMap::kElementsPerAccess; + bool column_guard = mask.predicates[column_vector_idx]; + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ReductionDetail::kRowsPerThread; ++row) { + + bool fetch; + if (ReductionDetail::kOobCheck) { + int row_idx = (row % ThreadMap::Iterations::kRow); + int residual = (row / ThreadMap::Iterations::kRow); + + int group_idx = (residual % ThreadMap::Iterations::kGroup); + residual = (residual / ThreadMap::Iterations::kGroup); + + int cluster_idx = (residual % ThreadMap::Iterations::kCluster); + + int row_offset = row_idx * ThreadMap::Delta::kRow + + group_idx * ThreadMap::Delta::kGroup + + cluster_idx * ThreadMap::Delta::kCluster; + + int output_row = destination_iterator.thread_start_row() + row_offset; + + fetch = (output_row < destination_iterator.extent_row() && column_guard); + } + else { + fetch = true; + } + + ElementCompute value = ElementCompute(); + if (fetch) { + value = compute_fragment[row * ReductionDetail::kColumnsPerThread + column]; + } + + reduction_fragment[column] = reduction_op( + reduction_fragment[column], + value); + } + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_with_visitor.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_with_visitor.h new file mode 100644 index 0000000..cc44803 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_with_visitor.h @@ -0,0 +1,411 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Generic epilogue for implementing certain kinds of fused epilogue behavior. +*/ + +#pragma once + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/semaphore.h" +#include "mctlass/epilogue/threadblock/epilogue_base.h" + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +class EpilogueFusedVisitorConcept { +public: + + static int const kIterations = 1; + static int const kElementsPerAccess = 4; + using ElementOutput = float; + using ElementAccumulator = float; + using AccumulatorFragment = Array; + + /// Arguments structure + struct Arguments { }; + + /// Params structure + struct Params { + + Params() { } + Params(Arguments const &args) { } + }; + + /// Shared storage + struct SharedStorage { }; + +public: + + MCTLASS_DEVICE + EpilogueFusedVisitorConcept( + Params const ¶ms, ///< Parameters routed to the epilogue + SharedStorage &shared_storage, ///< Shared storage needed by the functors here + MatrixCoord const &problem_size, ///< Problem size of the output + int thread_idx, ///< Thread index within the threadblock + int warp_idx, ///< Warp index within the threadblock + int lane_idx, ///< Lane index within the warp + MatrixCoord const &threadblock_offset = MatrixCoord(0, 0)) { ///< Coordinate + + } + + /// Helper to indicate split-K behavior + MCTLASS_DEVICE + void set_k_partition( + int split_k_index, ///< Index of this threadblock within split-K partitioned scheme + int split_k_slices) { ///< Total number of split-K slices + + } + + /// Called to set the batch index + MCTLASS_DEVICE + void set_batch_index(int batch_idx) { + + } + + /// Called at the start of the epilogue just before iterating over accumulator slices + MCTLASS_DEVICE + void begin_epilogue() { + + } + + /// Called at the start of one step before starting accumulator exchange + MCTLASS_DEVICE + void begin_step(int step_idx) { + + } + + /// Called at the start of a row + MCTLASS_DEVICE + void begin_row(int row_idx) { + + } + + /// Called after accumulators have been exchanged for each accumulator vector + MCTLASS_DEVICE + void visit( + int iter_idx, + int row_idx, + int column_idx, + int frag_idx, + AccumulatorFragment const &accum) { + + } + + /// Called at the end of a row + MCTLASS_DEVICE + void end_row(int row_idx) { + + } + + /// Called after all accumulator elements have been visited + MCTLASS_DEVICE + void end_step(int step_idx) { + + } + + /// Called after all steps have been completed + MCTLASS_DEVICE + void end_epilogue() { + + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator +template < + typename Visitor_, ///< Functor containing fused operations (satisfies EpilogueFusedVisitorConcept) + typename Shape_, ///< Shape of threadblock tile (concept: GemmShape) + typename WarpMmaOperator_, ///< Warp-level MMA operator (concept: gemm::warp::MmaTensorOp) + int PartitionsK, ///< Number of partitions of the K dimension + typename AccumulatorFragmentIterator_, ///< Fragment iterator selecting accumulators + typename WarpTileIterator_, ///< Warp-scoped tile iterator writing accumulators to SMEM + typename SharedLoadIterator_, ///< Threadblock-scoped tile iterator loading from SMEM + typename Padding_, ///< Padding added to SMEM allocation to avoid bank conflicts (concept: MatrixShape) + int FragmentsPerPartition = 1, ///< Used to coarsten the epilogue granularity + int IterationsUnroll = ///< Used to reduce binary size when epilogue op is large + (true || !IsEpilogueFunctorHeavy::value) +> +class EpilogueWithVisitor : + public EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_, + FragmentsPerPartition> { + +public: + + using Visitor = Visitor_; + + using Base = EpilogueBase< + Shape_, + typename WarpMmaOperator_::Shape, + PartitionsK, + AccumulatorFragmentIterator_, + WarpTileIterator_, + Padding_, + FragmentsPerPartition>; + + using Shape = Shape_; + using WarpMmaOperator = WarpMmaOperator_; + static int const kPartitionsK = PartitionsK; + + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using WarpTileIterator = WarpTileIterator_; + using SharedLoadIterator = SharedLoadIterator_; + using Padding = Padding_; + + using Layout = layout::RowMajor; + using LongIndex = typename Layout::LongIndex; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename Base::AccumulatorTile; + + /// Accumulator element + using ElementAccumulator = typename WarpTileIterator::Element; + + /// Output access size + static int const kElementsPerAccess = Visitor::kElementsPerAccess; + + /// Tensor reference to sync tensor + using SyncTensorRef = typename mctlass::TensorRef; + + /// Array type used by output functor + using AccumulatorAccessType = Array< + typename WarpTileIterator::Element, kElementsPerAccess>; + + /// Number of warps + using WarpCount = typename Base::WarpCount; + + static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1 ? Base::kFragmentsPerIteration : kPartitionsK; + static int constexpr kSmemPointerOffset = Base::SharedStorage::StorageShape::kCount / kSmemTiles; + + using SharedStorage = typename Base::SharedStorage; + +private: + + /// Loads fragment from shared memory aligned with output tensor + SharedLoadIterator shared_load_iterator_; + +public: + + /// Constructor + MCTLASS_DEVICE + EpilogueWithVisitor( + SharedStorage &shared_storage, ///< Shared storage object + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx ///< Id of thread within warp + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + shared_load_iterator_(shared_storage.reference(), thread_idx) + { + + } + + /// Streams the result to global memory + MCTLASS_DEVICE + void operator()( + Visitor & visitor, + AccumulatorTile const &accumulators) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + + visitor.begin_epilogue(); + + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + + //#pragma unroll(IterationsUnroll ? Visitor::kIterations : 1) + //mxcc not support using unroll with parentheses, and we would support with unroll with Visitor::kIterations someday + #pragma unroll 1 + for (int iter_idx = 0; iter_idx < Visitor::kIterations; ++iter_idx) { + + // + // Load the source + // + + visitor.begin_step(iter_idx); + + // + // Convert and store fragment + // + + __syncthreads(); + + acc2smem_source_needed>::push( + iter_idx, accum_fragment_iterator, this->warp_tile_iterator_); + + __syncthreads(); + + // + // Load fragments from shared memory + // + + typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK]; + + shared_load_iterator_.load(aligned_accum_fragment[0]); + + // If the number of k-slices is > 1 - perform a reduction amongst the k-slices + if (kPartitionsK > 1) { + + plus add_fragments; + + MCTLASS_PRAGMA_UNROLL + for ( int i = 1; i < kPartitionsK; ++i) { + shared_load_iterator_.add_pointer_offset(kSmemPointerOffset); + shared_load_iterator_.load(aligned_accum_fragment[i]); + aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]); + } + + shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * kSmemPointerOffset); + } + + // + // Iterate over output fragments + // + + AccumulatorAccessType const *accum_frag_ptr = + reinterpret_cast(&aligned_accum_fragment[0]); + + int const kAccumulatorFragmentCount = AccumulatorTile::kElements / (Visitor::kIterations * AccumulatorAccessType::kElements); + + MCTLASS_PRAGMA_UNROLL + for (int idx = 0; idx < kAccumulatorFragmentCount; ++idx) { + + int row_idx = idx / SharedLoadIterator::ThreadMap::Iterations::kColumn; + int col_idx = idx % SharedLoadIterator::ThreadMap::Iterations::kColumn; + + // Start a new row of the output fragment + if (!col_idx) { + visitor.begin_row(row_idx); + } + + visitor.visit( + iter_idx, + row_idx, + col_idx, + idx, + accum_frag_ptr[idx] + ); + + // End the row of the output fragment + if (col_idx + 1 == SharedLoadIterator::ThreadMap::Iterations::kColumn) { + visitor.end_row(row_idx); + } + } + + // + // Conclude the step + // + + visitor.end_step(iter_idx); + } + + visitor.end_epilogue(); + } + +private: + + + template + struct acc2smem_source_needed; + + template + struct acc2smem_source_needed> { + template + MCTLASS_DEVICE + static void helper(AccumulatorFragmentIterator accum_fragment_iterator, + WarpTileIterator &warp_tile_iterator) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Advance; i++) { + ++accum_fragment_iterator; + } + + typename AccumulatorFragmentIterator::Fragment accum_fragment; + accum_fragment_iterator.load(accum_fragment); + warp_tile_iterator.store(accum_fragment); + } + + MCTLASS_DEVICE + static void push(size_t pos, + AccumulatorFragmentIterator const &iterator_begin, + WarpTileIterator &warp_tile_iterator) { + int dummy[] = {(pos == Seq) && (helper(iterator_begin, warp_tile_iterator), 0)...}; + } + }; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Helper to create an EpilogueWithVisitor from an existing epilogue +template +struct EpilogueWithVisitorFromExistingEpilogue { + + using Epilogue = EpilogueWithVisitor< + Visitor_, + typename Existing_::Shape, + typename Existing_::WarpMmaOperator, + Existing_::kPartitionsK, + typename Existing_::AccumulatorFragmentIterator, + typename Existing_::WarpTileIterator, + typename Existing_::SharedLoadIterator, + typename Existing_::Padding, + Existing_::kFragmentsPerIteration, + IterationsUnroll + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_workspace.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_workspace.h new file mode 100644 index 0000000..4bc88c6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/epilogue_workspace.h @@ -0,0 +1,197 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs. + + This does not attempt to target any particular output layout. Instead, each threadblock + streams out its accumulator elements using 128b store operations. This assumes all threadblocks + have unique output tiles. + + The target data layout is: + - threadblock indices mapped to linear offsets as (m, n, k), where m is fastest-changing + - threadblock output space partitioned into warps; each warp's region is contiguous + - per-thread accumulators partitioned into 128b accesses + - output memory striped across the threads of a warp + + This enables very fast streaming of data, completely limited by the memory system. No predication + or data exchange is performed, and each threadblock is assumed to have a full region of memory + to write to. + + This epilogue establishes an upper bound for epilogue performance and is suitable for + reductions across the GEMM K dimension which require a separate workspace. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, ///< shape of accumulator tile (concept: MatrixShape) + int WarpCount, ///< number of warps + typename FragmentC_ ///< warp-level GEMM operator (concept: gemm::warp::Mma) +> +class EpilogueWorkspace { +public: + + using Shape = Shape_; + using FragmentC = FragmentC_; + using ElementC = typename FragmentC::value_type; + + static int const kWarpCount = WarpCount; + + /// Optimize for 128b accesses + static int const kAccessSizeInBits = 128; + + /// Warp size from the perspective of memory operations + static int const kWarpSize = 32; + + /// Vector length of accesses + static int const kElementsPerAccess = + kAccessSizeInBits / sizeof_bits::value; + + /// Number of stores per thread + static int const kIterations = FragmentC::kElements / kElementsPerAccess; + + static_assert( + !(FragmentC::kElements % kElementsPerAccess), + "The number of accumulators must be divisible by the access size."); + + /// Total number of vectorized accesses in warp (in units of vector) + static int const kWarpAccesses = kIterations * kWarpSize; + + /// Total number of vectorized accesses in threadblock tile (in units of vector) + static int const kThreadblockAccesses = kWarpAccesses * kWarpCount; + + /// Parameters structure + struct Params { + + /// Pointer to C matrix + ElementC *ptr_C; + + /// Stride between tiles along the GEMM N dimension (in units of vectors) + int stride_n; + + /// Stride between tiles along the GEMM K dimension (in units of vectors) + int stride_k; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params( + ElementC *ptr_C, ///< Pointer to C matrix + int stride_n_, ///< Stride between tiles along the GEMM N dimension (in units of ElementC) + int stride_k_ ///< Stride between tiles along the GEMM K dimension (in units of ElementC) + ): + ptr_C(ptr_C), stride_n(stride_n_ / kElementsPerAccess), stride_k(stride_k_ / kElementsPerAccess) { + + } + }; + + /// Shared storage allocation needed by the epilogue + struct SharedStorage { + // Intentionally empty + }; + +private: + + struct alignas((kAccessSizeInBits / 8)) AccessType { + Array storage; + }; + + /// Constant reference to parameters object + AccessType *pointer_; + + /// Stride between tiles along the n dimension (in vectors) + int stride_n_; + + /// Stride between tiles along the k dimension (in vectors) + int stride_k_; + +public: + + /// Constructor + MCTLASS_DEVICE + EpilogueWorkspace( + Params const ¶ms, ///< Host-constructable params object + SharedStorage &, ///< Shared storage object + int warp_idx, ///< ID of warp within threadblock + int lane_idx ///< Id of thread within warp + + ): + pointer_(reinterpret_cast(params.ptr_C)), + stride_n_(params.stride_n), + stride_k_(params.stride_k) { + + // Add per-thread offset + pointer_ += lane_idx + warp_idx * kWarpAccesses; + } + + /// Streams the result to global memory + MCTLASS_DEVICE + void operator()( + mctlass::gemm::GemmCoord problem_size, ///< Problem size of GEMM (units of ElementC) + mctlass::gemm::GemmCoord tb_tile_coord, ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + FragmentC const &accum) { ///< Accumulator tile + + // Compute offset for entire threadblock (note, per-thread offset has been folded in already) + AccessType *pointer = pointer_ + + tb_tile_coord.m() * kThreadblockAccesses + + tb_tile_coord.n() * stride_n_ + + tb_tile_coord.k() * stride_k_; + + // Cast to vectorized view of accumulator fragments + AccessType const * src_pointer = reinterpret_cast(&accum); + + // Write out accumulators at full speed + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kIterations; ++i) { + pointer[i * kWarpSize] = src_pointer[i]; + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/interleaved_epilogue.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/interleaved_epilogue.h new file mode 100644 index 0000000..cc6a3e0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/interleaved_epilogue.h @@ -0,0 +1,432 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/vector.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/tensor_coord.h" +#include "mctlass/aligned_buffer.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_iterator.h" + +#include "mctlass/epilogue/threadblock/epilogue_base_streamk.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +template +MCTLASS_DEVICE +T shfl_sync(unsigned long mask, T var, unsigned int lane, int width = 64) { + return var; +} + +MCTLASS_DEVICE +int32_t shfl_sync(unsigned long mask, int32_t var, unsigned int lane, int width = 64) { + return __shfl_sync(mask, var, lane, width); +} + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Epilogue operator without splitk +template < + /// Shape of threadblock tile (concept: GemmShape) + typename Shape_, + /// Warp-level MMA operator (concept: gemm::warp::MmaTensorOp) + typename WarpMmaOperator_, + /// Number of partitions of the K dimension + int PartitionsK, + /// Tile iterator reading and writing output tensors + typename OutputTileIterator_, + /// Fragment iterator selecting accumulators + typename AccumulatorFragmentIterator_, + /// Output operator + typename OutputOp_, + /// Number of interleaved k + int InterleavedK> +class InterleavedEpilogue : + public EpilogueBaseStreamK< + Shape_, + PartitionsK, + WarpMmaOperator_, + AccumulatorFragmentIterator_> +{ + public: + + using BaseStreamK = EpilogueBaseStreamK< + Shape_, + PartitionsK, + WarpMmaOperator_, + AccumulatorFragmentIterator_>; + + using Shape = Shape_; + using WarpMmaOperator = WarpMmaOperator_; + static int const kPartitionsK = PartitionsK; + using AccumulatorFragmentIterator = AccumulatorFragmentIterator_; + using OutputTileIterator = OutputTileIterator_; + using OutputOp = OutputOp_; + + /// The complete warp-level accumulator tile + using AccumulatorTile = typename AccumulatorFragmentIterator::AccumulatorTile; + + /// Fragment type used by the accumulator tile's fragment iterator + using AccumulatorFragment = typename AccumulatorFragmentIterator::Fragment; + + /// Accumulator element + using ElementAccumulator = typename AccumulatorTile::Element; + + /// Output element + using ElementOutput = typename OutputTileIterator::Element; + + /// Output access size + static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess; + + /// Tensor reference to destination tensor + using TensorRef = typename OutputTileIterator::TensorRef; + + /// Tensor reference to sync tensor + using SyncTensorRef = + typename mctlass::TensorRef; + + /// Const tensor reference to source tensor + using ConstTensorRef = typename OutputTileIterator::ConstTensorRef; + + /// Array type used to output + using OutputAccessType = Array; + + /// Array type used by output functor + using AccumulatorAccessType = + Array; + + /// Number of warps + using WarpCount = + gemm::GemmShape; + + public: + + static_assert(OutputTileIterator::kElementsPerAccess, + "This must not be zero."); + + static_assert(!(OutputTileIterator::Fragment::kElements % + OutputTileIterator::kElementsPerAccess), + "Divisibility"); + +public: + + /// Aspect for when epilogue source is not needed + struct SourceAspectNotNeeded + { + /// Constructor + MCTLASS_DEVICE + SourceAspectNotNeeded() + {} + + /// Invoke the output functor over each vector of output + MCTLASS_DEVICE + void apply_output_operator( + typename OutputTileIterator::Fragment &output_fragment, + OutputOp const &output_op, + typename AccumulatorFragmentIterator::Fragment const &aligned_accum_fragment) + { + OutputAccessType *output_frag_ptr = + reinterpret_cast(&output_fragment); + + AccumulatorAccessType const *compute_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) + { + // Call the output operator + output_frag_ptr[i] = output_op(compute_frag_ptr[i]); + } + } + }; + + + /// Aspect for when epilogue source is needed + struct SourceAspectNeeded + { + OutputTileIterator source_iterator; + + typename OutputTileIterator::Fragment source_fragment; + + /// Invoke the output functor over each vector of output + MCTLASS_DEVICE + static void apply_output_operator( + typename OutputTileIterator::Fragment &output_fragment, + OutputOp const &output_op, + typename AccumulatorFragmentIterator::Fragment const &aligned_accum_fragment, + typename OutputTileIterator::Fragment const &source_fragment) + { + OutputAccessType *output_frag_ptr = + reinterpret_cast(&output_fragment); + + AccumulatorAccessType const *compute_frag_ptr = + reinterpret_cast(&aligned_accum_fragment); + + OutputAccessType const *source_frag_ptr = + reinterpret_cast(&source_fragment); + + int const kOutputOpIterations = + OutputTileIterator::Fragment::kElements / OutputTileIterator::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kOutputOpIterations; ++i) + { + // Call the output operator + output_frag_ptr[i] = output_op(compute_frag_ptr[i], source_frag_ptr[i]); + } + } + + /// Constructor + MCTLASS_DEVICE + SourceAspectNeeded(OutputTileIterator source_iterator) : + source_iterator(source_iterator) + { + source_fragment.clear(); + } + + /// Invoke the output functor over each vector of output + MCTLASS_DEVICE + void apply_output_operator( + typename OutputTileIterator::Fragment &output_fragment, + OutputOp const &output_op, + typename AccumulatorFragmentIterator::Fragment const &aligned_accum_fragment) + { + // Load addend source fragment from global memory + source_iterator.load(source_fragment); + ++source_iterator; + + apply_output_operator(output_fragment, output_op, aligned_accum_fragment, source_fragment); + } + }; + + /// Shared storage allocation needed by the epilogue + struct SharedStorage {}; + + + public: + /// Constructor + MCTLASS_DEVICE + InterleavedEpilogue( + SharedStorage &shared_storage, ///< Shared storage object + int thread_idx, ///< ID of a thread within the threadblock + int warp_idx, ///< ID of warp within threadblock + int lane_idx) ///< Id of thread within warp + : + BaseStreamK(thread_idx) + {} + /// Aggregates the accumulator sets shared by peer blocks in the global workspace, + /// performing epilogue computations, writing to output + MCTLASS_DEVICE + void reduce( + int peer_idx_begin, + int peer_idx_end, + int reduce_fragment_idx, + void *element_workspace, + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + OutputTileIterator source_iterator) ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles) + { + // Redcuce peer accumulator fragments into one fragment + AccumulatorFragment accum_fragment; + BaseStreamK::reduce(accum_fragment, peer_idx_begin, peer_idx_end, reduce_fragment_idx, element_workspace); + + // Source-fragment data (zero-initialized for scenarios where the + // output operator allows us to skip loading it from global input) + typename OutputTileIterator::Fragment source_fragment; + source_fragment.clear(); + + if (output_op.is_source_needed()) + { + source_iterator += reduce_fragment_idx; + source_iterator.load(source_fragment); + } + + // Compute the output result + typename OutputTileIterator::Fragment output_fragment; + + // Apply the output operator + SourceAspectNeeded::apply_output_operator(output_fragment, output_op, accum_fragment, source_fragment); + + // Store the final result + destination_iterator += reduce_fragment_idx; + destination_iterator.store(output_fragment); + } + + + /// Perform the epilogue computations and stream the result to global memory. + MCTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators) ///< Complete warp-level accumulator tile + { + operator()(output_op, destination_iterator, accumulators, SourceAspectNotNeeded()); + } + + /// Streams the result to global memory + MCTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator) { ///< Threadblock tile coordina te in GEMM (in units of threadblock tiles) + + Array *accum = reinterpret_cast *>(&accumulators); + + const int cur_lane_id = __lane_id() & 31; + const int src_lane_id = 2 * (cur_lane_id % 4) + ((cur_lane_id / 16) * 16); + const int index_n = (cur_lane_id % 16) / 4; + Array x, y, z, w; + //For s8n_s8t_s8n cases we need to shfl value layout same with cuda, not a good choice + for (int i = 0; i < accumulators.size() / 4; ++i) { + + for (int j = 0; j < 4; ++j) { + x[j] = shfl_sync(UINT64_MAX, accum[i][j], src_lane_id); + y[j] = shfl_sync(UINT64_MAX, accum[i][j], (src_lane_id + 1)); + z[j] = shfl_sync(UINT64_MAX, accum[i][j], (src_lane_id + 32)); + w[j] = shfl_sync(UINT64_MAX, accum[i][j], (src_lane_id + 33)); + } + + accum[i][0] = x[index_n]; + accum[i][1] = y[index_n]; + accum[i][2] = z[index_n]; + accum[i][3] = w[index_n]; + + } + + if (!output_op.is_source_needed()) { + operator()(output_op, destination_iterator, accumulators, SourceAspectNeeded(source_iterator)); + } + else { + operator()(output_op, destination_iterator, accumulators, SourceAspectNotNeeded()); + } + } + + /// Perform the epilogue computations and stream the result to global memory. Implements a + /// single codepath, regardless of whether the output op requires addend data to be loaded + MCTLASS_DEVICE + void unified( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + OutputTileIterator source_iterator ) ///< Tile iterator for addend source + { + if (!output_op.is_source_needed()) + { + source_iterator.clear_mask(); + __syncthreads(); // Dummy (CUDA 11.0) + } + + operator()(output_op, destination_iterator, accumulators, SourceAspectNeeded(source_iterator)); + } + + /// Streams the result to global memory + template + MCTLASS_DEVICE + void operator()( + OutputOp const &output_op, ///< Output operator + OutputTileIterator destination_iterator, ///< Tile iterator for destination + AccumulatorTile const &accumulators, ///< Complete warp-level accumulator tile + SourceAspect source) + { + + // + // Iterator over warp-level accumulator fragment + // + + AccumulatorFragmentIterator accum_fragment_iterator(accumulators); + + // + // Iterate over accumulator tile + // + + MCTLASS_PRAGMA_UNROLL + for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) { + + // + // Convert fragment + // + + typename AccumulatorFragmentIterator::Fragment accum_fragment; + + accum_fragment_iterator.load(accum_fragment); + ++accum_fragment_iterator; + + // + // Compute the output result + // + + typename OutputTileIterator::Fragment output_fragment; + source.apply_output_operator(output_fragment, output_op, accum_fragment); + + // + // Store the final result + // + + destination_iterator.set_iteration_index(iter); + destination_iterator.store(output_fragment); + ++destination_iterator; + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/output_iterator_parameter.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/output_iterator_parameter.h new file mode 100644 index 0000000..5488b4f --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/output_iterator_parameter.h @@ -0,0 +1,92 @@ +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/conv/convolution.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/tensor_ref.h" + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +template< + typename TensorLayout_, ///! The original output tensor layout + typename OutputIteratorLayout_, ///! Layout used by epilogue output iterator + typename TensorRef_, ///! Input tensor to epilogue output iterator + conv::Operator ConvOperator, ///! Convolutional operator (Fprop, Dgrad, Wgrad) + typename ConvProblemSize_ ///! Convolutional operator on 2D or 3D problem +> +struct ConvOutputIteratorParameter { + + using TensorLayout = TensorLayout_; + using OutputIteratorLayout = OutputIteratorLayout_; + using OutputTensorCoord = typename OutputIteratorLayout::TensorCoord; + using TensorRef = TensorRef_; + static conv::Operator const kConvolutionalOperator = ConvOperator; + using ConvProblemSize = ConvProblemSize_; + + /// Wgrad stride idx for implicit gemm algorithm + // Conv2d row-major matrix (KxRSC) + // Conv3d row-major matrix (KxTRSC) + static int const kWgradStrideIdx = + platform::is_same::value ? 2 : 3; + + /// This chooses the appropriate stride element of the C tensor. + static int const kTensorStrideIdx = + (kConvolutionalOperator == conv::Operator::kWgrad ? kWgradStrideIdx : 0); + + + MCTLASS_HOST_DEVICE + static OutputIteratorLayout layout(const TensorRef & ref) { + return ref.stride(kTensorStrideIdx); + } + + MCTLASS_HOST_DEVICE + static OutputTensorCoord extent(ConvProblemSize problem_size) { + return conv::implicit_gemm_problem_size(kConvolutionalOperator, problem_size).mn(); + } + +}; + + + +template < + int InterleavedK, + typename TensorRef_, + conv::Operator ConvOperator, + typename ConvProblemSize_ +> +struct ConvOutputIteratorParameter< + layout::TensorNCxHWx, + layout::TensorNCxHWx, + TensorRef_, + ConvOperator, + ConvProblemSize_> +{ + + using TensorLayout = typename layout::TensorNCxHWx; + using OutputIteratorLayout = typename layout::TensorNCxHWx; + using OutputTensorCoord = typename OutputIteratorLayout::TensorCoord; + using TensorRef = TensorRef_; + static conv::Operator const kConvolutionalOperator = ConvOperator; + using ConvProblemSize = ConvProblemSize_; + + MCTLASS_HOST_DEVICE + static OutputIteratorLayout layout(const TensorRef & ref) { + return ref.stride(); + } + + MCTLASS_HOST_DEVICE + static OutputTensorCoord extent(ConvProblemSize problem_size) { + return problem_size.output_extent(); + } + +}; + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/output_tile_thread_map.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/output_tile_thread_map.h new file mode 100644 index 0000000..cfa0f6e --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/output_tile_thread_map.h @@ -0,0 +1,730 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Metaprogram for determining the mapping of output elements to threads for epilogue tiles. + + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/fast_math.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tuple defining point in output tile +template < + int Column, + int Row, + int Group, + int Cluster, + int Tile +> +struct OutputTileShape { + static int const kColumn = Column; + static int const kRow = Row; + static int const kGroup = Group; + static int const kCluster = Cluster; + static int const kTile = Tile; + + static int const kCount = kColumn * kRow * kGroup * kCluster * kTile; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template +struct OutputTileThreadMapHelpers { + + /// Determines the iteration index of a vector access according to the thread map + MCTLASS_HOST_DEVICE + static void iteration_index( + int &column_idx, + int &row_idx, + int &group_idx, + int &cluster_idx, + int &tile_idx, + int iter_idx) { + + column_idx = iter_idx % Iterations::kColumn; + int residual = iter_idx / Iterations::kColumn; + + row_idx = residual % Iterations::kRow; + residual = residual / Iterations::kRow; + + group_idx = residual % Iterations::kGroup; + residual = residual / Iterations::kGroup; + + cluster_idx = residual % Iterations::kCluster; + tile_idx = residual / Iterations::kCluster; + } + + /// Computes the offset of a given vector access + MCTLASS_HOST_DEVICE + static MatrixCoord iteration_offset(int iter_idx) { + + int column_idx; + int row_idx; + int group_idx; + int cluster_idx; + int tile_idx; + + iteration_index(column_idx, row_idx, group_idx, cluster_idx, tile_idx, iter_idx); + + return + MatrixCoord( + row_idx * Delta::kRow + + group_idx * Delta::kGroup + + cluster_idx * Delta::kCluster + + tile_idx * Delta::kTile, + + column_idx * Delta::kColumn); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + + +template < + typename ThreadMap_, + typename Shape_, + typename Iterations_, + typename Delta_, + typename Count_ +> +struct OutputTileThreadMap : public OutputTileThreadMapHelpers { + + /// Conventional thread map (concept: ThreadMap) + using ThreadMap = ThreadMap_; + + /// Number of threads participating in the operation + static int const kThreads = ThreadMap::kThreads; + + /// Number of scalar elements per access + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + + /// Shape of the tile + using Shape = Shape_; + + /// Iterations performed by each thread + using Iterations = Iterations_; + + /// Delta between accesses + using Delta = Delta_; + + /// Number of iterator iterations + using Count = Count_; + + /// Initial offset function + MCTLASS_HOST_DEVICE + static MatrixCoord initial_offset(int thread_idx) { + + using Index = typename layout::PitchLinearCoord::Index; + + layout::PitchLinearCoord coord = ThreadMap::initial_offset(thread_idx); + + Index cluster = coord.strided() / (Shape::kGroup * Shape::kRow); + Index cluster_residual = coord.strided() % (Shape::kGroup * Shape::kRow); + + Index group = cluster_residual / (Shape::kRow); + Index row = cluster_residual % (Shape::kRow); + + return MatrixCoord{ + row + group * Shape::kRow * Count::kRow + + cluster * Shape::kGroup * Count::kGroup * Shape::kRow * Count::kRow, + coord.contiguous() + }; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +/// RowArrangement determines how one or more warps cover a region of consecutive rows. +template < + typename Shape, + int WarpsRemaining, + int ElementsPerAccess, + int ElementSize, + bool Is2dTile +> +struct RowArrangement; + +/// RowArrangement in which each warp's access is a 1D tiled arrangement. +template < + typename Shape, + int WarpsRemaining, + int ElementsPerAccess, + int ElementSize +> +struct RowArrangement { + static int const kWarpSize = 32; + static int const kElementsPerAccess = ElementsPerAccess; + static int const kElementSize = ElementSize; + + static int const kIterationsRow = 1; + static int const kDeltaRow = 1; + static int const kIterationsColumn = Shape::kColumn / kElementsPerAccess / kWarpSize; + static int const kDeltaColumn = kWarpSize * kElementsPerAccess; + + static int const kAccessWidth = kWarpSize; + static int const kAccessRows = 1; + static int const kWarpPartitionsRow = 1; + static int const kWarpPartitionsColumn = WarpsRemaining; +}; + +/// RowArrangement in which each warp's access is a 2D tiled arrangement. +template < + typename Shape, + int WarpsRemaining, + int ElementsPerAccess, + int ElementSize +> +struct RowArrangement { + + static int const kMemoryAccessSize = 256; // Preferred access size + static int const kWarpSize = 32; + + static int const kElementsPerAccess = ElementsPerAccess; + static int const kElementSize = ElementSize; + + struct Detail { + static int const kShapeRow = Shape::kRow / WarpsRemaining; + static int const kShapeWidth = Shape::kColumn / kElementsPerAccess; + + static int const kTargetMemoryAccessWidth = + kMemoryAccessSize / (kElementsPerAccess * kElementSize / 8); + + static int const kTargetAccessRows = kWarpSize / kTargetMemoryAccessWidth; + }; + + static int const kAccessWidth = + (Detail::kTargetAccessRows > Detail::kShapeRow ? + kWarpSize / Detail::kShapeRow + : const_min( + Detail::kShapeWidth, + const_min(kWarpSize, kMemoryAccessSize / (kElementsPerAccess * kElementSize / 8)) + )); + + static int const kAccessRows = + (Detail::kTargetAccessRows > Detail::kShapeRow ? + Detail::kShapeRow + : const_min(Shape::kRow, kWarpSize / kAccessWidth)); + + static int const kIterationsRow = Detail::kShapeRow / kAccessRows; + static int const kDeltaRow = kAccessRows; + + static int const kIterationsColumn = Detail::kShapeWidth / kAccessWidth; + static int const kDeltaColumn = kAccessWidth * kElementsPerAccess; + + static_assert( kAccessWidth * kElementsPerAccess <= Shape::kColumn, "Accessing too many elements per access"); + static_assert( kIterationsColumn > 0, "Iteration Count Column must be > 0" ); + static_assert( kIterationsRow > 0, "Iteration Count Row must be > 0" ); + + static int const kWarpPartitionsRow = 1; + static int const kWarpPartitionsColumn = 1; +}; + +} + +//////////////////////////////////////////////////////////////////////////////// + +/// Template metaprogram for partitioning a 4D space across warps to achieve several performance +/// objectives: +/// +/// - coalesced memory accesses in units of 128 Byte lines +/// - minimal address arithmetic +/// - minimal predicate calculations +/// +template < + typename Shape_, + typename Count_, + int Threads, + int ElementsPerAccess, + int ElementSize +> +struct OutputTileOptimalThreadMap { + + using Shape = Shape_; + using Count = Count_; + + static int const kWarpSize = 32; + static int const kThreads = Threads; + static int const kWarpCount = kThreads / kWarpSize; + + static int const kElementsPerAccess = ElementsPerAccess; + static int const kElementSize = ElementSize; + + // + // Metaprogram computation + // + + struct Detail { + + // Clusters + static int const kIterationsCluster = + ((Shape::kCluster > kWarpCount) ? + Shape::kCluster / kWarpCount + : 1); + + static int const kDeltaCluster = + ((Shape::kCluster > kWarpCount) ? + Shape::kRow * Count::kRow * Shape::kGroup * Count::kGroup * Shape::kCluster / kIterationsCluster + : 1); + + static int const kCompactedDeltaCluster = + ((Shape::kCluster > kWarpCount) ? + Shape::kRow * Shape::kGroup * Shape::kCluster / kIterationsCluster + : 1); + + static int const kWarpPartitionsCluster = + ((Shape::kCluster > kWarpCount) ? + kWarpCount + : kWarpCount / Shape::kCluster); + + static int const kWarpsRemainingForGroups = + ((Shape::kCluster > kWarpCount) ? 1 : kWarpCount / Shape::kCluster); + + // Groups + static int const kIterationsGroup = + ((Shape::kGroup > kWarpsRemainingForGroups) ? + Shape::kGroup / kWarpsRemainingForGroups + : 1); + + static int const kDeltaGroup = + ((Shape::kGroup > kWarpsRemainingForGroups) ? + Shape::kRow * Count::kRow * Shape::kGroup / kIterationsGroup + : 1); + + static int const kCompactedDeltaGroup = + ((Shape::kGroup > kWarpsRemainingForGroups) ? + Shape::kRow * Shape::kGroup / kIterationsGroup + : 1); + + static int const kWarpPartitionsGroup = + ((Shape::kGroup > kWarpsRemainingForGroups) ? + 1 + : kWarpsRemainingForGroups / Shape::kGroup); + + static int const kWarpsRemainingForRows = + ((Shape::kGroup > kWarpsRemainingForGroups) ? + 1 + : kWarpsRemainingForGroups / Shape::kGroup); + + // Rows + using RowArrangement = detail::RowArrangement< + Shape, + kWarpsRemainingForRows, + kElementsPerAccess, + kElementSize, + (Shape::kRow > kWarpsRemainingForRows) + >; + + // Warp partitions + using WarpPartitions = OutputTileShape< + RowArrangement::kWarpPartitionsColumn, + RowArrangement::kWarpPartitionsRow, + kWarpPartitionsGroup, + kWarpPartitionsCluster, + 1>; + + static int const kAccessWidth = RowArrangement::kAccessWidth; + static int const kAccessRows = RowArrangement::kAccessRows; + }; + + // + // Output + // + + using Iterations = OutputTileShape< + Detail::RowArrangement::kIterationsColumn, + Detail::RowArrangement::kIterationsRow, + Detail::kIterationsGroup, + Detail::kIterationsCluster, + 1>; + + using Delta = OutputTileShape< + Detail::RowArrangement::kDeltaColumn, + Detail::RowArrangement::kDeltaRow, + Detail::kDeltaGroup, + Detail::kDeltaCluster, + 1>; + + /// Initial offset function + MCTLASS_DEVICE + static MatrixCoord initial_offset(int thread_idx) { + + int warp_idx = __shfl_sync(0xffffffff, (thread_idx % kThreads) / kWarpSize, 0); + int lane_idx = thread_idx % kWarpSize; + + // Compute warp location + int cluster_idx = warp_idx / Detail::WarpPartitions::kCluster; + int residual_cluster = warp_idx % Detail::WarpPartitions::kCluster; + + int group_idx = residual_cluster / Detail::WarpPartitions::kGroup; + int residual_group = residual_cluster % Detail::WarpPartitions::kGroup; + + int row_idx = residual_group / Detail::WarpPartitions::kRow; + int col_idx = residual_group % Detail::WarpPartitions::kRow; + + // Compute per-lane offset + int lane_row_offset = lane_idx / Detail::kAccessWidth; + int lane_col_offset = lane_idx % Detail::kAccessWidth; + + // Compute coordinate in output space + int cluster_offset = cluster_idx * Shape::kRow * Count::kRow * Shape::kGroup * Count::kGroup; + int group_offset = group_idx * Shape::kRow * Count::kRow; + int row_offset = row_idx * Iterations::kRow * Detail::kAccessRows; + int column_offset = col_idx * Iterations::kColumn * Detail::kAccessWidth * kElementsPerAccess; + + return MatrixCoord( + cluster_offset + group_offset + row_offset + lane_row_offset, + column_offset + lane_col_offset * kElementsPerAccess + ); + } + + static inline unsigned int __shfl_sync_host(unsigned long mask, + unsigned long var, int srcLane, + int width = warpSize) { + + static_assert(sizeof(unsigned long) == 2 * sizeof(unsigned int), ""); + static_assert(sizeof(unsigned long) == sizeof(uint64_t), ""); + + unsigned int tmp[2]; + std::memcpy(tmp, &var, sizeof(tmp)); + tmp[0] = __shfl_sync_host(mask, tmp[0], srcLane, width); + tmp[1] = __shfl_sync_host(mask, tmp[1], srcLane, width); + + uint64_t tmp0 = + (static_cast(tmp[1]) << 32ull) | static_cast(tmp[0]); + unsigned long tmp1; + std::memcpy(&tmp1, &tmp0, sizeof(tmp0)); + return tmp1; + +} + + static MatrixCoord initial_offset_host(int thread_idx) { + + int warp_idx = __shfl_sync_host(0xffffffff, thread_idx / kWarpSize, 0); + int lane_idx = thread_idx % kWarpSize; + + // Compute warp location + int cluster_idx = warp_idx / Detail::WarpPartitions::kCluster; + int residual_cluster = warp_idx % Detail::WarpPartitions::kCluster; + + int group_idx = residual_cluster / Detail::WarpPartitions::kGroup; + int residual_group = residual_cluster % Detail::WarpPartitions::kGroup; + + int row_idx = residual_group / Detail::WarpPartitions::kRow; + int col_idx = residual_group % Detail::WarpPartitions::kRow; + + // Compute per-lane offset + int lane_row_offset = lane_idx / Detail::kAccessWidth; + int lane_col_offset = lane_idx % Detail::kAccessWidth; + + // Compute coordinate in output space + int cluster_offset = cluster_idx * Shape::kRow * Count::kRow * Shape::kGroup * Count::kGroup; + int group_offset = group_idx * Shape::kRow * Count::kRow; + int row_offset = row_idx * Iterations::kRow * Detail::kAccessRows; + int column_offset = col_idx * Iterations::kColumn * Detail::kAccessWidth * kElementsPerAccess; + + return MatrixCoord( + cluster_offset + group_offset + row_offset + lane_row_offset, + column_offset + lane_col_offset * kElementsPerAccess + ); + } + + /// Computes the offset of a given vector access + MCTLASS_HOST_DEVICE + static MatrixCoord iteration_offset(int iter_idx) { + return OutputTileThreadMapHelpers::iteration_offset(iter_idx); + } + + /// Compacted thread map in which the 4D region is contiguous + struct CompactedThreadMap { + + + using Shape = Shape_; + + using TileShape = MatrixShape< + Shape::kTile * Shape::kCluster * Shape::kGroup * Shape::kRow, + Shape::kColumn + >; + + using Iterations = OutputTileShape< + Detail::RowArrangement::kIterationsColumn, + Detail::RowArrangement::kIterationsRow, + Detail::kIterationsGroup, + Detail::kIterationsCluster, + 1>; + + using Delta = OutputTileShape< + Detail::RowArrangement::kDeltaColumn, + Detail::RowArrangement::kDeltaRow, + Detail::kCompactedDeltaGroup, + Detail::kCompactedDeltaCluster, + 1>; + + /// Number of elements within each vector access + static int const kElementsPerAccess = ElementsPerAccess; + + /// Number of threads + static int const kThreads = Threads; + + /// Function to compute each thread's initial offset + MCTLASS_DEVICE + static MatrixCoord initial_offset(int thread_idx) { + + int warp_idx = __shfl_sync(0xffffffff, (thread_idx % kThreads) / kWarpSize, 0); + int lane_idx = thread_idx % kWarpSize; + + // Compute warp location + int cluster_idx = warp_idx / Detail::WarpPartitions::kCluster; + int residual_cluster = warp_idx % Detail::WarpPartitions::kCluster; + + int group_idx = residual_cluster / Detail::WarpPartitions::kGroup; + int residual_group = residual_cluster % Detail::WarpPartitions::kGroup; + + int row_idx = residual_group / Detail::WarpPartitions::kRow; + int col_idx = residual_group % Detail::WarpPartitions::kRow; + + // Compute per-lane offset + int lane_row_offset = lane_idx / Detail::kAccessWidth; + int lane_col_offset = lane_idx % Detail::kAccessWidth; + + // Compute coordinate in output space + int cluster_offset = cluster_idx * Shape::kRow * Shape::kGroup; + int group_offset = group_idx * Shape::kRow; + int row_offset = row_idx * Iterations::kRow * Detail::kAccessRows; + int column_offset = col_idx * Iterations::kColumn * Detail::kAccessWidth * kElementsPerAccess; + + MatrixCoord coord( + cluster_offset + group_offset + row_offset + lane_row_offset, + column_offset + lane_col_offset * kElementsPerAccess + ); + + return coord; + } + + static inline unsigned int __shfl_sync_host(unsigned long mask, + unsigned long var, int srcLane, + int width = warpSize) { + + static_assert(sizeof(unsigned long) == 2 * sizeof(unsigned int), ""); + static_assert(sizeof(unsigned long) == sizeof(uint64_t), ""); + + unsigned int tmp[2]; + std::memcpy(tmp, &var, sizeof(tmp)); + tmp[0] = __shfl_sync_host(mask, tmp[0], srcLane, width); + tmp[1] = __shfl_sync_host(mask, tmp[1], srcLane, width); + + uint64_t tmp0 = + (static_cast(tmp[1]) << 32ull) | static_cast(tmp[0]); + unsigned long tmp1; + std::memcpy(&tmp1, &tmp0, sizeof(tmp0)); + return tmp1; + + } + + static MatrixCoord initial_offset_host(int thread_idx) { + + int warp_idx = __shfl_sync_host(0xffffffff, thread_idx / kWarpSize, 0); + int lane_idx = thread_idx % kWarpSize; + + // Compute warp location + int cluster_idx = warp_idx / Detail::WarpPartitions::kCluster; + int residual_cluster = warp_idx % Detail::WarpPartitions::kCluster; + + int group_idx = residual_cluster / Detail::WarpPartitions::kGroup; + int residual_group = residual_cluster % Detail::WarpPartitions::kGroup; + + int row_idx = residual_group / Detail::WarpPartitions::kRow; + int col_idx = residual_group % Detail::WarpPartitions::kRow; + + // Compute per-lane offset + int lane_row_offset = lane_idx / Detail::kAccessWidth; + int lane_col_offset = lane_idx % Detail::kAccessWidth; + + // Compute coordinate in output space + int cluster_offset = cluster_idx * Shape::kRow * Shape::kGroup; + int group_offset = group_idx * Shape::kRow; + int row_offset = row_idx * Iterations::kRow * Detail::kAccessRows; + int column_offset = col_idx * Iterations::kColumn * Detail::kAccessWidth * kElementsPerAccess; + + MatrixCoord coord( + cluster_offset + group_offset + row_offset + lane_row_offset, + column_offset + lane_col_offset * kElementsPerAccess + ); + + return coord; + } + }; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template metaprogram for partitioning a 3D interleaved layout across warps +/// to achieve several performance objectives: +/// +/// - coalesced memory accesses in units of 64 Byte lines +/// - minimal address arithmetic +/// - minimal predicate calculations +/// +template +struct InterleavedOutputTileThreadMap { + using WarpCount = WarpCount_; + + static int const kWarpSize = 64; + static int const kThreads = Threads; + static int const kWarpCount = kThreads / kWarpSize; + + static int const kElementsPerAccess = ElementsPerAccess; + static int const kElementSize = ElementSize; + + // + // Metaprogram computation + // + + struct Detail {}; + + // + // Output + // + + using Iterations = Iterations_; + + using Delta = layout::PitchLinearShape<(kWarpSize / 2) * kElementsPerAccess, 1>; + + /// Initial offset function + MCTLASS_HOST_DEVICE + static layout::PitchLinearCoord initial_offset(int thread_idx) { + int warp_idx = thread_idx / kWarpSize; + int lane_idx = thread_idx % kWarpSize; + + // Compute warp location + layout::PitchLinearCoord warp_footprint{ + Delta::kContiguous * Iterations::kContiguous, + Delta::kStrided * Iterations::kStrided}; + + layout::PitchLinearCoord warp_offset{warp_idx % WarpCount::kContiguous, + warp_idx / WarpCount::kContiguous}; + + // Compute per-lane offset + layout::PitchLinearCoord thread_offset_in_warp{ + lane_idx * kElementsPerAccess, 0}; + + layout::PitchLinearCoord thread_offset_in_threadblock_tile = + warp_footprint * warp_offset + thread_offset_in_warp; + + return thread_offset_in_threadblock_tile; + } +}; + + +//////////////////////////////////////////////////////////////////////////////// + +/// Template metaprogram for partitioning a 4D interleaved layout across warps +/// to achieve several performance objectives: +/// +/// - coalesced memory accesses in units of 64 Byte lines +/// - minimal address arithmetic +/// - minimal predicate calculations +/// +template +struct InterleavedConvOutputTileThreadMap { + using WarpCount = WarpCount_; + + static int const kWarpSize = 64; + static int const kThreads = Threads; + static int const kWarpCount = kThreads / kWarpSize; + + static int const kElementsPerAccess = ElementsPerAccess; + static int const kElementSize = ElementSize; + + // + // Metaprogram computation + // + + struct Detail {}; + + // + // Output + // + + using Iterations = Iterations_; + + using Delta = MatrixShape; + + /// Initial offset function + MCTLASS_HOST_DEVICE + static MatrixCoord initial_offset(int thread_idx) { + int warp_idx = thread_idx / kWarpSize; + int lane_idx = thread_idx % kWarpSize; + + // Compute warp location + MatrixCoord warp_footprint{ + Delta::kRow * Iterations::kRow, + Delta::kColumn * Iterations::kColumn, + }; + + MatrixCoord warp_offset{warp_idx % WarpCount::kRow, + warp_idx / WarpCount::kRow}; + + // Compute per-lane offset + MatrixCoord thread_offset_in_warp{lane_idx / 4, + (lane_idx % 4) * kElementsPerAccess}; + + MatrixCoord thread_offset_in_threadblock_tile = + warp_footprint * warp_offset + thread_offset_in_warp; + + return thread_offset_in_threadblock_tile; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator.h new file mode 100644 index 0000000..5b96094 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator.h @@ -0,0 +1,1343 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/permute.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/epilogue/threadblock/output_tile_thread_map.h" +#include "mctlass/arch/arch.h" +#include "mctlass/arch/memory.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////// + +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load and store output tile from global memory in epilogue. +/// +/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator +/// +template < + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_, ///< Element data type + bool ScatterD = false, ///< Scatter D operand or not + typename PermuteDLayout = layout::NoPermute, ///< Permute D operand or not + bool UseCUDAStore = false +> +class PredicatedTileIterator { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static bool constexpr PermuteD = !layout::is_trivial_permute; + + static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0"); + static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0"); + static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0"); + static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0"); + + /// Fragment object + using Fragment = Array< + Element, + ThreadMap::Iterations::kColumn * + ThreadMap::Iterations::kRow * + ThreadMap::Iterations::kGroup * + ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + + /// Uses a non-template class + struct Params : PredicatedTileIteratorParams { + using Base = PredicatedTileIteratorParams; + + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Layout const &layout): + PredicatedTileIteratorParams( + layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess, + make_OutputTileThreadMapDesc() + ) + { } + + MCTLASS_HOST_DEVICE + Params(Base const &base) : + Base(base) { } + }; + + /// Mask object + struct Mask { + + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + MCTLASS_HOST_DEVICE + Mask() { + enable(); + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_HOST_DEVICE void clear() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = false; + } + } + + ///< MCTLASS_HOST_DEVICE enables all accesses guarded by mask + MCTLASS_DEVICE void enable() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = true; + } + } + }; + +private: + + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + PredicatedTileIteratorParams params_; + + /// Byte-level pointer. This pointer is usually for both load() and store(), unless PermuteD is performed. When having PermuteD, byte_pointer_ is only for load(). + uint8_t *byte_pointer_; + + /// Byte-level pointer for store(). Due to PermuteD Op, store_byte_pointer_ may be with different address computation compared to byte_pointer_. + uint8_t *store_byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Index extent_row_; + + /// Extent of the matrix tile in rows + Index extent_column_; + + /// A thread's starting row position (assuming steady-state predicates have been computed) + Index thread_start_row_; + + /// A thread's starting column + Index thread_start_column_; + + /// Internal state counter + int state_[3]; + + /// Scatter indices + int const *indices_; + + /// PermuteDLayout + PermuteDLayout permute_layout_; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides"); + +private: + + // + // Methods + // + +public: + + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + PredicatedTileIterator( + PredicatedTileIteratorParams const & params, + Element *pointer, + TensorCoord extent, + int thread_idx, + TensorCoord threadblock_offset = TensorCoord(), + int const *indices = nullptr + ): + params_(params), indices_(indices), + permute_layout_(PitchLinearCoord(extent.column(), extent.row()), params_.stride * kElementsPerAccess / sizeof(AccessType)) + { + + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + extent_row_ = extent.row(); + extent_column_ = extent.column(); + + thread_start_row_ = thread_offset.row(); + thread_start_column_ = thread_offset.column(); + + // Initialize predicates + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + + mask_.predicates[c] = ((thread_offset.column() + + ThreadMap::Delta::kColumn * c) < extent.column()); + } + + // Null pointer performs no accesses + if (!pointer) { + mask_.clear(); + } + + if (ScatterD && !indices) { + mask_.clear(); + } + + // Initialize byte_pointer_ + byte_pointer_ = reinterpret_cast(pointer) + + LongIndex(thread_offset.row()) * LongIndex(params_.stride) + + LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; + + if (ScatterD) { + byte_pointer_ = reinterpret_cast(pointer) + + LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; + } + + // store_byte_pointer_ is set to be the same with byte_pointer_ unless PermuteD is used. + store_byte_pointer_ = PermuteD ? reinterpret_cast(pointer) : byte_pointer_; + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + store_byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, int64_t byte_offset) const { + + uint8_t *byte_pointer = byte_pointer_; + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + AccessType *memory_pointer = reinterpret_cast(byte_pointer + byte_offset); + + if (ScatterD && row_guard) { + assert(indices_); + + memory_pointer = reinterpret_cast(byte_pointer + byte_offset + + LongIndex(indices_[row_offset + thread_start_row_]) * LongIndex(params_.stride)); + } + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + bool guard = row_guard && mask_.predicates[column]; + + mctlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + + column], + (void *)&memory_pointer[column * ThreadMap::Delta::kColumn / + kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + if (!ScatterD) { + byte_pointer += params_.increment_row; + } + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) const { + + load_with_byte_offset(frag, 0); + } + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) const { + uint8_t *byte_pointer = store_byte_pointer_; + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + AccessType *memory_pointer = reinterpret_cast(byte_pointer + byte_offset); + + if (ScatterD && row_guard) { + assert(indices_); + + memory_pointer = reinterpret_cast(byte_pointer + byte_offset + + LongIndex(indices_[row_offset + thread_start_row_]) * LongIndex(params_.stride)); + } + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + bool guard = row_guard && mask_.predicates[column]; + + if (PermuteD) { + + int col_offset = column * ThreadMap::Delta::kColumn; + + int col = col_offset + thread_start_column_; + int row = row_offset + thread_start_row_; + + // Locate memory_pointer + memory_pointer = reinterpret_cast(byte_pointer + byte_offset + + permute_layout_(PitchLinearCoord(col, row)) * sizeof(AccessType) / kElementsPerAccess); + } + + if (UseCUDAStore) { + if (guard) { + memory_pointer[0] = + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column]; + } + } else { + mctlass::arch::global_store( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void *)&memory_pointer[0], + guard); + } + + if (!PermuteD) { + memory_pointer += (ThreadMap::Delta::kColumn / kElementsPerAccess); + } + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + if (!ScatterD && !PermuteD) { + byte_pointer += params_.increment_row; + } + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) const { + if (threadIdx.x < kThreads) { + store_with_byte_offset(frag, 0); + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void downsample_load_with_byte_offset(Fragment &frag, int64_t byte_offset, int convolution_P, int convolution_Q, int add_P, int add_Q, int problem_N) const { + + uint8_t *byte_pointer = byte_pointer_; + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + int output_row = row_offset + thread_start_row_; + int output_N = output_row / (convolution_P * convolution_Q); + int output_PQ = output_row % (convolution_P * convolution_Q); + int output_P = output_PQ / convolution_Q; + int output_Q = output_PQ % convolution_Q; + + int input_row = output_N * 2 * convolution_P * 2 * convolution_Q + + (2 * output_P + add_P) * 2 * convolution_Q + 2 * output_Q + add_Q; + + int64_t byte_offset = (input_row-output_row)*problem_N*sizeof(float); + + AccessType *memory_pointer = reinterpret_cast(byte_pointer + byte_offset); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + bool guard = row_guard && mask_.predicates[column]; + + mctlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + + column], + (void *)&memory_pointer[column * ThreadMap::Delta::kColumn / + kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + byte_pointer += params_.increment_row; + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void upsample_load_with_byte_offset(Fragment &frag, int64_t byte_offset, int convolution_P, int convolution_Q, int add_P, int add_Q, int problem_N) const { + + uint8_t *byte_pointer = byte_pointer_; + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + int output_row = row_offset + thread_start_row_; + int output_N = output_row / (convolution_P * convolution_Q); + int output_PQ = output_row % (convolution_P * convolution_Q); + int output_P = output_PQ / convolution_Q; + int output_Q = output_PQ % convolution_Q; + int row_add_P = add_P; + int row_add_Q = add_Q; + if (output_P > convolution_P - 2) row_add_P = 0; + if (output_Q > convolution_Q - 2) row_add_Q = 0; + + int input_row = output_N * (convolution_P/2) * (convolution_Q/2) + + ((output_P + row_add_P)/2) * (convolution_Q/2) + (output_Q + row_add_Q)/2; + + int64_t byte_offset = (input_row-output_row)*problem_N*sizeof(float); + + AccessType *memory_pointer = reinterpret_cast(byte_pointer + byte_offset); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + bool guard = row_guard && mask_.predicates[column]; + + mctlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + + column], + (void *)&memory_pointer[column * ThreadMap::Delta::kColumn / + kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + byte_pointer += params_.increment_row; + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + MCTLASS_DEVICE + MatrixCoord thread_start() const { + return MatrixCoord(thread_start_row_, thread_start_column_); + } + + /// Need to get the thread start row from the tile iterator + MCTLASS_DEVICE + int32_t thread_start_row() const { + return thread_start_row_; + } + + /// Need to get the thread start row from the tile iterator + MCTLASS_DEVICE + int32_t thread_start_column() const { + return thread_start_column_; + } + + /// Extent of the matrix in rows + MCTLASS_DEVICE + Index extent_row() const { + return extent_row_; + } + + /// Extent of the matrix in columns + MCTLASS_DEVICE + Index extent_column() const { + return extent_column_; + } + + /// Advances to the next position to load or store + MCTLASS_HOST_DEVICE + PredicatedTileIterator &operator++() { + + ++state_[0]; + + if (!ScatterD && !PermuteD) { + store_byte_pointer_ += params_.advance_row; + } + + if (!ScatterD) { + byte_pointer_ += params_.advance_row; + } + + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + + state_[0] = 0; + ++state_[1]; + byte_pointer_ += params_.advance_group; + store_byte_pointer_ += params_.advance_group; + + thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * + ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + + state_[1] = 0; + ++state_[2]; + byte_pointer_ += params_.advance_cluster; + store_byte_pointer_ += params_.advance_cluster; + + thread_start_row_ += ThreadMap::Count::kGroup * + ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { + state_[2] = 0; + byte_pointer_ += params_.advance_tile; + store_byte_pointer_ += params_.advance_tile; + thread_start_row_ += ThreadMap::Shape::kGroup * ThreadMap::Shape::kRow + * ThreadMap::Shape::kCluster * ThreadMap::Shape::kTile; + } + } + } + + return *this; + } + + /// Advances a number of positions to load or store + MCTLASS_HOST_DEVICE + PredicatedTileIterator &operator+=(int increment) + { + // Row + state_[0] += increment; + int increment_row = state_[0] / ThreadMap::Count::kRow; + state_[0] = state_[0] % ThreadMap::Count::kRow; + + byte_pointer_ += (params_.advance_row * increment); + store_byte_pointer_ += (params_.advance_row * increment); + thread_start_row_ += (ThreadMap::Shape::kRow * increment); + + // Group + state_[1] += increment_row; + int increment_group = state_[1] / ThreadMap::Count::kGroup; + state_[1] = state_[1] % ThreadMap::Count::kGroup; + + byte_pointer_ += (params_.advance_group * increment_row); + store_byte_pointer_ += (params_.advance_group * increment_row); + thread_start_row_ += + (ThreadMap::Shape::kGroup - 1) * + ThreadMap::Shape::kRow * + ThreadMap::Count::kRow * + increment_row; + + + // Cluster + state_[2] += increment_group; + int increment_cluster = state_[2] / ThreadMap::Count::kCluster; + state_[2] = state_[2] % ThreadMap::Count::kCluster; + + byte_pointer_ += (params_.advance_cluster * increment_group); + store_byte_pointer_ += (params_.advance_cluster * increment_group); + thread_start_row_ += + ThreadMap::Count::kGroup * + ThreadMap::Shape::kGroup * + ThreadMap::Count::kRow * + ThreadMap::Shape::kRow * + increment_group; + + // Tile + byte_pointer_ += (params_.advance_tile * increment_cluster); + store_byte_pointer_ += (params_.advance_tile * increment_cluster); + thread_start_row_ += + ThreadMap::Shape::kGroup * + ThreadMap::Shape::kRow * + ThreadMap::Shape::kCluster * + ThreadMap::Shape::kTile * + increment_cluster; + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_DEVICE void clear_mask() { + mask_.clear(); + } + + ///< Efficiently enables all accesses guarded by mask + MCTLASS_DEVICE void enable_mask() { + mask_.enable(); + } + + ///< Sets the mask + MCTLASS_DEVICE void get_mask(Mask &mask) const { + mask = mask_; + } + + ///< Sets the mask + MCTLASS_DEVICE void set_mask(Mask const &mask) { + mask_ = mask; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load output tile from global memory in epilogue. +/// +/// Satisfies: ReadableTileIterator | InterleavedPredicatedTileIterator | ForwardTileIterator +/// +template < + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_, ///< Element data type + int InterleavedN ///< Number of Interleaved N +> +class InterleavedPredicatedTileIterator { +public: + using ThreadMap = ThreadMap_; + + using Element = Element_; + + using Layout = layout::ColumnMajorInterleaved; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = layout::PitchLinearCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Iterations::kCount; + + /// Fragment object + using Fragment = Array; + + /// Memory access size + using AccessType = AlignedArray; + + /// Uses a non-template class + struct Params : InterleavedPredicatedTileIteratorParams { + using Base = InterleavedPredicatedTileIteratorParams; + + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Layout const &layout): + Base( + layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess, + make_InterleavedPredicatedTileIteratorDesc() + ) { } + + MCTLASS_HOST_DEVICE + Params(Base const &base) : + Base(base) { } + }; + + /// Mask object + struct Mask { + static int const kCount = (ThreadMap::Iterations::kContiguous < 8) + ? 8 + : ThreadMap::Iterations::kContiguous; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + MCTLASS_HOST_DEVICE + Mask() { + enable(); + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_HOST_DEVICE void clear() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = false; + } + } + + ///< MCTLASS_HOST_DEVICE enables all accesses guarded by mask + MCTLASS_DEVICE void enable() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = true; + } + } + }; + +private: + + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + Params params_; + + /// Byte-level pointer + uint8_t *byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in columns + Index extent_col_; + + /// A thread's starting column position (assuming steady-state predicates have + /// been computed) + Index thread_start_col_; + + /// Internal iteration counter + int iteration_contiguous_; + + int iteration_strided_; + +private: + + // + // Methods + // + +public: + + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + InterleavedPredicatedTileIterator( + Params const & params, + Element *pointer, + TensorCoord extent, + int thread_idx, + TensorCoord threadblock_offset, + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ): + params_(params) { + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + + TensorCoord(threadblock_offset.contiguous() * InterleavedN, + threadblock_offset.strided() / InterleavedN); + + extent_col_ = extent.strided() / InterleavedN; + thread_start_col_ = thread_offset.strided(); + + // Initialize predicates + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + mask_.predicates[c] = + ((thread_offset.contiguous() + ThreadMap::Delta::kContiguous * c) < + (extent.contiguous() * InterleavedN)); + } + + // Initialize pointer + byte_pointer_ = reinterpret_cast(pointer) + + LongIndex(thread_offset.strided()) * LongIndex(params_.stride) + + LongIndex(thread_offset.contiguous()) * sizeof(AccessType) / kElementsPerAccess; + + // Initialize internal state counter + iteration_contiguous_ = iteration_strided_ = 0; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + + uint8_t *byte_pointer = byte_pointer_; + AccessType *frag_ptr = reinterpret_cast(&frag); + AccessType *memory_pointer = reinterpret_cast(byte_pointer); + + int col_offset = iteration_strided_ * ThreadMap::Delta::kStrided; + + bool col_guard = ((thread_start_col_ + col_offset) < extent_col_); + + bool guard = col_guard && mask_.predicates[iteration_contiguous_]; + + mctlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + *frag_ptr, + (void *)memory_pointer, + guard); + } + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + if (__lane_id() >= 32) return; + uint8_t *byte_pointer = byte_pointer_; + AccessType const *frag_ptr = reinterpret_cast(&frag); + AccessType *memory_pointer = reinterpret_cast(byte_pointer); + + int col_offset = iteration_strided_ * ThreadMap::Delta::kStrided; + + bool col_guard = ((thread_start_col_ + col_offset) < extent_col_); + + bool guard = col_guard && mask_.predicates[iteration_contiguous_]; + + mctlass::arch::global_store( + *frag_ptr, (void *)memory_pointer, guard); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int iteration) { + iteration_contiguous_ = iteration % ThreadMap::Iterations::kContiguous; + iteration_strided_ = iteration / ThreadMap::Iterations::kContiguous; + } + + /// Advances to the next position to load or store + MCTLASS_HOST_DEVICE + InterleavedPredicatedTileIterator &operator++() { + + ++iteration_contiguous_; + byte_pointer_ += params_.advance_row; + + if (iteration_contiguous_ == ThreadMap::Iterations::kContiguous) { + + iteration_contiguous_ = 0; + ++iteration_strided_; + byte_pointer_ += params_.advance_column; + + if (iteration_strided_ == ThreadMap::Iterations::kStrided) { + iteration_strided_ = 0; + } + } + + return *this; + } + /// Advances a number of positions to load or store + MCTLASS_HOST_DEVICE + InterleavedPredicatedTileIterator &operator+=(int increment) + { + // Contiguous + iteration_contiguous_ += increment; + int increment_strided = iteration_contiguous_ / ThreadMap::Iterations::kContiguous; + iteration_contiguous_ = iteration_contiguous_ % ThreadMap::Iterations::kContiguous; + byte_pointer_ += (params_.advance_row * increment); + + // Strided + iteration_strided_ += increment_strided; + byte_pointer_ += (params_.advance_column * increment_strided); + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_DEVICE void clear_mask() { + mask_.clear(); + } + + ///< Efficiently enables all accesses guarded by mask + MCTLASS_DEVICE void enable_mask() { + mask_.enable(); + } + + ///< Sets the mask + MCTLASS_DEVICE void get_mask(Mask &mask) { + mask = mask_; + } + + ///< Sets the mask + MCTLASS_DEVICE void set_mask(Mask const &mask) { + mask_ = mask; + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load output tile from global memory in epilogue. +/// +/// Satisfies: ReadableTileIterator | InterleavedMaskedTileIterator | ForwardTileIterator +/// +template < + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_, ///< Element data type + int InterleavedN ///< Number of Interleaved N +> +class InterleavedConvPredicatedTileIterator { +public: + using ThreadMap = ThreadMap_; + + using Element = Element_; + + using Layout = layout::TensorNCxHWx; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = Tensor4DCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Iterations::kCount; + + /// Fragment object + using Fragment = Array; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + + struct Params { + + // + // Data members + // + + LongIndex stride_col; ///< stride in bytes between columns + LongIndex stride_row; ///< stride in bytes between rows + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Status initialize(typename Layout::Stride stride_) { + stride_col = stride_[1]; + stride_row = stride_[2]; + + return Status::kSuccess; + } + + MCTLASS_HOST_DEVICE + Params() { + initialize(mctlass::make_Coord(0, 0, 0)); + } + + MCTLASS_HOST_DEVICE + Params(Layout const &layout) { + + initialize(layout.stride()); + } + }; + + /// Mask object + struct Mask { + static int const kCount = + (ThreadMap::Iterations::kRow < 8) ? 8 : ThreadMap::Iterations::kRow; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + MCTLASS_HOST_DEVICE + Mask() { + enable(); + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_HOST_DEVICE void clear() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = false; + } + } + + ///< MCTLASS_HOST_DEVICE enables all accesses guarded by mask + MCTLASS_DEVICE void enable() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = true; + } + } + }; + +private: + + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + Params params_; + + /// Byte-level pointer + uint8_t *byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in columns + Index extent_col_; + + /// Extent of the matrix tile in rows + Index extent_row_; + + /// Extent of the matrix tile in pq + Index extent_pq_; + + /// A thread's starting row position (assuming steady-state predicates have + /// been computed) + Index thread_start_row_; + + /// A thread's starting column position (assuming steady-state predicates have + /// been computed) + Index thread_start_col_; + + /// Internal iteration counter + LongIndex iteration_row_; + LongIndex iteration_col_; + + uint32_t pq_mul_; + + uint32_t pq_shr_; + +private: + + // + // Methods + // + +public: + + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + InterleavedConvPredicatedTileIterator( + Params const & params, + Element *pointer, + TensorCoord extent, + int thread_idx, + MatrixCoord threadblock_offset + ): + params_(params) { + MatrixCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + extent_col_ = extent.c(); + extent_pq_ = extent.h() * extent.w(); + extent_row_ = extent.n() * extent_pq_; + + find_divisor(pq_mul_, pq_shr_, extent_pq_); + + thread_start_row_ = thread_offset.row(); + thread_start_col_ = thread_offset.column(); + + // Initialize predicates + MCTLASS_PRAGMA_UNROLL + for (int r = 0; r < ThreadMap::Iterations::kRow; ++r) { + mask_.predicates[r] = + ((thread_offset.row() + ThreadMap::Delta::kRow * r) < extent_row_); + } + + // Initialize pointer + byte_pointer_ = reinterpret_cast(pointer) + + ((thread_start_col_ / InterleavedN) * params_.stride_col + + (thread_start_col_ % InterleavedN)) * + sizeof_bits::value / 8; + + // Initialize internal state counter + iteration_row_ = iteration_col_ = 0; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + + int col_offset = iteration_col_ * ThreadMap::Delta::kColumn; + bool col_guard = ((thread_start_col_ + col_offset) < extent_col_); + bool guard = col_guard && mask_.predicates[iteration_row_]; + + int n, pq_rem; + + fast_divmod(n, pq_rem, + thread_start_row_ + iteration_row_ * ThreadMap::Delta::kRow, + extent_pq_, pq_mul_, pq_shr_); + + uint8_t *byte_pointer = + byte_pointer_ + (n * params_.stride_row + pq_rem * InterleavedN) * + sizeof_bits::value / 8; + AccessType *frag_ptr = reinterpret_cast(&frag); + AccessType const *memory_pointer = + reinterpret_cast(byte_pointer); + + mctlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + *frag_ptr, + (void *)memory_pointer, + guard); + } + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + if (__lane_id() >= 32) { + return; + } + int col_offset = iteration_col_ * ThreadMap::Delta::kColumn; + bool col_guard = ((thread_start_col_ + col_offset) < extent_col_); + bool guard = col_guard && mask_.predicates[iteration_row_]; + + int n, pq_rem; + + fast_divmod(n, pq_rem, + thread_start_row_ + iteration_row_ * ThreadMap::Delta::kRow, + extent_pq_, pq_mul_, pq_shr_); + + uint8_t *byte_pointer = + byte_pointer_ + (n * params_.stride_row + pq_rem * InterleavedN) * + sizeof_bits::value / 8; + AccessType const *frag_ptr = reinterpret_cast(&frag); + AccessType *memory_pointer = reinterpret_cast(byte_pointer); + + mctlass::arch::global_store( + *frag_ptr, (void *)memory_pointer, guard); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int iteration) { + iteration_row_ = iteration % ThreadMap::Iterations::kRow; + iteration_col_ = iteration / ThreadMap::Iterations::kRow; + } + + /// Advances to the next position to load or store + MCTLASS_HOST_DEVICE + InterleavedConvPredicatedTileIterator &operator++() { + + ++iteration_row_; + + if (iteration_row_ == ThreadMap::Iterations::kRow) { + + iteration_row_ = 0; + ++iteration_col_; + byte_pointer_ += params_.stride_col; + + if (iteration_col_ == ThreadMap::Iterations::kColumn) { + iteration_col_ = 0; + } + } + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_DEVICE void clear_mask() { + mask_.clear(); + } + + ///< Efficiently enables all accesses guarded by mask + MCTLASS_DEVICE void enable_mask() { + mask_.enable(); + } + + ///< Sets the mask + MCTLASS_DEVICE void get_mask(Mask &mask) { + mask = mask_; + } + + ///< Sets the mask + MCTLASS_DEVICE void set_mask(Mask const &mask) { + mask_ = mask; + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_affine.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_affine.h new file mode 100644 index 0000000..34e1c13 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_affine.h @@ -0,0 +1,616 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/epilogue/threadblock/output_tile_thread_map.h" +#include "mctlass/arch/arch.h" +#include "mctlass/arch/memory.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////// + +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load and store output tile from global memory in epilogue. +/// +/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator +/// +/// It provides a fast path for the case Rank = 2 which does not need div/rem to +/// calculate modes. + +template < + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_, ///< Element data type + int Rank +> +class PredicatedTileIteratorAffineRankN { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::AffineRankN; + using TensorRef = TensorRef; + using TensorView = TensorView; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = typename Layout::TensorCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0"); + static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0"); + static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0"); + static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0"); + static_assert( !(Layout::kRank % 2), + "Layout rank must be even. This assumes the first half of the modes correspond to the 'row' " + "and the second half of the modes correspond to the 'column'"); + + static bool const kBigEndian = false; + + /// Fragment object + using Fragment = Array< + Element, + ThreadMap::Iterations::kColumn * + ThreadMap::Iterations::kRow * + ThreadMap::Iterations::kGroup * + ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + + /// Parameters structure + struct Params { + + // + // Data members + // + + Layout layout; + + /// Stride in units of bytes along M modes + Coord stride_m; + + /// Stride in units of bytes along N modes + Coord stride_n; + + /// Fast divmod objects divided by tensor extents + FastDivmod divmod_m[(Layout::kRank == 2) ? 1 : (Layout::kRank/2 - 1)]; + + /// Fast divmod objects divided by tensor extents + FastDivmod divmod_n[(Layout::kRank == 2) ? 1 : (Layout::kRank/2 - 1)]; + + int64_t rank2_inc_col; + int64_t rank2_inc_row; + + // + // Methods + // + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(TensorCoord const &extent, Layout const &layout_): layout(layout_) { + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2; ++i) { + stride_m[i] = OffsetBytes(layout_.stride()[i]); + stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2]); + } + + if (kBigEndian) { + // "Big Endian" scheme + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2 - 1; ++i) { + divmod_m[i] = FastDivmod(extent[i + 1]); + divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2 + 1]); + } + } + else { + // "Little Endian" scheme + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2 - 1; ++i) { + divmod_m[i] = FastDivmod(extent[i]); + divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2]); + } + } + + #if 0 + // + // Debug print statements to verify extents and strides are passed correctly. + // + printf("PredicatedTileIteratorAffine::Params() entered\n"); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank; ++i) { + printf(" extent[%d]: %d\n", i, extent[i]); + } + for (int i = 0; i < Layout::kRank; ++i) { + printf(" stride[%d]: %ld\n", i, layout_.stride()[i]); + } + printf("PredicatedTileIteratorAffine::Params() returning\n"); + #endif + } + + MCTLASS_HOST_DEVICE + Params(Layout const &layout_): layout(layout_) { + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2; ++i) { + stride_m[i] = OffsetBytes(layout_.stride()[i]); + stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2]); + } + + rank2_inc_col = ThreadMap::Delta::kColumn * stride_n[0]; + rank2_inc_row = ThreadMap::Delta::kRow * stride_m[0]; + } + }; + + /// Mask object + struct Mask { + + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + MCTLASS_HOST_DEVICE + Mask() { + enable(); + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_HOST_DEVICE void clear() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = false; + } + } + + ///< MCTLASS_HOST_DEVICE enables all accesses guarded by mask + MCTLASS_DEVICE void enable() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = true; + } + } + }; + +private: + + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + Params params_; + + /// Byte-level pointer + uint8_t *byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Index extent_row_; + + /// Extent of the matrix tile in columns + Index extent_col_; + + /// A thread's starting row position (assuming steady-state predicates have been computed) + Index thread_start_row_; + + /// A thread's starting column position (assuming steady-state predicates have been computed) + Index thread_start_column_; + + /// Internal state counter + int state_[3]; + + /// Offsets in columns, cached for performance + int64_t offset_modes_n_[ThreadMap::Iterations::kColumn]; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + +private: + + // + // Methods + // + +public: + + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + PredicatedTileIteratorAffineRankN( + Params const & params, + Element *pointer, + MatrixCoord extent, + int thread_idx, + MatrixCoord threadblock_offset = MatrixCoord(), + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ): + params_(params) + { + + MatrixCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + extent_row_ = extent.row(); + extent_col_ = extent.column(); + + thread_start_row_ = thread_offset.row(); + thread_start_column_ = thread_offset.column(); + + if (Layout::kRank > 2) { + // Initialize predicates + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + + // + // Compute coordinate and decompose into N modes + // + + int coord_n = thread_start_column_ + c * ThreadMap::Delta::kColumn; + + mask_.predicates[c] = coord_n < extent.column(); + + Coord modes_n; + + int64_t offset_modes_n = 0; + + if (kBigEndian) { + modes_n = CoordinateDecomposition(coord_n, params_.divmod_n); + + offset_modes_n = dot(modes_n, params_.stride_n); + } + else { + modes_n = CoordinateDecompositionLittleEndian(coord_n, params_.divmod_n); + + offset_modes_n = dot(modes_n, params_.stride_n); + } + + offset_modes_n_[c] = offset_modes_n; + + } + + if (!pointer) { + mask_.clear(); + } + } + + // Initialize pointer + byte_pointer_ = reinterpret_cast(pointer); + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, int64_t byte_offset) { + uint8_t const *byte_pointer = byte_pointer_; + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + int row_begin = thread_start_row_ + group * ThreadMap::Delta::kGroup + cluster * ThreadMap::Delta::kCluster; + int64_t offset_modes_m = row_begin * params_.stride_m[0]; + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + // + // Compute coordinate and decompose into M modes + // + + int coord_m = row * ThreadMap::Delta::kRow + row_begin; + + Coord modes_m; + + if (Layout::kRank > 2) { + if (kBigEndian) { + modes_m = CoordinateDecomposition(coord_m, params_.divmod_m); + } else { + modes_m = CoordinateDecompositionLittleEndian(coord_m, params_.divmod_m); + } + + offset_modes_m = dot(modes_m, params_.stride_m); + } + + // + // Compute the offset due to modes M + // + + bool row_guard = (coord_m < extent_row_); + int64_t offset_modes_n = thread_start_column_ * params_.stride_n[0]; + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + // + // Compute coordinate and decompose into N modes + // + + if (Layout::kRank > 2) { + offset_modes_n = offset_modes_n_[column]; + } + + // + // Compute the pointer and access + // + bool guard; + + if (Layout::kRank > 2) { + guard = row_guard && mask_.predicates[column]; + } else { + guard = (coord_m < extent_row_) && + ((thread_start_column_ + ThreadMap::Delta::kColumn * column) < extent_col_); + } + + mctlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void *)(byte_pointer + offset_modes_m + offset_modes_n + byte_offset), + guard + ); + + if (Layout::kRank == 2) { + offset_modes_n += params_.rank2_inc_col; + } + } + + if (Layout::kRank == 2) { + offset_modes_m += params_.rank2_inc_row; + } + } + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + + load_with_byte_offset(frag, 0); + } + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) { + uint8_t *byte_pointer = byte_pointer_; + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + int row_begin = thread_start_row_ + group * ThreadMap::Delta::kGroup + cluster * ThreadMap::Delta::kCluster; + int64_t offset_modes_m = row_begin * params_.stride_m[0]; + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + // + // Compute coordinate and decompose into M modes + // + + int coord_m = row * ThreadMap::Delta::kRow + row_begin; + + Coord modes_m; + + if (Layout::kRank > 2) { + if (kBigEndian) { + modes_m = CoordinateDecomposition(coord_m, params_.divmod_m); + } else { + modes_m = CoordinateDecompositionLittleEndian(coord_m, params_.divmod_m); + } + + offset_modes_m = dot(modes_m, params_.stride_m); + } + + // + // Compute the offset due to modes M + // + + bool row_guard = (coord_m < extent_row_); + int64_t offset_modes_n = thread_start_column_ * params_.stride_n[0]; + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + // + // Compute coordinate and decompose into N modes + // + + if (Layout::kRank > 2) { + offset_modes_n = offset_modes_n_[column]; + } + + // + // Compute the pointer and access + // + bool guard; + if (Layout::kRank > 2) { + guard = row_guard && mask_.predicates[column]; + } else { + guard = (coord_m < extent_row_) && ((thread_start_column_ + ThreadMap::Delta::kColumn * column) < extent_col_); + } + + mctlass::arch::global_store( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void *)(byte_pointer + offset_modes_m + offset_modes_n + byte_offset), + guard); + + if (Layout::kRank == 2) { + offset_modes_n += params_.rank2_inc_col; + } + } + + if (Layout::kRank == 2) { + offset_modes_m += params_.rank2_inc_row; + } + } + } + } + } + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + if (threadIdx.x < kThreads) { + store_with_byte_offset(frag, 0); + } + } + + /// Advances to the next position to load or store + MCTLASS_HOST_DEVICE + PredicatedTileIteratorAffineRankN &operator++() { + + ++state_[0]; + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + + state_[0] = 0; + ++state_[1]; + + thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * + ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + + state_[1] = 0; + ++state_[2]; + + thread_start_row_ += ThreadMap::Count::kGroup * + ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { + state_[2] = 0; + } + } + } + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_DEVICE void clear_mask() { + mask_.clear(); + } + + ///< Efficiently enables all accesses guarded by mask + MCTLASS_DEVICE void enable_mask() { + mask_.enable(); + } + + ///< Sets the mask + MCTLASS_DEVICE void get_mask(Mask &mask) { + mask = mask_; + } + + ///< Sets the mask + MCTLASS_DEVICE void set_mask(Mask const &mask) { + mask_ = mask; + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_affine_layout_params.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_affine_layout_params.h new file mode 100644 index 0000000..546a914 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_affine_layout_params.h @@ -0,0 +1,156 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/fast_math.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int Rank +> +struct PredicatedTileIteratorAffineLayoutRankNParams { + using Layout = layout::AffineRankN; + using TensorCoord = typename Layout::TensorCoord; + + static bool const kBigEndian = false; + + // + // Data members + // + + Layout layout; + + /// Stride in units of bytes along M modes + Coord stride_m; + + /// Stride in units of bytes along N modes + Coord stride_n; + + /// Fast divmod objects divided by tensor extents + FastDivmod divmod_m[(Layout::kRank == 2) ? 1 : (Layout::kRank/2 - 1)]; + + /// Fast divmod objects divided by tensor extents + FastDivmod divmod_n[(Layout::kRank == 2) ? 1 : (Layout::kRank/2 - 1)]; + + int64_t rank2_inc_col; + int64_t rank2_inc_row; + + // + // Methods + // + MCTLASS_HOST_DEVICE + PredicatedTileIteratorAffineLayoutRankNParams() { } + + MCTLASS_HOST_DEVICE + PredicatedTileIteratorAffineLayoutRankNParams(TensorCoord const &extent, + Layout const &layout_, + int64_t element_sizeof_bits) + : layout(layout_) + { + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2; ++i) { + stride_m[i] = OffsetBytes(layout_.stride()[i], element_sizeof_bits); + stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2], element_sizeof_bits); + } + + if (kBigEndian) { + // "Big Endian" scheme + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2 - 1; ++i) { + divmod_m[i] = FastDivmod(extent[i + 1]); + divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2 + 1]); + } + } + else { + // "Little Endian" scheme + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2 - 1; ++i) { + divmod_m[i] = FastDivmod(extent[i]); + divmod_n[i] = FastDivmod(extent[i + Layout::kRank / 2]); + } + } + + #if 0 + // + // Debug print statements to verify extents and strides are passed correctly. + // + printf("PredicatedTileIteratorAffine::Params() entered\n"); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank; ++i) { + printf(" extent[%d]: %d\n", i, extent[i]); + } + for (int i = 0; i < Layout::kRank; ++i) { + printf(" stride[%d]: %ld\n", i, layout_.stride()[i]); + } + printf("PredicatedTileIteratorAffine::Params() returning\n"); + #endif + } + + MCTLASS_HOST_DEVICE + PredicatedTileIteratorAffineLayoutRankNParams(Layout const &layout_, + int32_t threadmap_delta_kColumn, + int32_t threadmap_delta_kRow, + int64_t element_sizeof_bits) + : layout(layout_) + { + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Layout::kRank / 2; ++i) { + stride_m[i] = OffsetBytes(layout_.stride()[i], element_sizeof_bits); + stride_n[i] = OffsetBytes(layout_.stride()[i + Layout::kRank / 2], element_sizeof_bits); + } + + rank2_inc_col = threadmap_delta_kColumn * stride_n[0]; + rank2_inc_row = threadmap_delta_kRow * stride_m[0]; + } +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_blas3.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_blas3.h new file mode 100644 index 0000000..afc88eb --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_blas3.h @@ -0,0 +1,635 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/epilogue/threadblock/output_tile_thread_map.h" +#include "mctlass/arch/arch.h" +#include "mctlass/arch/memory.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////// + +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load and store output tile from global memory in epilogue. +/// +/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator +/// +template < + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_, ///< Element data type + BlasMode BlasMode_ = BlasMode::kGemm ///< Tile Iterator for a Symmetric or Hermitian Kernel +> +class PredicatedTileIteratorBlas3 { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static BlasMode const kBlasMode = BlasMode_; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0"); + static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0"); + static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0"); + static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0"); + + /// Fragment object + using Fragment = Array< + Element, + ThreadMap::Iterations::kColumn * + ThreadMap::Iterations::kRow * + ThreadMap::Iterations::kGroup * + ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>; + + /// Memory access size + using AccessType = AlignedArray; + static_assert( AccessType::kElements == 1, "BLAS3 Epilogue must use AccessType::kElements as 1"); + + // + // Parameters struct + // + + /// Uses a non-template class + struct Params : PredicatedTileIteratorParams { + + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Layout const &layout): + PredicatedTileIteratorParams( + layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess, + make_OutputTileThreadMapDesc() + ) + { + + } + }; + + /// Mask object + struct Mask { + + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + MCTLASS_HOST_DEVICE + Mask() { + enable(); + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_HOST_DEVICE void clear() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = false; + } + } + + ///< MCTLASS_HOST_DEVICE enables all accesses guarded by mask + MCTLASS_DEVICE void enable() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = true; + } + } + }; + +private: + + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + PredicatedTileIteratorParams params_; + + /// Byte-level pointer + uint8_t *byte_pointer_; + + /// Fill Mode for a tile on diagonal of a symmetric kernel + mctlass::FillMode fill_mode; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Index extent_row_; + + /// A thread's starting row position (assuming steady-state predicates have been computed) + Index thread_start_row_; + + /// Internal state counter + int state_[3]; + + /// Starting address of the matrix + size_t matrix_start_addr; + + static_assert((kBlasMode == BlasMode::kSymmetric || kBlasMode == BlasMode::kHermitian), + "Unsupported blas3 mode."); + +private: + + // + // Methods + // + +public: + + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + PredicatedTileIteratorBlas3( + PredicatedTileIteratorParams const & params, + Element *pointer, + TensorCoord extent, + int thread_idx, + TensorCoord threadblock_offset + , mctlass::FillMode fill_mode + ): + params_(params), fill_mode(fill_mode) + { + + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + extent_row_ = extent.row(); + thread_start_row_ = thread_offset.row(); + + // Initialize predicates + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + + mask_.predicates[c] = ((thread_offset.column() + + ThreadMap::Delta::kColumn * c) < extent.column()); + } + + // Check Symmetric kernel modes (Lower and Upper - for diagonal CTAs, None for rest CTAs) + if ((kBlasMode == BlasMode::kSymmetric || kBlasMode == BlasMode::kHermitian) && + fill_mode == mctlass::FillMode::kInvalid) { + arch::device_breakpoint(); + } + + // Starting address of the matrix + matrix_start_addr = reinterpret_cast(pointer); + + // Initialize pointer + byte_pointer_ = reinterpret_cast(pointer) + + LongIndex(thread_offset.row()) * LongIndex(params_.stride) + + LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, int64_t byte_offset) { + + uint8_t *byte_pointer = byte_pointer_; + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + AccessType *memory_pointer = reinterpret_cast(byte_pointer + byte_offset); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + bool guard = row_guard && mask_.predicates[column]; + + mctlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + + column], + (void *)&memory_pointer[column * ThreadMap::Delta::kColumn / + kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + byte_pointer += params_.increment_row; + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Loads a fragment on the diagonal of a symmetric kernel to memory + MCTLASS_DEVICE + void load_symmetric_with_byte_offset(Fragment &frag, int64_t byte_offset) { + + uint8_t *byte_pointer = byte_pointer_; + AccessType *frag_ptr = reinterpret_cast(&frag); + + bool isLowerMode = (fill_mode == mctlass::FillMode::kLower) ? true : false; + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + AccessType *memory_pointer = reinterpret_cast(byte_pointer + byte_offset); + + // Offset of row from beginning of the matrix per thread + size_t row_start_offset = (size_t)memory_pointer - matrix_start_addr; + + // Absolute row index + int row_index = int(row_start_offset/params_.stride); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + bool guard = row_guard && mask_.predicates[column]; + + // Offset of column from beginning of row per thread + size_t col_start_offset = row_start_offset + + (column * ThreadMap::Delta::kColumn / kElementsPerAccess) * sizeof(AccessType); + + // Absolute column index + size_t col_index = (col_start_offset%params_.stride)/sizeof(AccessType); + guard = guard && ( (isLowerMode && row_index >= col_index) || + (!isLowerMode && row_index <= col_index) ); + + mctlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + + column], + (void *)&memory_pointer[column * ThreadMap::Delta::kColumn / + kElementsPerAccess], + guard); + + // The imaginary parts of the diagonal elements of a complex element are assumed and set to zero + if (guard && kBlasMode == BlasMode::kHermitian && mctlass::is_complex::value) { + Element *scalar_ptr = reinterpret_cast(frag_ptr); + + if (row_index == col_index) { + scalar_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column] = + real(scalar_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column]); + } + } + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + byte_pointer += params_.increment_row; + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + + if (fill_mode == mctlass::FillMode::kNone) { + load_with_byte_offset(frag, 0); + } + else { + load_symmetric_with_byte_offset(frag, 0); + } + } + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) { + uint8_t *byte_pointer = byte_pointer_; + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + AccessType *memory_pointer = reinterpret_cast(byte_pointer + byte_offset); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + bool guard = row_guard && mask_.predicates[column]; + + mctlass::arch::global_store( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void *)&memory_pointer[column * ThreadMap::Delta::kColumn / kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + byte_pointer += params_.increment_row; + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Stores a fragment on the diagonal of a symmetric kernel to memory + MCTLASS_DEVICE + void store_symmetric_with_byte_offset(Fragment const &frag, int64_t byte_offset) { + uint8_t *byte_pointer = byte_pointer_; + AccessType const *frag_ptr = reinterpret_cast(&frag); + + bool isLowerMode = (fill_mode == mctlass::FillMode::kLower) ? true : false; + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + AccessType *memory_pointer = reinterpret_cast(byte_pointer + byte_offset); + + // Offset of row from beginning of the matrix per thread + size_t row_start_offset = (size_t)memory_pointer - matrix_start_addr; + + // Absolute row index + int row_index = int(row_start_offset/params_.stride); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + bool guard = row_guard && mask_.predicates[column]; + + // Offset of column from beginning of row per thread + size_t col_start_offset = row_start_offset + + (column * ThreadMap::Delta::kColumn / kElementsPerAccess) * sizeof(AccessType); + + // Absolute column index + size_t col_index = (col_start_offset%params_.stride)/sizeof(AccessType); + + guard = guard && ( (isLowerMode && row_index >= col_index) || + (!isLowerMode && row_index <= col_index) ); + + // The imaginary parts of the diagonal elements of a complex element are assumed and set to zero + if (guard && kBlasMode == BlasMode::kHermitian && mctlass::is_complex::value) { + + AccessType *frag_ptr_modify = const_cast(frag_ptr); + Element *scalar_ptr = reinterpret_cast(frag_ptr_modify); + + if (row_index == col_index) { + scalar_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column] = + real(scalar_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column]); + } + } + + mctlass::arch::global_store( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + + column], + (void *)&memory_pointer[column * ThreadMap::Delta::kColumn / + kElementsPerAccess], + guard); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + byte_pointer += params_.increment_row; + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + + if (threadIdx.x < kThreads) { + if (fill_mode == mctlass::FillMode::kNone) { + store_with_byte_offset(frag, 0); + } + else { + store_symmetric_with_byte_offset(frag, 0); + } + } + + } + + /// Advances to the next position to load or store + MCTLASS_HOST_DEVICE + PredicatedTileIteratorBlas3 &operator++() { + + ++state_[0]; + byte_pointer_ += params_.advance_row; + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + + state_[0] = 0; + ++state_[1]; + byte_pointer_ += params_.advance_group; + + thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * + ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + + state_[1] = 0; + ++state_[2]; + byte_pointer_ += params_.advance_cluster; + + thread_start_row_ += ThreadMap::Count::kGroup * + ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { + state_[2] = 0; + byte_pointer_ += params_.advance_tile; + } + } + } + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_DEVICE void clear_mask() { + mask_.clear(); + } + + ///< Efficiently enables all accesses guarded by mask + MCTLASS_DEVICE void enable_mask() { + mask_.enable(); + } + + ///< Sets the mask + MCTLASS_DEVICE void get_mask(Mask &mask) { + mask = mask_; + } + + ///< Sets the mask + MCTLASS_DEVICE void set_mask(Mask const &mask) { + mask_ = mask; + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_direct_conv.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_direct_conv.h new file mode 100644 index 0000000..5c3f2da --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_direct_conv.h @@ -0,0 +1,445 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/permute.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/epilogue/threadblock/output_tile_thread_map.h" +#include "mctlass/arch/arch.h" +#include "mctlass/arch/memory.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_params.h" +#include "mctlass/conv/conv2d_problem_size.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////// + +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load and store output tile from global memory in epilogue. +/// +/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator +/// +template < + typename ThreadMap_, ///< Thread map (conept: PitchLinearThreadMap) + typename Element_, ///< Element data type + typename ThreadOutputShape_ = mctlass::conv::TensorNHWCShape<1, 1, 1, 1>, + typename ThreadBlockOutputShape_ = mctlass::conv::TensorNHWCShape<1, 1, 1, 1> +> +class PredicatedTileIteratorDirectConv { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + using ThreadOutputShape = ThreadOutputShape_; + using ThreadBlockOutputShape = ThreadBlockOutputShape_; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + + using ConvProblemSize = typename mctlass::conv::Conv2dProblemSize; + + /// Fragment object + using Fragment = Array; + + /// Memory access size + using AccessType = AlignedArray; + + static int const kLoadsPerAccess = AccessType::kElements / AccessType::kElements; + + using ThreadTileCount = MatrixShape< + ThreadBlockOutputShape::kH / ThreadOutputShape::kH, + ThreadBlockOutputShape::kW / ThreadOutputShape::kW + >; + + // + // Parameters struct + // + + /// Uses a non-template class + struct Params : PredicatedTileIteratorDirect2dConvParams { + using Base = PredicatedTileIteratorDirect2dConvParams; + + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Layout const &layout, mctlass::conv::Conv2dProblemSize const &problem_size): + PredicatedTileIteratorDirect2dConvParams( + layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess, + problem_size, + {ThreadBlockOutputShape::kH, ThreadBlockOutputShape::kW} + ) + { } + + MCTLASS_HOST_DEVICE + Params(Base const &base) : + Base(base) { } + }; + + /// Mask object + struct Mask { + + static int const kCount = ThreadMap::Iterations::kContiguous; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + MCTLASS_HOST_DEVICE + Mask() { + enable(); + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_HOST_DEVICE void clear() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = false; + } + } + + ///< MCTLASS_HOST_DEVICE enables all accesses guarded by mask + MCTLASS_DEVICE void enable() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = true; + } + } + }; + +private: + + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + PredicatedTileIteratorDirect2dConvParams params_; + + /// Byte-level pointer + uint8_t *byte_pointer_; + + /// + Element *pointer_; + + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Index extent_row_; + + /// Extent of the matrix tile in rows + Index extent_column_; + + /// A thread's starting row position (assuming steady-state predicates have been computed) + Index thread_start_row_; + + /// A thread's starting column + Index thread_start_column_; + + /// Initial thread output location + int thread_start_n_, thread_start_p_, thread_start_q_; + + /// Current threadblock tile index + int tile_index_; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(PredicatedTileIteratorDirect2dConvParams::stride) == 8, "Expected 64b strides"); + +private: + + // + // Methods + // + + + +public: + + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + PredicatedTileIteratorDirectConv( + PredicatedTileIteratorDirect2dConvParams const & params, + Element *pointer, + TensorCoord extent, + int thread_idx, + TensorCoord threadblock_offset = TensorCoord() + ): + params_(params), pointer_(pointer) + { + + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx); + + extent_row_ = extent.row(); + extent_column_ = extent.column(); + + // stride dim (PQ) + thread_start_row_ = thread_offset.column(); + // contiguous dim (Channels) + thread_start_column_ = threadblock_offset.column() + thread_offset.row(); + + tile_index_ = threadblock_offset.row(); + + set_tile_index(0); + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void set_tile_index(const int index) { + + int residual; + params_.pq_divmod(thread_start_n_, residual, tile_index_ + index); + params_.q_divmod(thread_start_p_, thread_start_q_, residual); + + // Compute the base output coord of ThreadBlock + thread_start_p_ *= ThreadBlockOutputShape::kH; + thread_start_q_ *= ThreadBlockOutputShape::kW; + + // Initialize predicates + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + mask_.predicates[c] = ((thread_start_column_ + + c * ThreadMap::Delta::kContiguous) < extent_column_); + } + + // Null pointer performs no accesses + if (!pointer_) { + mask_.clear(); + } + + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, int64_t byte_offset) const { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + int frag_base_idx = s * ThreadMap::Iterations::kContiguous + c; + + int current_row = thread_start_row_ + s * ThreadMap::Delta::kStrided; + int p = current_row / ThreadBlockOutputShape::kW; + int q = current_row % ThreadBlockOutputShape::kW; + + int current_p = thread_start_p_ + p; + int current_q = thread_start_q_ + q; + + bool row_guard = (current_p) < params_.P && (current_q) < params_.Q && + (thread_start_n_ < params_.N) && current_row < ThreadMap::Shape::kStrided; + + int output_row_offset = + thread_start_n_ * params_.stride_n + current_p * params_.stride_p + current_q; + + uint8_t *byte_pointer = + reinterpret_cast(pointer_) + + LongIndex(output_row_offset) * LongIndex(params_.stride) + + LongIndex(thread_start_column_ + c * ThreadMap::Delta::kContiguous) * + sizeof(AccessType) / kElementsPerAccess; + + AccessType *frag_ptr = reinterpret_cast(&frag); + + AccessType *memory_pointer = reinterpret_cast(byte_pointer + byte_offset); + + bool guard = row_guard && mask_.predicates[c]; + + mctlass::arch::global_load( + frag_ptr[frag_base_idx], (void *)&memory_pointer[0], guard); + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) const { + load_with_byte_offset(frag, 0); + } + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) const { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + int frag_base_idx = s * ThreadMap::Iterations::kContiguous + c; + + int current_row = thread_start_row_ + s * ThreadMap::Delta::kStrided; + int p = current_row / ThreadBlockOutputShape::kW; + int q = current_row % ThreadBlockOutputShape::kW; + + int current_p = thread_start_p_ + p; + int current_q = thread_start_q_ + q; + + bool row_guard = (current_p) < params_.P && (current_q) < params_.Q && + (thread_start_n_ < params_.N) && current_row < ThreadMap::Shape::kStrided; + + int output_row_offset = + thread_start_n_ * params_.stride_n + current_p * params_.stride_p + current_q; + + uint8_t *byte_pointer = + reinterpret_cast(pointer_) + + LongIndex(output_row_offset) * LongIndex(params_.stride) + + LongIndex(thread_start_column_ + c * ThreadMap::Delta::kContiguous) * + sizeof(AccessType) / kElementsPerAccess; + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + AccessType *memory_pointer = reinterpret_cast(byte_pointer + byte_offset); + + bool guard = row_guard && mask_.predicates[c]; + + mctlass::arch::global_store( + frag_ptr[frag_base_idx], (void *)&memory_pointer[0], guard); + } + } + } + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) const { + + store_with_byte_offset(frag, 0); + } + + MCTLASS_DEVICE + MatrixCoord thread_start() const { + return MatrixCoord(thread_start_row_, thread_start_column_); + } + + /// Need to get the thread start row from the tile iterator + MCTLASS_DEVICE + int32_t thread_start_row() const { + return thread_start_row_; + } + + /// Need to get the thread start row from the tile iterator + MCTLASS_DEVICE + int32_t thread_start_column() const { + return thread_start_column_; + } + + /// Extent of the matrix in rows + MCTLASS_DEVICE + Index extent_row() const { + return extent_row_; + } + + /// Extent of the matrix in columns + MCTLASS_DEVICE + Index extent_column() const { + return extent_column_; + } + + /// Advances to the next position to load or store + MCTLASS_HOST_DEVICE + PredicatedTileIteratorDirectConv &operator++() { + // do nothing + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_DEVICE void clear_mask() { + mask_.clear(); + } + + ///< Efficiently enables all accesses guarded by mask + MCTLASS_DEVICE void enable_mask() { + mask_.enable(); + } + + ///< Sets the mask + MCTLASS_DEVICE void get_mask(Mask &mask) const { + mask = mask_; + } + + ///< Sets the mask + MCTLASS_DEVICE void set_mask(Mask const &mask) { + mask_ = mask; + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_params.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_params.h new file mode 100644 index 0000000..518b51f --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_params.h @@ -0,0 +1,475 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/matrix.h" + +#include "mctlass/conv/conv2d_problem_size.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +struct OutputTileShapeDesc { + + int column; + int row; + int group; + int cluster; + int tile; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + OutputTileShapeDesc(): column(0), row(0), group(0), cluster(0), tile(0) { } + + /// Ctor + MCTLASS_HOST_DEVICE + OutputTileShapeDesc( + int column_, + int row_, + int group_, + int cluster_, + int tile_ + ): + column(column_), + row(row_), + group(group_), + cluster(cluster_), + tile(tile_) { } + + /// Total number of points in the 5D space + MCTLASS_HOST_DEVICE + int count() const { + return column * row * group * cluster * tile; + } + + #if 0 + MCTLASS_HOST_DEVICE + void print() const { + printf("{%d, %d, %d, %d, %d}", column, row, group, cluster, tile); + } + #endif +}; + +/// Helper template to construct an OutputTileShapeDesc from a OutputTileShape template. +template +MCTLASS_HOST_DEVICE +OutputTileShapeDesc make_OutputTileShapeDesc() { + return OutputTileShapeDesc( + Shape::kColumn, + Shape::kRow, + Shape::kGroup, + Shape::kCluster, + Shape::kTile + ); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Thread map description +struct OutputTileThreadMapDesc { + + int threads; + int elements_per_access; + OutputTileShapeDesc shape; + OutputTileShapeDesc iterations; + OutputTileShapeDesc delta; + OutputTileShapeDesc count; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + OutputTileThreadMapDesc() { } + + MCTLASS_HOST_DEVICE + OutputTileThreadMapDesc( + int threads_, + int elements_per_access_, + OutputTileShapeDesc shape_, + OutputTileShapeDesc iterations_, + OutputTileShapeDesc delta_, + OutputTileShapeDesc count_ + ): + threads(threads_), + elements_per_access(elements_per_access_), + shape(shape_), + iterations(iterations_), + delta(delta_), + count(count_) + { + + } +}; + +/// Helper template to construct an OutputTileShapeDesc from a OutputTileThreadMap template. +template +MCTLASS_HOST_DEVICE +OutputTileThreadMapDesc make_OutputTileThreadMapDesc() { + return OutputTileThreadMapDesc( + ThreadMap::kThreads, + ThreadMap::kElementsPerAccess, + make_OutputTileShapeDesc(), + make_OutputTileShapeDesc(), + make_OutputTileShapeDesc(), + make_OutputTileShapeDesc() + ); +} +/////////////////////////////////////////////////////////////////////////////// + +// +// Parameters struct for PredicatedTileIterator +// + +struct PredicatedTileIteratorParams { + + using Index = int32_t; + using LongIndex = int64_t; + + // + // Data members + // + + LongIndex stride; ///< stride in bytes between rows + + LongIndex increment_row; ///< increment quantity (in bytes) to advance when moving between rows + LongIndex increment_group; ///< increment quantity (in bytes) to advance when moving to the next group + LongIndex increment_cluster; ///< increment quantity (in bytes) to advance when moving to the next cluster + + LongIndex advance_row; ///< amount to add to move to the next 'row' position + LongIndex advance_group; ///< amount to add to move to the next 'group' position + LongIndex advance_cluster; ///< amount to add to move to the next 'cluster' position + LongIndex advance_tile; ///< amount to add to move to the next 'tile' + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Status initialize(LongIndex stride_, OutputTileThreadMapDesc thread_map) { + + stride = stride_; + + increment_row = stride * thread_map.delta.row; + + increment_group = stride * thread_map.delta.group + - stride * thread_map.delta.row * (thread_map.iterations.row - 1); + + increment_cluster = stride * thread_map.delta.cluster + - stride * thread_map.delta.group * (thread_map.iterations.group - 1) + - stride * thread_map.delta.row * (thread_map.iterations.row - 1); + + advance_row = stride * thread_map.shape.row; + + advance_group = + stride * + (thread_map.shape.group - 1) * thread_map.shape.row * thread_map.count.row; + + advance_cluster = + stride * + thread_map.count.group * + thread_map.shape.group * + thread_map.count.row * + thread_map.shape.row; + + advance_tile = + stride * + thread_map.shape.group * + thread_map.shape.row * + thread_map.shape.cluster * + thread_map.shape.tile; + + return Status::kSuccess; + } + + MCTLASS_HOST_DEVICE + Status initialize(Index stride_, OutputTileThreadMapDesc thread_map) { + return initialize(LongIndex(stride_), thread_map); + } + + MCTLASS_HOST_DEVICE + PredicatedTileIteratorParams() { + initialize(LongIndex(0), OutputTileThreadMapDesc()); + } + + MCTLASS_HOST_DEVICE + PredicatedTileIteratorParams(Index stride, OutputTileThreadMapDesc thread_map) { + initialize(stride, thread_map); + } + + MCTLASS_HOST_DEVICE + PredicatedTileIteratorParams(LongIndex stride, OutputTileThreadMapDesc thread_map) { + initialize(stride, thread_map); + } +}; + + + +/////////////////////////////////////////////////////////////////////////////// + +// +// Parameters struct for PredicatedTileIteratorDirect2dConv +// + +struct PredicatedTileIteratorDirect2dConvParams{ + using Index = int32_t; + using LongIndex = int64_t; + + // + // Data members + // + FastDivmod pq_divmod; + FastDivmod q_divmod; + + LongIndex stride; + LongIndex stride_n; + LongIndex stride_p; + + int N; + int P; + int Q; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Status initialize(LongIndex stride_, + mctlass::conv::Conv2dProblemSize const &problem_size, + MatrixCoord threadblock_output_shape) { + stride = stride_; // The stride per row of output tensor (bytes) + stride_n = problem_size.P * problem_size.Q; + stride_p = problem_size.Q ; + + N = problem_size.N; + P = problem_size.P; + Q = problem_size.Q; + + // Fastdivmod for output O, P, Q + if(threadblock_output_shape.row() != 0 && threadblock_output_shape.column() !=0 ){ + int tiles_p = + (problem_size.P + (threadblock_output_shape.row() - 1)) / (threadblock_output_shape.row()); + int tiles_q = (problem_size.Q + (threadblock_output_shape.column() - 1)) / + (threadblock_output_shape.column()); + + pq_divmod = FastDivmod(tiles_p * tiles_q); + q_divmod = FastDivmod(tiles_q); + } + + return Status::kSuccess; + } + + MCTLASS_HOST_DEVICE + Status initialize( + Index stride_, + mctlass::conv::Conv2dProblemSize const &problem_size = mctlass::conv::Conv2dProblemSize(), + MatrixCoord threadblock_output_shape = MatrixCoord()) { + return initialize(LongIndex(stride_), problem_size, threadblock_output_shape); + } + + MCTLASS_HOST_DEVICE + PredicatedTileIteratorDirect2dConvParams() { initialize(LongIndex(0)); } + + MCTLASS_HOST_DEVICE + PredicatedTileIteratorDirect2dConvParams(Index stride, + mctlass::conv::Conv2dProblemSize const &problem_size, + MatrixCoord threadblock_output_shape) { + initialize(stride, problem_size, threadblock_output_shape); + } + + MCTLASS_HOST_DEVICE + PredicatedTileIteratorDirect2dConvParams(LongIndex stride, + mctlass::conv::Conv2dProblemSize const &problem_size, + MatrixCoord threadblock_output_shape) { + initialize(stride, problem_size, threadblock_output_shape); + } +}; + +/////////////////////////////////////////////////////////////////////////////// +// InterleavedPredicatedTileIterator +/////////////////////////////////////////////////////////////////////////////// + + +/// Predicated tile access iterator descriptor object containing template dependent state +struct InterleavedPredicatedTileIteratorDesc { + + int element_size_bits; + int elements_per_access; + int threadmap_warp_size; + layout::PitchLinearCoord threadmap_iterations; + layout::PitchLinearCoord threadmap_delta; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + InterleavedPredicatedTileIteratorDesc() { } + + MCTLASS_HOST_DEVICE + InterleavedPredicatedTileIteratorDesc( + int element_size_bits_, + int elements_per_access_, + int threadmap_warp_size_, + layout::PitchLinearCoord threadmap_iterations_, + layout::PitchLinearCoord threadmap_delta_ + ): + element_size_bits(element_size_bits_), + elements_per_access(elements_per_access_), + threadmap_warp_size(threadmap_warp_size_), + threadmap_iterations(threadmap_iterations_), + threadmap_delta(threadmap_delta_) { } +}; + +// +// Parameters struct InterleavedPredicatedTileIterator +// + +struct InterleavedPredicatedTileIteratorParams { + + using Index = int32_t; + using LongIndex = int64_t; + + // + // Data members + // + + LongIndex stride; ///< stride in bytes between rows + LongIndex advance_row; ///< amount to add to move to the next 'row' position + LongIndex advance_column; ///< amount to add to move to the next 'column' position + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Status initialize(LongIndex stride_, InterleavedPredicatedTileIteratorDesc desc) { + + stride = stride_; + + advance_row = desc.threadmap_delta.contiguous() * desc.element_size_bits / 8; + + advance_column = stride_ - desc.threadmap_iterations.contiguous() * + desc.elements_per_access * + desc.element_size_bits * + desc.threadmap_warp_size / 16; + + return Status::kSuccess; + } + + MCTLASS_HOST_DEVICE + InterleavedPredicatedTileIteratorParams() { + initialize(LongIndex(0), InterleavedPredicatedTileIteratorDesc()); + } + + MCTLASS_HOST_DEVICE + InterleavedPredicatedTileIteratorParams(Index stride, InterleavedPredicatedTileIteratorDesc desc) { + initialize(stride, desc); + } + + MCTLASS_HOST_DEVICE + InterleavedPredicatedTileIteratorParams(LongIndex stride, InterleavedPredicatedTileIteratorDesc desc) { + initialize(stride, desc); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Helper template to construct an OutputTileShapeDesc from a OutputTileThreadMap template. +template +MCTLASS_HOST_DEVICE +InterleavedPredicatedTileIteratorDesc make_InterleavedPredicatedTileIteratorDesc() { + return InterleavedPredicatedTileIteratorDesc( + sizeof_bits::value, + ThreadMap::kElementsPerAccess, + ThreadMap::kWarpSize, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided} + ); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Helper template to construct an MakePredicatedTileIteratorDesc from a template +// dependent state +template + struct MakePredicatedTileIteratorDesc; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for layout::RowMajor output data. +template +struct MakePredicatedTileIteratorDesc < + Element, layout::RowMajor, ThreadMap> { + + MCTLASS_HOST_DEVICE + OutputTileThreadMapDesc operator()() { + + return make_OutputTileThreadMapDesc(); + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for layout::ColumnMajorInterleaved output data. +template +struct MakePredicatedTileIteratorDesc < + Element, layout::ColumnMajorInterleaved, ThreadMap> { + + MCTLASS_HOST_DEVICE + InterleavedPredicatedTileIteratorDesc operator()() { + + return make_InterleavedPredicatedTileIteratorDesc(); + } +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_predicates.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_predicates.h new file mode 100644 index 0000000..8d87958 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_predicates.h @@ -0,0 +1,309 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief PredicatedTileIteratorPredicates. + + PredicatedTileIteratorPredicates enables both upper and lower bounds for predicates. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/epilogue/threadblock/output_tile_thread_map.h" +#include "mctlass/arch/arch.h" +#include "mctlass/arch/memory.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////// + +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator predicates used to bound computations in epilogue. +/// +/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator +/// +template < + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_ ///< Element data type +> +class PredicatedTileIteratorPredicates { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0"); + static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0"); + static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0"); + static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0"); + + /// Fragment object + using Fragment = Array< + Element, + ThreadMap::Iterations::kColumn * + ThreadMap::Iterations::kRow * + ThreadMap::Iterations::kGroup * + ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + + /// Uses a non-template class + struct Params : PredicatedTileIteratorParams { + + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Layout const &layout): + PredicatedTileIteratorParams( + layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess, + make_OutputTileThreadMapDesc() + ) + { + + } + }; + + /// Mask object + struct Mask { + + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + MCTLASS_HOST_DEVICE + Mask() { + enable(); + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_HOST_DEVICE void clear() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = false; + } + } + + ///< MCTLASS_HOST_DEVICE enables all accesses guarded by mask + MCTLASS_DEVICE void enable() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = true; + } + } + }; + +private: + + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + PredicatedTileIteratorParams params_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Index lower_extent_row_; + Index upper_extent_row_; + + /// A thread's starting row position (assuming steady-state predicates have been computed) + Index thread_start_row_; + + /// Internal state counter + int state_[3]; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(lower_extent_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(upper_extent_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides"); + +private: + + // + // Methods + // + +public: + + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + PredicatedTileIteratorPredicates( + PredicatedTileIteratorParams const & params, + TensorCoord lower_extent, + TensorCoord upper_extent, + int thread_idx, + TensorCoord threadblock_offset = TensorCoord() + ): + params_(params) + { + + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + lower_extent_row_ = lower_extent.row(); + upper_extent_row_ = upper_extent.row(); + thread_start_row_ = thread_offset.row(); + + // Initialize predicates + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + + mask_.predicates[c] = ((thread_offset.column() + + ThreadMap::Delta::kColumn * c) < upper_extent.column()) && + ((thread_offset.column() + ThreadMap::Delta::kColumn * c) >= lower_extent.column()); + } + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + } + + /// Advances to the next position to load or store + MCTLASS_HOST_DEVICE + PredicatedTileIteratorPredicates &operator++() { + + ++state_[0]; + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + + state_[0] = 0; + ++state_[1]; + + thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * + ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + + state_[1] = 0; + ++state_[2]; + + thread_start_row_ += ThreadMap::Count::kGroup * + ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { + state_[2] = 0; + } + } + } + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_DEVICE void clear_mask() { + mask_.clear(); + } + + ///< Efficiently enables all accesses guarded by mask + MCTLASS_DEVICE void enable_mask() { + mask_.enable(); + } + + ///< Gets the mask + MCTLASS_DEVICE void get_mask(Mask &mask) { + mask = mask_; + } + + ///< Sets the mask + MCTLASS_DEVICE void set_mask(Mask const &mask) { + mask_ = mask; + } + + ///< Gets lower_extent_row_ + MCTLASS_DEVICE Index get_lower_extent_row() { + return lower_extent_row_; + } + + ///< Gets upper_extent_row_ + MCTLASS_DEVICE Index get_upper_extent_row() { + return upper_extent_row_; + } + + ///< Gets thread_start_row_ + MCTLASS_DEVICE Index get_thread_start_row() { + return thread_start_row_; + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_row_broadcast.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_row_broadcast.h new file mode 100644 index 0000000..58b286b --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_row_broadcast.h @@ -0,0 +1,519 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/permute.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/epilogue/threadblock/output_tile_thread_map.h" +#include "mctlass/arch/arch.h" +#include "mctlass/arch/memory.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////// + +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load and store output tile from global memory in epilogue. +/// +/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator +/// +template < + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_, ///< Element data type + bool ScatterD = false, ///< Scatter D operand or not + typename PermuteDLayout = layout::NoPermute, ///< Permute D operand or not + bool UseCUDAStore = false +> +class PredicatedTileIteratorRowBroadcast { + static_assert(!ScatterD); + static_assert(std::is_same::value); + +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0"); + static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0"); + static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0"); + static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0"); + + /// Fragment object + using Fragment = Array< + Element, + ThreadMap::Iterations::kColumn * + ThreadMap::Iterations::kRow * + ThreadMap::Iterations::kGroup * + ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + + /// Uses a non-template class + struct Params : PredicatedTileIteratorParams { + using Base = PredicatedTileIteratorParams; + + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Layout const &layout): + PredicatedTileIteratorParams( + layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess, + make_OutputTileThreadMapDesc() + ) + { } + + MCTLASS_HOST_DEVICE + Params(Base const &base) : + Base(base) { } + }; + + /// Mask object + struct Mask { + + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + MCTLASS_HOST_DEVICE + Mask() { + enable(); + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_HOST_DEVICE void clear() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = false; + } + } + + ///< MCTLASS_HOST_DEVICE enables all accesses guarded by mask + MCTLASS_DEVICE void enable() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = true; + } + } + }; + +private: + + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + PredicatedTileIteratorParams params_; + + /// Byte-level pointer. + uint8_t *byte_pointer_; + + /// Byte-level pointer for store(). + uint8_t *store_byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Index extent_row_; + + /// Extent of the matrix tile in rows + Index extent_column_; + + /// A thread's starting row position (assuming steady-state predicates have been computed) + Index thread_start_row_; + + /// A thread's starting column + Index thread_start_column_; + + /// Internal state counter + int state_[3]; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides"); + +private: + + // + // Methods + // + +public: + + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + PredicatedTileIteratorRowBroadcast( + PredicatedTileIteratorParams const & params, + Element *pointer, + TensorCoord extent, + int thread_idx, + TensorCoord threadblock_offset = TensorCoord(), + int const *indices = nullptr + ): + params_(params) + { + + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + extent_row_ = extent.row(); + extent_column_ = extent.column(); + + thread_start_row_ = thread_offset.row(); + thread_start_column_ = thread_offset.column(); + + // Initialize predicates + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + + mask_.predicates[c] = ((thread_offset.column() + + ThreadMap::Delta::kColumn * c) < extent.column()); + } + + // Null pointer performs no accesses + if (!pointer) { + mask_.clear(); + } + + // Initialize byte_pointer_ + byte_pointer_ = reinterpret_cast(pointer) + + LongIndex(thread_offset.row()) * LongIndex(params_.stride) + + LongIndex(thread_offset.column()) * sizeof(AccessType) / kElementsPerAccess; + + // store_byte_pointer_ is set to be the same with byte_pointer_ + store_byte_pointer_ = byte_pointer_; + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + + byte_pointer_ = reinterpret_cast(pointer) + + LongIndex(thread_offset.row()) * LongIndex(params_.stride); + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + store_byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, int64_t byte_offset) const { + uint8_t *byte_pointer = byte_pointer_; + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + bool guard = row_guard && mask_.predicates[column]; + + /* + mctlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + + column], + (void *)&memory_pointer[column * ThreadMap::Delta::kColumn / + kElementsPerAccess], + guard); + */ + if (guard) { + Element *bias = reinterpret_cast(byte_pointer + byte_offset); + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column].fill(*bias); + } + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + byte_pointer += params_.increment_row; + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) const { + + load_with_byte_offset(frag, 0); + } + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) const { + uint8_t *byte_pointer = store_byte_pointer_; + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + bool row_guard = ((row_offset + thread_start_row_) < extent_row_); + + AccessType *memory_pointer = reinterpret_cast(byte_pointer + byte_offset); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + bool guard = row_guard && mask_.predicates[column]; + + if (UseCUDAStore) { + if (guard) { + memory_pointer[0] = + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column]; + } + } else { + mctlass::arch::global_store( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void *)&memory_pointer[0], + guard); + } + + memory_pointer += (ThreadMap::Delta::kColumn / kElementsPerAccess); + } + + if (row + 1 < ThreadMap::Iterations::kRow) { + byte_pointer += params_.increment_row; + } + } + + if (group + 1 < ThreadMap::Iterations::kGroup) { + byte_pointer += params_.increment_group; + } + } + + if (cluster + 1 < ThreadMap::Iterations::kCluster) { + byte_pointer += params_.increment_cluster; + } + } + } + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) const { + + store_with_byte_offset(frag, 0); + } + + MCTLASS_DEVICE + MatrixCoord thread_start() const { + return MatrixCoord(thread_start_row_, thread_start_column_); + } + + /// Need to get the thread start row from the tile iterator + MCTLASS_DEVICE + int32_t thread_start_row() const { + return thread_start_row_; + } + + /// Need to get the thread start row from the tile iterator + MCTLASS_DEVICE + int32_t thread_start_column() const { + return thread_start_column_; + } + + /// Extent of the matrix in rows + MCTLASS_DEVICE + Index extent_row() const { + return extent_row_; + } + + /// Extent of the matrix in columns + MCTLASS_DEVICE + Index extent_column() const { + return extent_column_; + } + + /// Advances to the next position to load or store + MCTLASS_HOST_DEVICE + PredicatedTileIteratorRowBroadcast &operator++() { + + ++state_[0]; + + store_byte_pointer_ += params_.advance_row; + + byte_pointer_ += params_.advance_row; + + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + + state_[0] = 0; + ++state_[1]; + byte_pointer_ += params_.advance_group; + store_byte_pointer_ += params_.advance_group; + + thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * + ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + + state_[1] = 0; + ++state_[2]; + byte_pointer_ += params_.advance_cluster; + store_byte_pointer_ += params_.advance_cluster; + + thread_start_row_ += ThreadMap::Count::kGroup * + ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { + state_[2] = 0; + byte_pointer_ += params_.advance_tile; + store_byte_pointer_ += params_.advance_tile; + + thread_start_row_ += ThreadMap::Shape::kGroup * ThreadMap::Shape::kRow + * ThreadMap::Shape::kCluster * ThreadMap::Shape::kTile; + } + } + } + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_DEVICE void clear_mask() { + mask_.clear(); + } + + ///< Efficiently enables all accesses guarded by mask + MCTLASS_DEVICE void enable_mask() { + mask_.enable(); + } + + ///< Sets the mask + MCTLASS_DEVICE void get_mask(Mask &mask) const { + mask = mask_; + } + + ///< Sets the mask + MCTLASS_DEVICE void set_mask(Mask const &mask) { + mask_ = mask; + } +}; + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_strided_dgrad.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_strided_dgrad.h new file mode 100644 index 0000000..988f400 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/predicated_tile_iterator_strided_dgrad.h @@ -0,0 +1,480 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/epilogue/threadblock/output_tile_thread_map.h" +#include "mctlass/arch/arch.h" +#include "mctlass/arch/memory.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////// + +namespace epilogue { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load and store output tile from global memory in epilogue. +/// +/// Satisfies: ReadableTileIterator | PredicatedTileIterator | ForwardTileIterator +/// +template < + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_ ///< Element data type +> +class PredicatedTileIteratorStridedDgrad { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + static int const kThreads = ThreadMap::kThreads; + static int const kIterations = ThreadMap::Count::kTile; + + static_assert( ThreadMap::Iterations::kRow > 0,"ThreadMap::Iterations::kRow must be > 0"); + static_assert( ThreadMap::Iterations::kGroup > 0,"ThreadMap::Iterations::kGroup must be > 0"); + static_assert( ThreadMap::Iterations::kCluster > 0,"ThreadMap::Iterations::kCluster must be > 0"); + static_assert( ThreadMap::Iterations::kColumn > 0,"ThreadMap::Iterations::kColumn must be > 0"); + + /// Fragment object + using Fragment = Array< + Element, + ThreadMap::Iterations::kColumn * + ThreadMap::Iterations::kRow * + ThreadMap::Iterations::kGroup * + ThreadMap::Iterations::kCluster * ThreadMap::kElementsPerAccess>; + + /// Memory access size + using AccessType = AlignedArray; + + // + // Parameters struct + // + + /// Uses a non-template class + struct Params : PredicatedTileIteratorParams { + + /// Convolution problem size + mctlass::conv::Conv2dProblemSize problem_size; + int tiled_rows_per_filter; + + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Layout const &layout, mctlass::conv::Conv2dProblemSize problem_size_, int threadblock_row): + problem_size(problem_size_), + PredicatedTileIteratorParams( + layout.stride(0) * int(sizeof(AccessType)) / kElementsPerAccess, + make_OutputTileThreadMapDesc() + ) + { + + int tile_m_per_filter = strided_dgrad_tile_m_per_filter(problem_size, threadblock_row); + + tiled_rows_per_filter = tile_m_per_filter * threadblock_row; + } + }; + + /// Mask object + struct Mask { + + static int const kCount = ThreadMap::Iterations::kColumn; + + /// Predicate state + bool predicates[kCount]; + + // + // Mask + // + MCTLASS_HOST_DEVICE + Mask() { + enable(); + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_HOST_DEVICE void clear() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = false; + } + } + + ///< MCTLASS_HOST_DEVICE enables all accesses guarded by mask + MCTLASS_DEVICE void enable() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kCount; ++i) { + predicates[i] = true; + } + } + }; + +private: + + // + // Data members + // + + /// Parameters structure containing reference and precomputed state. + Params params_; + + /// Byte-level pointer + uint8_t *byte_pointer_; + + /// Array of boolean values to contain steady-state predicates + Mask mask_; + + /// Extent of the matrix tile in rows + Index extent_row_; + + /// Starting Dx h and w dimenstion for strided dgrad mapping + int start_h_, start_w_; + + /// Effective Dy P and Q dimenstions for strided dgrad mapping + int p_, q_; + + /// A thread's starting row position (assuming steady-state predicates have been computed) + Index thread_start_row_; + + /// A thread's starting column position (assuming steady-state predicates have been computed) + Index thread_start_column_; + + /// Internal state counter + int state_[3]; + + // + // Static asserts about internal strides + // + + static_assert(sizeof(extent_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(thread_start_row_) == 4, "Expected 32b extents"); + static_assert(sizeof(PredicatedTileIteratorParams::stride) == 8, "Expected 64b strides"); + +private: + + // + // Methods + // + +public: + + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + PredicatedTileIteratorStridedDgrad( + Params const & params, + Element *pointer, + TensorCoord extent, + int thread_idx, + FastDivmod const &stride_h_divmod, FastDivmod const &stride_w_divmod, + int start_r, int start_s, + TensorCoord threadblock_offset = TensorCoord() + ): + params_(params) + { + + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx) + threadblock_offset; + + int r = start_r; + int s = start_s; + + if (params_.problem_size.mode == mctlass::conv::Mode::kConvolution) { + r = (params_.problem_size.R - 1 - r); + s = (params_.problem_size.S - 1 - s); + } + + // compute starting coordinates in Dx start_h_ and start_w_ + strided_dgrad_starting_coords( + params_.problem_size, + stride_h_divmod, stride_w_divmod, + r, s, + start_h_, start_w_); + + p_ = (params_.problem_size.H - start_h_ + params_.problem_size.stride_h - 1) / params_.problem_size.stride_h; + q_ = (params_.problem_size.W - start_w_ + params_.problem_size.stride_w - 1) / params_.problem_size.stride_w; + + extent_row_ = extent.row(); + thread_start_row_ = thread_offset.row(); + thread_start_column_ = thread_offset.column(); + + // Initialize predicates + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kColumn; ++c) { + + mask_.predicates[c] = ((thread_offset.column() + + ThreadMap::Delta::kColumn * c) < extent.column()); + } + + // Null pointer performs no accesses + if (!pointer) { + mask_.clear(); + } + + // Initialize pointer + byte_pointer_ = reinterpret_cast(pointer); + + // Initialize internal state counter + state_[0] = state_[1] = state_[2] = 0; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, int64_t byte_offset) { + + uint8_t *byte_pointer = byte_pointer_; + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + // remapping rows to find the mapped_row_offset + int npq_offset = (row_offset + thread_start_row_) % params_.tiled_rows_per_filter; + + // (STEP 4.a) [order NHW rows to be loaded and stored in output Dx NHWxC layout] + int n = npq_offset / (p_ * q_); + int residual = npq_offset % (p_ * q_); + int p = residual / q_; + int q = residual % q_; + + int mapped_row_offset = n * (params_.problem_size.H * params_.problem_size.W) + + (start_h_ + p * params_.problem_size.stride_h) * params_.problem_size.W + + (start_w_ + q * params_.problem_size.stride_w); + bool row_guard = mapped_row_offset < extent_row_; + + int64_t row_byte_offset = mapped_row_offset * params_.stride; + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + int64_t column_byte_offset = (thread_start_column_ + column * ThreadMap::Delta::kColumn) * (sizeof_bits::value / 8); + + bool guard = row_guard && mask_.predicates[column]; + + mctlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + + column], + (void *)(byte_pointer + row_byte_offset + column_byte_offset + byte_offset), + guard); + } + } + } + } + } + + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + + load_with_byte_offset(frag, 0); + } + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, int64_t byte_offset) { + uint8_t *byte_pointer = byte_pointer_; + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + int row_offset = row * ThreadMap::Delta::kRow + + group * ThreadMap::Delta::kGroup + + cluster * ThreadMap::Delta::kCluster; + + // remapping rows to find the mapped_row_offset + int npq_offset = (row_offset + thread_start_row_) % params_.tiled_rows_per_filter; + + // (STEP 4.a) [order NHW rows to be loaded and stored in output Dx NHWxC layout] + int n = npq_offset / (p_ * q_); + int residual = npq_offset % (p_ * q_); + int p = residual / q_; + int q = residual % q_; + + int mapped_row_offset = n * (params_.problem_size.H * params_.problem_size.W) + + (start_h_ + p * params_.problem_size.stride_h) * params_.problem_size.W + + (start_w_ + q * params_.problem_size.stride_w); + bool row_guard = mapped_row_offset < extent_row_; + + int64_t row_byte_offset = mapped_row_offset * params_.stride; + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + int64_t column_byte_offset = (thread_start_column_ + column * ThreadMap::Delta::kColumn) * (sizeof_bits::value / 8); + + bool guard = row_guard && mask_.predicates[column]; + + mctlass::arch::global_store( + frag_ptr[frag_row_idx * ThreadMap::Iterations::kColumn + column], + (void *)(byte_pointer + row_byte_offset + column_byte_offset + byte_offset), + guard); + } + } + } + } + } + + + /// Stores a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + if (threadIdx.x < kThreads) { + store_with_byte_offset(frag, 0); + } + } + + /// Advances to the next position to load or store + MCTLASS_HOST_DEVICE + PredicatedTileIteratorStridedDgrad &operator++() { + + ++state_[0]; + + thread_start_row_ += ThreadMap::Shape::kRow; + + if (state_[0] == ThreadMap::Count::kRow) { + + state_[0] = 0; + ++state_[1]; + + thread_start_row_ += (ThreadMap::Shape::kGroup - 1) * + ThreadMap::Shape::kRow * ThreadMap::Count::kRow; + + if (state_[1] == ThreadMap::Count::kGroup) { + + state_[1] = 0; + ++state_[2]; + + thread_start_row_ += ThreadMap::Count::kGroup * + ThreadMap::Shape::kGroup * ThreadMap::Count::kRow * ThreadMap::Shape::kRow; + + if (state_[2] == ThreadMap::Count::kCluster) { + state_[2] = 0; + } + } + } + + return *this; + } + + ///< Efficiently disables all accesses guarded by mask + MCTLASS_DEVICE void clear_mask() { + mask_.clear(); + } + + ///< Efficiently enables all accesses guarded by mask + MCTLASS_DEVICE void enable_mask() { + mask_.enable(); + } + + ///< Sets the mask + MCTLASS_DEVICE void get_mask(Mask &mask) { + mask = mask_; + } + + ///< Sets the mask + MCTLASS_DEVICE void set_mask(Mask const &mask) { + mask_ = mask; + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/shared_load_iterator.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/shared_load_iterator.h new file mode 100644 index 0000000..e4c783d --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/shared_load_iterator.h @@ -0,0 +1,223 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" + +#include "mctlass/epilogue/threadblock/output_tile_thread_map.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load output tile from shared memory in epilogue. +/// +/// Satisfies: ReadableTileIterator +/// +template < + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_, ///< Element data type + int MaxAlignment = ThreadMap_::kElementsPerAccess * sizeof_bits::value / 8 +> +class SharedLoadIterator { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::TileShape; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + + static int const kMinAlignment = ThreadMap_::kElementsPerAccess * sizeof_bits::value / 8; + + static int const kAlignment = (MaxAlignment < kMinAlignment ? MaxAlignment : kMinAlignment); + + static int const kThreads = ThreadMap::kThreads; + + /// Fragment object + using Fragment = Array< + Element, + ThreadMap::Iterations::kColumn * + ThreadMap::Iterations::kRow * + ThreadMap::Iterations::kGroup * + ThreadMap::Iterations::kCluster * + ThreadMap::kElementsPerAccess>; + + /// Memory access size + using AccessType = AlignedArray< + Element, + ThreadMap::kElementsPerAccess, + kAlignment>; + + /// Vector type used for SMEM loads + using LoadType = AlignedArray< + Element, + const_min(128 / sizeof_bits::value, ThreadMap::kElementsPerAccess), + const_min(16, kAlignment) + >; + + static int const kLoadsPerAccess = AccessType::kElements / LoadType::kElements; + +private: + + // + // Data members + // + + /// Byte-level pointer + uint8_t *byte_pointer_; + + /// Stride along adjacent rows + int stride_; + +public: + + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + SharedLoadIterator( + TensorRef ref, + int thread_idx + ): + byte_pointer_(reinterpret_cast(ref.data())), + stride_((ref.stride(0) * sizeof_bits::value) / 8) { + + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx); + + // Initialize pointer + byte_pointer_ += + thread_offset.row() * stride_ + + thread_offset.column() * sizeof(AccessType) / kElementsPerAccess; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &offset) { + byte_pointer_ += + offset.row() * Shape::kRow * stride_ + + offset.column() * Shape::kColumn * sizeof_bits::value / 8; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + uint8_t const *byte_pointer = byte_pointer_ + + row * ThreadMap::Delta::kRow * stride_ + + group * ThreadMap::Delta::kGroup* stride_ + + cluster * ThreadMap::Delta::kCluster * stride_ + + pointer_offset * sizeof_bits::value / 8; + + int frag_row_idx = + (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + LoadType *frag_ptr = reinterpret_cast(&frag); + LoadType const *memory_pointer = reinterpret_cast(byte_pointer); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kLoadsPerAccess; ++v) { + frag_ptr[frag_idx * kLoadsPerAccess + v] = + memory_pointer[(column * ThreadMap::Delta::kColumn / kElementsPerAccess) * kLoadsPerAccess + v]; + } + } + } + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void set_smem_base_address(Index address) { + } + + /// Loads a fragment + MCTLASS_DEVICE + void load(Fragment &frag) const { + + load_with_pointer_offset(frag, 0); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/shared_load_iterator_mixed.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/shared_load_iterator_mixed.h new file mode 100644 index 0000000..51ba3cc --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/shared_load_iterator_mixed.h @@ -0,0 +1,594 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops optimized for mixed-precision. + + This assumes the shared memory tile is in a permuted layout which avoids bank conflicts on loading. + + When the fragment is loaded into registers, it matches the row-major thread map assumed by + the predicated tile iterator writing to global memory. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" + +#include "mctlass/epilogue/threadblock/output_tile_thread_map.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load output tile from shared memory in epilogue. +/// +/// Satisfies: ReadableTileIterator +/// +template < + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_, ///< Accumulator data type + int ElementSizeBits_, ///< Size of accumulator in bits + int OutputSizeBits_, ///< Size of output element in bits + int ElementsPerAccess, ///< Vector length of output vector + int ContiguousLanes, ///< Number of lanes in the warp writing to contiguous elements + /// in the global memory tensor + bool EightBitsOutputOrLess = (OutputSizeBits_ <= 8) +> +class SharedLoadIteratorMixed; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load output tile from shared memory in epilogue. +/// +/// Satisfies: ReadableTileIterator +/// +template < + typename ThreadMap_, ///< Thread map (conept: OutputTileThreadMap) + typename Element_ ///< Accumulator data type +> +class SharedLoadIteratorMixed { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + + static int const kAlignment = ThreadMap::kElementsPerAccess * sizeof_bits::value / 8; + + static int const kThreads = ThreadMap::kThreads; + + /// Fragment object + using Fragment = Array< + Element, + ThreadMap::Iterations::kColumn * + ThreadMap::Iterations::kRow * + ThreadMap::Iterations::kGroup * + ThreadMap::Iterations::kCluster * + ThreadMap::kElementsPerAccess>; + + /// Memory access size + using AccessType = AlignedArray< + Element, + ThreadMap::kElementsPerAccess, + kAlignment>; + + /// Vector type used for SMEM loads + using LoadType = AlignedArray< + Element, + const_min(128 / sizeof_bits::value, ThreadMap::kElementsPerAccess), + const_min(16, kAlignment) + >; + + static int const kLoadsPerAccess = AccessType::kElements / LoadType::kElements; + +private: + + // + // Data members + // + + /// Byte-level pointer + LoadType const *pointers_[kLoadsPerAccess]; + + /// Stride along adjacent rows in units of LoadType + int stride_; + +public: + + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + SharedLoadIteratorMixed( + TensorRef ref, + int thread_idx + ): + stride_((ref.stride(0) / LoadType::kElements)) { + + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx); + + // Initialize pointers + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kLoadsPerAccess; ++i) { + pointers_[i] = reinterpret_cast(ref.data()); + + int col_idx = (thread_offset.column() / kElementsPerAccess) * kLoadsPerAccess; + int bank_offset = (col_idx * int(sizeof(LoadType)) / 128) % kLoadsPerAccess; + + col_idx += (bank_offset + i) % kLoadsPerAccess; + + pointers_[i] += thread_offset.row() * stride_ + col_idx; + } + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kLoadsPerAccess; ++i) { + pointers_[i] += pointer_offset / LoadType::kElements; + } + } + + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &offset) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kLoadsPerAccess; ++i) { + pointers_[i] += + offset.row() * Shape::kRow * stride_ + + offset.column() * Shape::kColumn / LoadType::kElements; + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int row_ptr_offset = + row * ThreadMap::Delta::kRow * stride_ + + group * ThreadMap::Delta::kGroup* stride_ + + cluster * ThreadMap::Delta::kCluster * stride_ + + pointer_offset / LoadType::kElements; + + int frag_row_idx = (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + LoadType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kLoadsPerAccess; ++v) { + + int vector_idx = (column * ThreadMap::Delta::kColumn / kElementsPerAccess * kLoadsPerAccess); + + LoadType const *memory_pointer = pointers_[v] + row_ptr_offset; + + frag_ptr[frag_idx * kLoadsPerAccess + v] = memory_pointer[vector_idx]; + } + } + } + } + } + } + + /// Set base smem address + MCTLASS_DEVICE + void set_smem_base_address(Index address) {} + + /// Loads a fragment + MCTLASS_DEVICE + void load(Fragment &frag) const { + + load_with_pointer_offset(frag, 0); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for +/// int32_t x 16 => int8_t/int4b_t x 16 +/// float x 16 => float_e4m3_t/float_e5m2_t x 16 +template < + typename ThreadMap_, ///< Thread map (concept: OutputTileThreadMap) + typename Element_, + int OutputSizeBits_ ///< Size of output element in bits +> +class SharedLoadIteratorMixed { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + static_assert(sizeof_bits::value == 32, "Element size in bits must be 32."); + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + + static int const kAlignment = 16; + + static int const kThreads = ThreadMap::kThreads; + + /// Fragment object + using Fragment = Array< + Element, + ThreadMap::Iterations::kColumn * + ThreadMap::Iterations::kRow * + ThreadMap::Iterations::kGroup * + ThreadMap::Iterations::kCluster * + ThreadMap::kElementsPerAccess>; + + /// Memory access size + using AccessType = AlignedArray< + Element, + 16, + kAlignment>; + + /// Vector type used for SMEM loads + using LoadType = AlignedArray< + Element, + 4, + 16 + >; + + static int const kLoadsPerAccess = 4; + +private: + + // + // Data members + // + + /// Byte-level pointer + LoadType const *pointers_[kLoadsPerAccess]; + + /// Stride along adjacent rows in units of LoadType + int stride_; + +public: + + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + SharedLoadIteratorMixed( + TensorRef ref, + int thread_idx + ): + stride_((ref.stride(0) / LoadType::kElements)) { + + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx); + + // Initialize pointers + LoadType const *base_ptr = reinterpret_cast(ref.data()) + thread_offset.row() * stride_; + + int lane_col_idx = thread_offset.column() / 16; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kLoadsPerAccess; ++i) { + int lane_offset = (lane_col_idx % 2) * 4 | ((lane_col_idx / 2) * 8) | ((lane_col_idx / 2) ^ i); + + pointers_[i] = base_ptr + lane_offset; + } + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kLoadsPerAccess; ++i) { + pointers_[i] += pointer_offset / LoadType::kElements; + } + } + + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &offset) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kLoadsPerAccess; ++i) { + pointers_[i] += + offset.row() * Shape::kRow * stride_ + + offset.column() * Shape::kColumn / LoadType::kElements; + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int row_ptr_offset = + row * ThreadMap::Delta::kRow * stride_ + + group * ThreadMap::Delta::kGroup* stride_ + + cluster * ThreadMap::Delta::kCluster * stride_ + + pointer_offset / LoadType::kElements; + + int frag_row_idx = (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + LoadType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kLoadsPerAccess; ++v) { + + LoadType const *memory_pointer = pointers_[v]; + + frag_ptr[frag_idx * kLoadsPerAccess + v] = memory_pointer[row_ptr_offset]; + } + } + } + } + } + } + + /// Set base smem address + MCTLASS_DEVICE + void set_smem_base_address(Index address) {} + + /// Loads a fragment + MCTLASS_DEVICE + void load(Fragment &frag) { + + load_with_pointer_offset(frag, 0); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for: +/// int32_t x 8 => int8_t/int4b_t x 8 and +/// float x 8 => float_e4m3_t/float_e5m2_t x 8 +template < + typename ThreadMap_, ///< Thread map (concept: OutputTileThreadMap) + typename Element_, + int OutputSizeBits_ +> +class SharedLoadIteratorMixed { +public: + using ThreadMap = ThreadMap_; + using Shape = typename ThreadMap::Shape; + + using Element = Element_; + static_assert(sizeof_bits::value == 32, "Element size in bits must be 32."); + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + + static int const kAlignment = 8; + + static int const kThreads = ThreadMap::kThreads; + + /// Fragment object + using Fragment = Array< + Element, + ThreadMap::Iterations::kColumn * + ThreadMap::Iterations::kRow * + ThreadMap::Iterations::kGroup * + ThreadMap::Iterations::kCluster * + ThreadMap::kElementsPerAccess>; + + /// Memory access size + using AccessType = AlignedArray< + Element, + 8, + kAlignment>; + + /// Vector type used for SMEM loads + using LoadType = AlignedArray< + Element, + 4, + 16 + >; + + static int const kLoadsPerAccess = 2; + +private: + + // + // Data members + // + + /// Byte-level pointer + LoadType const *pointers_[kLoadsPerAccess]; + + /// Stride along adjacent rows in units of LoadType + int stride_; + +public: + + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + SharedLoadIteratorMixed( + TensorRef ref, + int thread_idx + ): + stride_((ref.stride(0) / LoadType::kElements)) { + + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx); + + // Initialize pointers + LoadType const *base_ptr = reinterpret_cast(ref.data()) + thread_offset.row() * stride_; + + int lane_col_idx = thread_offset.column() / 8; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kLoadsPerAccess; ++i) { + int lane_offset = (lane_col_idx % 8) * 2 | ((lane_col_idx / 4) ^ i); + + pointers_[i] = base_ptr + lane_offset; + } + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kLoadsPerAccess; ++i) { + pointers_[i] += pointer_offset / LoadType::kElements; + } + } + + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &offset) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kLoadsPerAccess; ++i) { + pointers_[i] += + offset.row() * Shape::kRow * stride_ + + offset.column() * Shape::kColumn / LoadType::kElements; + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + MCTLASS_PRAGMA_UNROLL + for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster; ++cluster) { + + MCTLASS_PRAGMA_UNROLL + for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) { + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) { + + int row_ptr_offset = + row * ThreadMap::Delta::kRow * stride_ + + group * ThreadMap::Delta::kGroup* stride_ + + cluster * ThreadMap::Delta::kCluster * stride_ + + pointer_offset / LoadType::kElements; + + int frag_row_idx = (row + ThreadMap::Iterations::kRow * (group + ThreadMap::Iterations::kGroup * cluster)); + + LoadType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int column = 0; column < ThreadMap::Iterations::kColumn; ++column) { + + int frag_idx = frag_row_idx * ThreadMap::Iterations::kColumn + column; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kLoadsPerAccess; ++v) { + + LoadType const *memory_pointer = pointers_[v]; + + frag_ptr[frag_idx * kLoadsPerAccess + v] = memory_pointer[row_ptr_offset]; + } + } + } + } + } + } + + /// Set base smem address + MCTLASS_DEVICE + void set_smem_base_address(Index address) {} + + /// Loads a fragment + MCTLASS_DEVICE + void load(Fragment &frag) { + + load_with_pointer_offset(frag, 0); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/threadblock/shared_load_iterator_pitch_liner.h b/csrc/mctlass/include/mctlass/epilogue/threadblock/shared_load_iterator_pitch_liner.h new file mode 100644 index 0000000..036b632 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/threadblock/shared_load_iterator_pitch_liner.h @@ -0,0 +1,194 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Epilogue for threadblock scoped GEMMs using Tensor Ops. + + This assumes the shared memory tile is in a permuted layout which avoids bank conflicts on loading. + + When the fragment is loaded into registers, it matches the row-major thread map assumed by + the predicated tile iterator writing to global memory. + + The epilogue rearranges the result of a matrix product through shared memory to match canonical + tensor layouts in global memory. Epilogues support conversion and reduction operations. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/epilogue/threadblock/output_tile_thread_map.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" +#include "mctlass/tensor_ref.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator used to load output tile from shared memory in epilogue. +/// +/// Satisfies: ReadableTileIterator +/// +template ::value / 8> +class SharedLoadIteratorPitchLiner { + public: + using ThreadMap = ThreadMap_; + using Element = Element_; + + using Layout = layout::RowMajor; + using TensorRef = TensorRef; + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using TensorCoord = MatrixCoord; + + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + + static int const kMinAlignment = + ThreadMap_::kElementsPerAccess * sizeof_bits::value / 8; + + static int const kAlignment = (MaxAlignment < kMinAlignment ? MaxAlignment : kMinAlignment); + + static int const kThreads = ThreadMap::kThreads; + + /// Fragment object + using Fragment = Array; + + /// Memory access size + using AccessType = AlignedArray; + + /// Vector type used for SMEM loads + using LoadType = + AlignedArray::value, ThreadMap::kElementsPerAccess), + const_min(16, kAlignment)>; + + static int const kLoadsPerAccess = AccessType::kElements / LoadType::kElements; + + private: + // + // Data members + // + + /// Byte-level pointer + uint8_t *byte_pointer_; + + /// Stride along adjacent rows + int stride_; + + /// Base address offset + Index base_smem_address_; + + public: + // + // Methods + // + + /// Constructor + MCTLASS_DEVICE + SharedLoadIteratorPitchLiner(TensorRef ref, int thread_idx) + : byte_pointer_(reinterpret_cast(ref.data())), + stride_((ref.stride(0) * sizeof_bits::value) / 8), + base_smem_address_(0) { + TensorCoord thread_offset = ThreadMap::initial_offset(thread_idx); + + // Initialize pointer + // thread_offset.row() is contiguous dim + // thread_offset.column() is stride dim + byte_pointer_ += thread_offset.row() * sizeof(AccessType) / kElementsPerAccess+ + thread_offset.column() * stride_ ; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_pointer_ += pointer_offset * sizeof_bits::value / 8; + } + + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &offset) { + byte_pointer_ += + offset.row() * ThreadMap::StorageShape::kContiguous * sizeof(AccessType) / kElementsPerAccess + + offset.column() * ThreadMap::StorageShape::kStrided * stride_; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + uint8_t const *byte_pointer = + byte_pointer_ + s * ThreadMap::Delta::kStrided * stride_ + + c * ThreadMap::Delta::kContiguous * ThreadMap::kElementsPerAccess * + sizeof_bits::value / 8 + + pointer_offset * sizeof_bits::value / 8 + base_smem_address_; + + int frag_base_idx = s * ThreadMap::Iterations::kContiguous + c; + + LoadType *frag_ptr = reinterpret_cast(&frag); + + LoadType const *memory_pointer = reinterpret_cast(byte_pointer); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kLoadsPerAccess; ++v) { + frag_ptr[frag_base_idx * kLoadsPerAccess + v] = memory_pointer[v]; + } + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void set_smem_base_address(Index address) { base_smem_address_ = address; } + + /// Loads a fragment + MCTLASS_DEVICE + void load(Fragment &frag) const { load_with_pointer_offset(frag, 0); } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_complex_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_complex_tensor_op.h new file mode 100644 index 0000000..bac91a5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_complex_tensor_op.h @@ -0,0 +1,187 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile + that participate in one warp-level store operation. + + Typically, the accumulator tile is the largest single block of register-backed storage + within the kernel. Storing it to memory is best accomplished by partitioning it into + smaller tiles and storing these sequentially. + + Round trips through shared memory during the Epilogue phase require partitioning, as + shared memory capacity is typically insufficient for a threadblock's total accumulator + size. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" + +#include "mctlass/epilogue/warp/tensor_op_policy.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +//////////////////////////////////////////////////////////////////////////////// + +/// +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename OperatorShape, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename OperatorElementC, ///< matrix multiply operation data type (concept: data type) + typename OperatorFragmentC, ///< matrix multiply operation fragment (concept: Array) + typename Layout ///< target shared memory layout +> +class FragmentIteratorComplexTensorOp; + +//////////////////////////////////////////////////////////////////////////////// + + +/// Partial specialization for row-major shared memory +template < + typename WarpShape_, ///< shape of the warp-level GEMM tile + typename OperatorShape_, ///< underlying real-valued matrix multiply operation shape (concept: gemm::GemmShape) + typename OperatorElementC_, ///< underlying real-valued matrix multiply operation data type + typename OperatorFragmentC_ ///< underlying real-valued matrix multiply operation fragment (concept: Array) +> +class FragmentIteratorComplexTensorOp { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using OperatorElementC = OperatorElementC_; + using OperatorFragmentC = OperatorFragmentC_; + using Layout = layout::RowMajor; + + using Policy = TensorOpPolicy; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array< + complex, + Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>; + + static int const kRealIndex = 0; + + /// Offset into the accumulator fragment + static int const kImaginaryIndex = + OperatorFragmentC::kElements * Policy::OperatorCount::kRow * Policy::OperatorCount::kColumn; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = Array; + + /// This is the complete warp-level accumulator tile. + using OutputAccumulatorTile = Array, kImaginaryIndex>; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + +private: + + /// Internal access type + using AccessType = Array; + + using FragmentAccessType = Array, Policy::kElementsPerAccess>; + +private: + + // + // Data members + // + + /// Accumulator tile + AccessType const *accumulators_; + + /// Internal index + int index_; + +public: + + /// Constructs an iterator + MCTLASS_HOST_DEVICE + FragmentIteratorComplexTensorOp(AccumulatorTile const &accum): + accumulators_(reinterpret_cast(&accum)), + index_(0) { + + } + + /// Increments + MCTLASS_HOST_DEVICE + FragmentIteratorComplexTensorOp &operator++() { + ++index_; + return *this; + } + + /// Decrements + MCTLASS_HOST_DEVICE + FragmentIteratorComplexTensorOp &operator--() { + --index_; + return *this; + } + + /// Loads a fragment from the referenced part of the accumulator tile + MCTLASS_HOST_DEVICE + void load(Fragment &frag, int index_offset = 0) const { + + int index = index_ + index_offset; + + FragmentAccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) { + + int accumulator_access_offset = + index + n * Policy::kAccumulatorColumnStride / Policy::kElementsPerAccess; + + auto const & real_accum_array = accumulators_[accumulator_access_offset + kRealIndex]; + auto const & imag_accum_array = accumulators_[accumulator_access_offset + kImaginaryIndex / Policy::kElementsPerAccess]; + + // Pack real and imaginary parts into a structure. This is likely to result in MOVs + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Policy::kElementsPerAccess; ++i) { + + frag_ptr[n][i].real() = real_accum_array[i]; + frag_ptr[n][i].imag() = imag_accum_array[i]; + } + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_gaussian_complex_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_gaussian_complex_tensor_op.h new file mode 100644 index 0000000..d87c351 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_gaussian_complex_tensor_op.h @@ -0,0 +1,194 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile + that participate in one warp-level store operation. + + Typically, the accumulator tile is the largest single block of register-backed storage + within the kernel. Storing it to memory is best accomplished by partitioning it into + smaller tiles and storing these sequentially. + + Round trips through shared memory during the Epilogue phase require partitioning, as + shared memory capacity is typically insufficient for a threadblock's total accumulator + size. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" + +#include "mctlass/epilogue/warp/tensor_op_policy.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +//////////////////////////////////////////////////////////////////////////////// + +/// +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename OperatorShape, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename OperatorElementC, ///< matrix multiply operation data type (concept: data type) + typename OperatorFragmentC, ///< matrix multiply operation fragment (concept: Array) + typename Layout ///< target shared memory layout +> +class FragmentIteratorGaussianComplexTensorOp; + +//////////////////////////////////////////////////////////////////////////////// + + +/// Partial specialization for row-major shared memory +template < + typename WarpShape_, ///< shape of the warp-level GEMM tile + typename OperatorShape_, ///< underlying real-valued matrix multiply operation shape (concept: gemm::GemmShape) + typename OperatorElementC_, ///< underlying real-valued matrix multiply operation data type + typename OperatorFragmentC_ ///< underlying real-valued matrix multiply operation fragment (concept: Array) +> +class FragmentIteratorGaussianComplexTensorOp { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using OperatorElementC = OperatorElementC_; + using OperatorFragmentC = OperatorFragmentC_; + using Layout = layout::RowMajor; + + using Policy = TensorOpPolicy; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array< + complex, + Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>; + + /// Size of one part of accumulator of 3-part accumulator in units of number of OperatorElementC + static int const kElementsAccumulatorPerPart = + OperatorFragmentC::kElements * Policy::OperatorCount::kRow * Policy::OperatorCount::kColumn; + + /// Offset into the accumulator fragment part 1 + static int const kPart1Index = kElementsAccumulatorPerPart * 0; + + /// Offset into the accumulator fragment part 2 + static int const kPart2Index = kElementsAccumulatorPerPart * 1; + + /// Offset into the accumulator fragment part 3 + static int const kPart3Index = kElementsAccumulatorPerPart * 2; + + /// This is the complete warp-level accumulator tile holding part1, part2, and part3 + using AccumulatorTile = Array; + + /// This is the complete warp-level accumulator tile holding final output of complex type + using OutputAccumulatorTile = Array, kElementsAccumulatorPerPart>; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + +private: + + /// Internal access type + using AccessType = Array; + + using FragmentAccessType = Array, Policy::kElementsPerAccess>; + +private: + + // + // Data members + // + + /// Accumulator tile + AccessType const *accumulators_; + + /// Internal index + int index_; + +public: + + /// Constructs an iterator + MCTLASS_HOST_DEVICE + FragmentIteratorGaussianComplexTensorOp(AccumulatorTile const &accum): + accumulators_(reinterpret_cast(&accum)), + index_(0) { + } + + /// Increments + MCTLASS_HOST_DEVICE + FragmentIteratorGaussianComplexTensorOp &operator++() { + ++index_; + return *this; + } + + /// Decrements + MCTLASS_HOST_DEVICE + FragmentIteratorGaussianComplexTensorOp &operator--() { + --index_; + return *this; + } + + /// Loads a fragment from the referenced part of the accumulator tile + MCTLASS_HOST_DEVICE + void load(Fragment &frag, int index_offset = 0) const { + + int index = index_ + index_offset; + + FragmentAccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) { + + int accumulator_access_offset = + index + n * Policy::kAccumulatorColumnStride / Policy::kElementsPerAccess; + + auto const & part1_accum_array = accumulators_[accumulator_access_offset + kPart1Index]; + auto const & part2_accum_array = accumulators_[accumulator_access_offset + kPart2Index / Policy::kElementsPerAccess]; + auto const & part3_accum_array = accumulators_[accumulator_access_offset + kPart3Index / Policy::kElementsPerAccess]; + + // Pack parts 1, 2, and 3 into a structure. This is likely to result in MOVs + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Policy::kElementsPerAccess; ++i) { + + frag_ptr[n][i].real() = part1_accum_array[i] - part3_accum_array[i]; + frag_ptr[n][i].imag() = part1_accum_array[i] + part2_accum_array[i]; + } + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_simt.h b/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_simt.h new file mode 100644 index 0000000..40803af --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_simt.h @@ -0,0 +1,164 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile + that participate in one warp-level store operation. + + Typically, the accumulator tile is the largest single block of register-backed storage + within the kernel. Storing it to memory is best accomplished by partitioning it into + smaller tiles and storing these sequentially. + + Round trips through shared memory during the Epilogue phase require partitioning, as + shared memory capacity is typically insufficient for a threadblock's total accumulator + size. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" + +#include "mctlass/epilogue/warp/simt_policy.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Fragment iterator for SIMT accumulator arrangements +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename Operator, ///< matrix multiply operation (concept: arch::Mma) + typename Layout, ///< target shared memory layout + typename MmaSimtPolicy ///< policy defining lane arrangement (concept: MmaSimtPolicy) +> +class FragmentIteratorSimt; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for row-major shared memory +template < + typename WarpShape_, ///< shape of the warp-level GEMM tile + typename Operator_ , ///< matrix multiply operator (concept: arch::Mma) + typename MmaSimtPolicy_ ///< policy defining lane arrangement (concept: MmaSimtPolicy) +> +class FragmentIteratorSimt { +public: + + using WarpShape = WarpShape_; + using Operator = Operator_; + using Layout = layout::RowMajor; + + /// Policy for warp-level epilogue components + using Policy = SimtPolicy; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array< + typename Operator::ElementC, + Policy::kElementsPerIteration>; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = Array< + typename Operator::ElementC, + Policy::kAccumulatorElementCount>; + + using OutputAccumulatorTile = AccumulatorTile; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + +private: + + /// Internal access type + using AccessType = Array; + +private: + + // + // Data members + // + + /// Accumulator tile + AccessType const *accumulators_; + + /// Internal index + int index_; + +public: + + /// Constructs an iterator + MCTLASS_HOST_DEVICE + FragmentIteratorSimt(AccumulatorTile const &accum): + accumulators_(reinterpret_cast(&accum)), + index_(0) { + + } + + /// Increments + MCTLASS_HOST_DEVICE + FragmentIteratorSimt &operator++() { + ++index_; + return *this; + } + + /// Decrements + MCTLASS_HOST_DEVICE + FragmentIteratorSimt &operator--() { + --index_; + return *this; + } + + /// Loads a fragment from the referenced part of the accumulator tile + MCTLASS_HOST_DEVICE + void load(Fragment &frag, int index_offset = 0) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::kAccessesPerIteration; ++n) { + + int accumulator_access_offset = index_ * Policy::kAccessesPerIteration + n; + + frag_ptr[n] = accumulators_[accumulator_access_offset]; + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_tensor_op.h new file mode 100644 index 0000000..34461cb --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_tensor_op.h @@ -0,0 +1,277 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile + that participate in one warp-level store operation. + + Typically, the accumulator tile is the largest single block of register-backed storage + within the kernel. Storing it to memory is best accomplished by partitioning it into + smaller tiles and storing these sequentially. + + Round trips through shared memory during the Epilogue phase require partitioning, as + shared memory capacity is typically insufficient for a threadblock's total accumulator + size. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" + +#include "mctlass/epilogue/warp/tensor_op_policy.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +//////////////////////////////////////////////////////////////////////////////// + +/// +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename OperatorShape, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename OperatorElementC, ///< matrix multiply operation data type (concept: data type) + typename OperatorFragmentC, ///< matrix multiply operation fragment (concept: Array) + typename Layout ///< target shared memory layout +> +class FragmentIteratorTensorOp; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for row-major shared memory +template < + typename WarpShape_, ///< shape of the warp-level GEMM tile + typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename OperatorElementC_, ///< matrix multiply operation data type (concept: data type) + typename OperatorFragmentC_ ///< matrix multiply operation fragment (concept: Array) +> +class FragmentIteratorTensorOp { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using OperatorElementC = OperatorElementC_; + using OperatorFragmentC = OperatorFragmentC_; + using Layout = layout::RowMajor; + + using Policy = TensorOpPolicy; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array< + OperatorElementC, + Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = Array< + OperatorElementC, + OperatorFragmentC::kElements * Policy::OperatorCount::kRow * Policy::OperatorCount::kColumn>; + + using OutputAccumulatorTile = AccumulatorTile; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + using TileIterations = typename Policy::TileIterations; + static int const kIterationsPerTile = kIterations / TileIterations::kCount; + +private: + + /// Internal access type + using AccessType = Array; + +private: + + // + // Data members + // + + /// Accumulator tile + AccessType const *accumulators_; + + /// Internal index + int index_; + +public: + + /// Constructs an iterator + MCTLASS_HOST_DEVICE + FragmentIteratorTensorOp(AccumulatorTile const &accum): + accumulators_(reinterpret_cast(&accum)), + index_(0) { + } + + /// Increments + MCTLASS_HOST_DEVICE + FragmentIteratorTensorOp &operator++() { + ++index_; + return *this; + } + + /// Decrements + MCTLASS_HOST_DEVICE + FragmentIteratorTensorOp &operator--() { + --index_; + return *this; + } + + /// Loads a fragment from the referenced part of the accumulator tile + MCTLASS_HOST_DEVICE + void load(Fragment &frag, int index_offset = 0) const { + + int index = index_ + index_offset; + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) { + + int accumulator_access_offset = + index + n * Policy::kAccumulatorColumnStride / Policy::kElementsPerAccess; + + frag_ptr[n] = accumulators_[accumulator_access_offset]; + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Dedicated to interleaved layout +template < + /// shape of the warp-level GEMM tile + typename WarpShape_, + /// matrix multiply operator shape (concept: gemm::GemmShape) + typename OperatorShape_, + /// matrix multiply operator data type (concept: data type) + typename OperatorElementC_, + /// matrix multiply operator fragment (concept: Array) + typename OperatorFragmentC_, + /// number of interleaved k + int InterleavedK> +class FragmentIteratorTensorOp> { + public: + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using OperatorElementC = OperatorElementC_; + using OperatorFragmentC = OperatorFragmentC_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::ColumnMajorInterleaved; + + using Policy = TensorOpPolicy; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = + Array; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = + Array; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + using TileIterations = typename Policy::TileIterations; + static int const kIterationsPerTile = kIterations / TileIterations::kCount; + + private: + /// Internal access type + using AccessType = + Array; + + private: + // + // Data members + // + + /// Accumulator tile + AccessType const *accumulators_; + + /// Internal index + int index_; + + public: + /// Constructs an iterator + MCTLASS_HOST_DEVICE + FragmentIteratorTensorOp(AccumulatorTile const &accum) + : accumulators_(reinterpret_cast(&accum)), + index_(0) {} + + /// Increments + MCTLASS_HOST_DEVICE + FragmentIteratorTensorOp &operator++() { + ++index_; + return *this; + } + + /// Decrements + MCTLASS_HOST_DEVICE + FragmentIteratorTensorOp &operator--() { + --index_; + return *this; + } + + /// Loads a fragment from the referenced part of the accumulator tile + MCTLASS_HOST_DEVICE + void load(Fragment &frag, int index_offset = 0) const { + int index = index_ + index_offset; + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < (InterleavedK / OperatorShape::kN); ++n) { + int index_m = index % (Policy::OperatorCount::kRow * + Policy::kIterationsPerInstruction); + int index_n = index / (Policy::OperatorCount::kRow * + Policy::kIterationsPerInstruction); + int accumulator_access_offset = + (index_m / Policy::kIterationsPerInstruction) * + (Policy::OperatorCount::kColumn * + Policy::kIterationsPerInstruction) + + (index_m % Policy::kIterationsPerInstruction) + + index_n * (InterleavedK / OperatorShape::kN) * + Policy::kIterationsPerInstruction + + n * Policy::kIterationsPerInstruction; + + frag_ptr[n] = accumulators_[accumulator_access_offset]; + } + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_volta_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_volta_tensor_op.h new file mode 100644 index 0000000..241e67c --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_volta_tensor_op.h @@ -0,0 +1,269 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile + that participate in one warp-level store operation. + + Typically, the accumulator tile is the largest single block of register-backed storage + within the kernel. Storing it to memory is best accomplished by partitioning it into + smaller tiles and storing these sequentially. + + Round trips through shared memory during the Epilogue phase require partitioning, as + shared memory capacity is typically insufficient for a threadblock's total accumulator + size. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/gemm/gemm.h" + +#include "mctlass/epilogue/warp/volta_tensor_op_policy.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename InterleavedTileShape, ///< shape of indivisible instruction-level arrangement (concept: GemmShape) + typename ElementC, ///< Accumulator layout + typename Layout ///< target shared memory layout +> +class FragmentIteratorVoltaTensorOp; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for row-major shared memory +template < + typename WarpShape_ ///< shape of warp-level GEMM (concept: MatrixShape) +> +class FragmentIteratorVoltaTensorOp, half_t, layout::RowMajor> { +public: + + using WarpShape = WarpShape_; + using InterleavedTileShape = gemm::GemmShape<32, 32, 4>; + using ElementC = half_t; + using Layout = layout::RowMajor; + + /// Policy operator + using Policy = VoltaTensorOpPolicy; + + /// Array type for aligned memory accesses + using AccessType = typename Policy::AccessType; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = typename Policy::Fragment; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = typename Policy::AccumulatorTile; + + using OutputAccumulatorTile = AccumulatorTile; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + +private: + +private: + + // + // Data members + // + + /// Accumulator tile + AccessType const *accumulators_; + + /// Internal index + int index_; + +public: + + /// Constructs an iterator + MCTLASS_HOST_DEVICE + FragmentIteratorVoltaTensorOp(AccumulatorTile const &accum): + accumulators_(reinterpret_cast(&accum)), + index_(0) { + + } + + /// Increments + MCTLASS_HOST_DEVICE + FragmentIteratorVoltaTensorOp &operator++() { + ++index_; + return *this; + } + + /// Decrements + MCTLASS_HOST_DEVICE + FragmentIteratorVoltaTensorOp &operator--() { + --index_; + return *this; + } + + /// Loads a fragment from the referenced part of the accumulator tile + MCTLASS_HOST_DEVICE + void load(Fragment &frag, int index_offset = 0) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + static int const kAccessesPerMma = Policy::kElementsPerMma / Policy::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; ++tile_n) { + + int tile_access_idx = + (tile_n * Policy::TileIterations::kRow + (index_ & 2) / 2) * Policy::MmaIterations::kCount * kAccessesPerMma; + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn * kAccessesPerMma; ++mma_n) { + + int mma_access_idx = ((mma_n & 1) * 2 + (index_ & 1)) * kAccessesPerMma + (mma_n & 2) / 2; + + frag_ptr[tile_n * Policy::MmaIterations::kColumn * kAccessesPerMma + + mma_n] = accumulators_[tile_access_idx + mma_access_idx]; + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for row-major shared memory +template < + typename WarpShape_ ///< shape of warp-level GEMM (concept: MatrixShape) +> +class FragmentIteratorVoltaTensorOp, float, layout::RowMajor> { +public: + + using WarpShape = WarpShape_; + using InterleavedTileShape = gemm::GemmShape<32, 32, 4>; + using ElementC = float; + using Layout = layout::RowMajor; + + /// Policy operator + using Policy = VoltaTensorOpPolicy; + + /// Array type for aligned memory accesses + using AccessType = typename Policy::AccessType; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = typename Policy::Fragment; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = typename Policy::AccumulatorTile; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + +private: + +private: + + // + // Data members + // + + /// Accumulator tile + AccessType const *accumulators_; + + /// Internal index + int index_; + +public: + + /// Constructs an iterator + MCTLASS_HOST_DEVICE + FragmentIteratorVoltaTensorOp(AccumulatorTile const &accum): + accumulators_(reinterpret_cast(&accum)), + index_(0) { + } + + /// Increments + MCTLASS_HOST_DEVICE + FragmentIteratorVoltaTensorOp &operator++() { + ++index_; + return *this; + } + + /// Decrements + MCTLASS_HOST_DEVICE + FragmentIteratorVoltaTensorOp &operator--() { + --index_; + return *this; + } + + /// Loads a fragment from the referenced part of the accumulator tile + MCTLASS_HOST_DEVICE + void load(Fragment &frag, int index_offset = 0) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + int const kRegsPerMmaRow = 2; + + MCTLASS_PRAGMA_UNROLL + for (int reg_row = 0; reg_row < Policy::kRowsPerMmaTile; ++reg_row) { + + MCTLASS_PRAGMA_UNROLL + for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; ++tile_n) { + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn * 2; ++mma_n) { + + int mma_idx = (index_ & 1) + (index_ & 2) * Policy::MmaIterations::kCount / 2 + + (tile_n * Policy::TileIterations::kRow) * Policy::MmaIterations::kCount + (mma_n & 1) * 2; + + int reg_offset = reg_row * kRegsPerMmaRow + (mma_n & 2) * 2; + int reg_idx = mma_idx * Policy::kElementsPerMma + reg_offset; + + *frag_ptr = accumulators_[reg_idx / Policy::kElementsPerAccess]; + ++frag_ptr; + } + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_wmma_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_wmma_tensor_op.h new file mode 100644 index 0000000..e5bb54f --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/fragment_iterator_wmma_tensor_op.h @@ -0,0 +1,165 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief This defines a "fragment" iterator for visiting the fragments of an accumulator tile + that participate in one warp-level store operation. + + Typically, the accumulator tile is the largest single block of register-backed storage + within the kernel. Storing it to memory is best accomplished by partitioning it into + smaller tiles and storing these sequentially. + + Round trips through shared memory during the Epilogue phase require partitioning, as + shared memory capacity is typically insufficient for a threadblock's total accumulator + size. +*/ + +#pragma once + +//#if !(defined(__clang__) && defined(__MACA__)) +#if defined(__MACA__) + +#include "mctlass/wmma_array.h" +#include "mctlass/layout/matrix.h" + +#include "mctlass/epilogue/warp/wmma_tensor_op_policy.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +//////////////////////////////////////////////////////////////////////////////// + +/// +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename OperatorShape, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename OperatorElementC, ///< matrix multiply operation data type (concept: data type) + typename OperatorFragmentC, ///< matrix multiply operation fragment + typename Layout ///< target shared memory layout +> +class FragmentIteratorWmmaTensorOp; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for row-major shared memory +template < + typename WarpShape_, ///< shape of the warp-level GEMM tile + typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename OperatorElementC_, ///< matrix multiply operation data type (concept: data type) + typename OperatorFragmentC_ ///< matrix multiply operation fragment +> +class FragmentIteratorWmmaTensorOp { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using OperatorElementC = OperatorElementC_; + using OperatorFragmentC = OperatorFragmentC_; + using Layout = layout::RowMajor; + + using Policy = WmmaTensorOpPolicy; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = WmmaFragmentArray; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = WmmaFragmentArray; + + using OutputAccumulatorTile = AccumulatorTile; + +private: + + /// Internal access type + using AccessType = WmmaFragmentArray; + +private: + + // + // Data members + // + + /// Accumulator tile + AccessType const *accumulators_; + + /// Internal index + int index_; + +public: + + /// Constructs an iterator + MCTLASS_HOST_DEVICE + FragmentIteratorWmmaTensorOp(AccumulatorTile const &accum): + accumulators_(reinterpret_cast(&accum)), + index_(0) { + } + + /// Increments + MCTLASS_HOST_DEVICE + FragmentIteratorWmmaTensorOp &operator++() { + ++index_; + return *this; + } + + /// Decrements + MCTLASS_HOST_DEVICE + FragmentIteratorWmmaTensorOp &operator--() { + --index_; + return *this; + } + + /// Loads a fragment from the referenced part of the accumulator tile + MCTLASS_HOST_DEVICE + void load(Fragment &frag, int index_offset = 0) const { + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for(int n=0; n < Policy::OperatorCount::kColumn; n++) { + + int accumulator_access_offset = index_ * Policy::OperatorCount::kColumn + n; + + frag_ptr[n] = accumulators_[accumulator_access_offset]; + } + } +}; + + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// + +#else +#error (defined(__clang__) && defined(__MACA__)) +#endif // !defined(__clang__) + diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/simt_policy.h b/csrc/mctlass/include/mctlass/epilogue/warp/simt_policy.h new file mode 100644 index 0000000..fd3056f --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/simt_policy.h @@ -0,0 +1,107 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic structures needed for implementing the warp-scoped phase of the epilogue. + These quantities assume a 'column-major' arrangement of SimtOp instructions, of which + a row-oriented slice is visible per iteration. +*/ + +#pragma once + +#include "mctlass/matrix_shape.h" +#include "mctlass/layout/matrix.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: GemmShape) + typename Operator, ///< matrix multiply operation (concept: arch::Mma) + typename Layout, ///< destination layout in shared memory + typename MmaSimtPolicy ///< policy defining lane arrangement (concept: MmaSimtPolicy) +> +struct SimtPolicy; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for row-major +template < + typename WarpShape_, ///< shape of warp-level GEMM (concept: MatrixShape) + typename Operator_, ///< matrix multiply operation (concept: arch::Mma) + typename MmaSimtPolicy_ ///< policy defining lane arrangement (concept: MmaSimtPolicy) +> +struct SimtPolicy { + + using WarpShape = WarpShape_; + using Operator = Operator_; + using MmaSimtPolicy = MmaSimtPolicy_; + + static_assert(!(WarpShape::kM % MmaSimtPolicy::WarpShape::kRow), "Divisibility"); + static_assert(!(WarpShape::kN % MmaSimtPolicy::WarpShape::kColumn), "Divisibility"); + + /// Number of iterations + static int const kIterations = WarpShape::kM / MmaSimtPolicy::WarpShape::kRow; + + /// Number of accumulators written per iteration + static int const kElementsPerIteration = + (WarpShape::kN / MmaSimtPolicy::WarpShape::kColumn); + + /// Total number of accumulators + static int const kAccumulatorElementCount = kElementsPerIteration * kIterations; + + /// Number of consecutive elements + static int const kElementsPerAccess = MmaSimtPolicy::LaneMmaShape::kN; + + /// Number of rows per epilogue iteration + static int const kRowsPerIteration = MmaSimtPolicy::WarpShape::kRow; + + /// Number of accesses made in one iteration + static int const kAccessesPerIteration = kElementsPerIteration / kElementsPerAccess; + + /// Number of elements in between accumulator chunks of (LaneMmaShape::kM x LaneMmaShape::kN) + using Delta = MatrixShape< + MmaSimtPolicy::WarpShape::kRow * MmaSimtPolicy::LaneMmaShape::kM, + MmaSimtPolicy::WarpShape::kColumn * MmaSimtPolicy::LaneMmaShape::kN + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/tensor_op_policy.h b/csrc/mctlass/include/mctlass/epilogue/warp/tensor_op_policy.h new file mode 100644 index 0000000..693a9b2 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/tensor_op_policy.h @@ -0,0 +1,148 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic structures needed for implementing the warp-scoped phase of the epilogue. + These quantities assume a 'column-major' arrangement of TensorOp instructions, of which + a row-oriented slice is visible per iteration. +*/ + +#pragma once + +#include "mctlass/matrix_shape.h" +#include "mctlass/layout/matrix.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +//////////////////////////////////////////////////////////////////////////////// + +/// Policy details related to the epilogue +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename OperatorShape, ///< matrix multiply operation shape (concept: gemm:GemmShape) + typename Layout ///< target shared memory layout +> +struct TensorOpPolicy; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for row-major +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename OperatorShape ///< matrix multiply operation shape (concept: gemm::GemmShape) +> +struct TensorOpPolicy { + + /// Number of operations + using OperatorCount = MatrixShape< + (WarpShape::kM + OperatorShape::kM - 1) / OperatorShape::kM, + (WarpShape::kN + OperatorShape::kN - 1) / OperatorShape::kN + >; + + // + // Hard-coded constants regarding Tensor Operations + // + + static int const kElementsPerAccess = 2; + static int const kRowsPerIteration = 8; + static bool const kDivisible = + !(WarpShape::kM % OperatorShape::kM) && !(WarpShape::kN % OperatorShape::kN); + + // + // Derived quantities + // + + // Number of 'externally visible' iterations per actual instruction + static int const kIterationsPerInstruction = OperatorShape::kM / kRowsPerIteration; + + // Number of externally visible iterations + static int const kIterations = OperatorCount::kRow * kIterationsPerInstruction; + + using TileIterations = MatrixShape; + + static int const kAccumulatorRowStride = kElementsPerAccess; + static int const kAccumulatorColumnStride = kElementsPerAccess * OperatorCount::kRow * kIterationsPerInstruction; + +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major-interleaved +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename OperatorShape, ///< matrix multiply operation (concept: arch::Mma) + int InterleavedK ///< number of interleaved k + > +struct TensorOpPolicy > { + /// Number of operations + using OperatorCount = MatrixShape; + + // + // Hard-coded constants regarding Tensor Operations + // + + static int const kElementsPerAccess = 2; + static int const kRowsPerIteration = 8; + + // + // Derived quantities + // + + // Number of 'externally visible' iterations per actual instruction + static int const kIterationsPerInstruction = + OperatorShape::kM / kRowsPerIteration; + + // Number of externally visible iterations + static int const kIterations = WarpShape::kN / InterleavedK * + OperatorCount::kRow * + kIterationsPerInstruction; + + static int const kElementsPerIteration = InterleavedK / OperatorShape::kN * kElementsPerAccess; + + static int const kAccessPerIteration = kElementsPerIteration / kElementsPerAccess; + + // Number of externally visible iterations + //static int const kTileIterations = OperatorCount::kRow * kIterationsPerInstruction; + using TileIterations = MatrixShape<1, WarpShape::kN / InterleavedK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_simt.h b/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_simt.h new file mode 100644 index 0000000..d04bcd1 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_simt.h @@ -0,0 +1,790 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" + +#include "mctlass/epilogue/warp/simt_policy.h" + +#define MCTLASS_SIMT_EPILOGUE_USE_SCALAR_STORES 1 + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename Operator, ///< matrix multiply operation (concept: arch::Mma) + typename Element, ///< data type of element to be written + typename Layout, ///< target shared memory layout + typename MmaSimtPolicy ///< policy defining lane arrangement (concept: MmaSimtPolicy) +> +class TileIteratorSimt; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template < + typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape) + typename Operator_, ///< matrix multiply operation (concept: arch::Mma) + typename Element_, ///< data type of element to be written + typename MmaSimtPolicy_ ///< policy defining lane arrangement (concept: MmaSimtPolicy) +> +class TileIteratorSimt { +public: + + using WarpShape = WarpShape_; + using Operator = Operator_; + using Element = Element_; + using Layout = layout::RowMajor; + + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + using Policy = SimtPolicy; + + /// Shape of the tile in memory + using Shape = MatrixShape< + Policy::kRowsPerIteration, + WarpShape::kN + >; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array< + typename Operator::ElementC, + Policy::kElementsPerIteration>; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = Array< + typename Operator::ElementC, + Policy::kAccumulatorElementCount>; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + + /// Padding quantity + using Padding = MatrixShape< + 0, + 4 * Policy::kElementsPerAccess +#if MCTLASS_SIMT_EPILOGUE_USE_SCALAR_STORES + + 1 +#endif + >; + +private: + +#if MCTLASS_SIMT_EPILOGUE_USE_SCALAR_STORES + /// Storage type for accessing memory + using AccessType = AlignedArray< + Element, + 1 + >; + +#else + /// Storage type for accessing memory + using AccessType = AlignedArray< + Element, + Policy::kElementsPerAccess + >; +#endif + + // + // Data members + // + + /// Internal pointer to memory + AccessType *pointer_; + + /// Internal layout object + Layout layout_; + +public: + + /// Default constructor + MCTLASS_HOST_DEVICE + TileIteratorSimt(): pointer_(nullptr) { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + TileIteratorSimt( + TensorRef const &ref, + unsigned lane_id + ): + pointer_(reinterpret_cast(ref.data())), + layout_(ref.stride()[0] / AccessType::kElements) { + + auto lane_layout = Policy::MmaSimtPolicy::get_lane_layout(); + MatrixCoord lane_offset = lane_layout.inverse(lane_id); + + pointer_ += layout_({ + lane_offset.row(), + lane_offset.column() * Policy::kElementsPerAccess / int(AccessType::kElements) + }); + } + + /// Adds a pointer offset + MCTLASS_HOST_DEVICE + TileIteratorSimt & add_pointer_offset(Index pointer_offset) { + pointer_ += pointer_offset / AccessType::kElements; + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorSimt & add_tile_offset(TensorCoord const &tile_offset) { + + pointer_ += layout_({ + tile_offset.row() * Shape::kRow, + (tile_offset.column() * Shape::kColumn / int(AccessType::kElements)) + }); + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorSimt & operator+=(TensorCoord const &tile_offset) { + + add_tile_offset(tile_offset); + + return *this; + } + + /// Store + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { +#if MCTLASS_SIMT_EPILOGUE_USE_SCALAR_STORES + // de-vectorized stores + using ScalarAccessType = AlignedArray; + ScalarAccessType const *scalarFragPtr = reinterpret_cast(&frag); + ScalarAccessType *scalarPointer = reinterpret_cast(pointer_) + pointer_offset; + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::kAccessesPerIteration; ++n) { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::kElementsPerAccess; s++) { + scalarPointer[n * Policy::MmaSimtPolicy::WarpShape::kColumn * Policy::kElementsPerAccess + s] = scalarFragPtr[n * Policy::kElementsPerAccess + s]; + } + } +#else + // original vector stores + AccessType const *frag_ptr = reinterpret_cast(&frag); + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::kAccessesPerIteration; ++n) { + pointer_[n * Policy::MmaSimtPolicy::WarpShape::kColumn + pointer_offset / int(AccessType::kElements)] = frag_ptr[n]; + } +#endif + } + + /// Store + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { +#if defined(__MACA_ARCH__) + if (__lane_id() < 32) { + store_with_pointer_offset(frag, 0); + } +#endif + } + + /// Load + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::kAccessesPerIteration; ++n) { + frag_ptr[n] = pointer_[n * Policy::MmaSimtPolicy::WarpShape::kColumn + pointer_offset / int(AccessType::kElements)]; + } + } + + /// Load + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Set smem base address + MCTLASS_HOST_DEVICE + void set_smem_base_address(Index address) { + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template +class TileIteratorSimtDirectConv { + public: + + using WarpShape = WarpShape_; + using Operator = Operator_; + using Element = Element_; + using Layout = layout::RowMajor; + + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + using Policy = SimtPolicy; + + /// Shape of the tile in memory + using Shape = MatrixShape; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = Array; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + + /// Padding quantity + using Padding = MatrixShape<0, + 0 + >; + +private: + /// Storage type for accessing memory + using AccessType = AlignedArray< + Element, + Policy::kElementsPerAccess + >; + + // + // Data members + // + + /// Internal pointer to memory + AccessType *pointer_; + + /// Internal layout object + Layout layout_; + + /// Base smem offset; + Index base_smem_address_; + + public: + /// Default constructor + MCTLASS_HOST_DEVICE + TileIteratorSimtDirectConv() : pointer_(nullptr) {} + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + TileIteratorSimtDirectConv( + TensorRef const &ref, + unsigned lane_id + ): + pointer_(reinterpret_cast(ref.data())), + layout_(ref.stride()[0] / AccessType::kElements) { + + auto lane_layout = Policy::MmaSimtPolicy::get_lane_layout(); + MatrixCoord lane_offset = lane_layout.inverse(lane_id); + + pointer_ += layout_({ + lane_offset.row(), + lane_offset.column() * Policy::kElementsPerAccess / int(AccessType::kElements) + }); + } + + /// Adds a pointer offset + MCTLASS_HOST_DEVICE + TileIteratorSimtDirectConv & add_pointer_offset(Index pointer_offset) { + pointer_ += pointer_offset / AccessType::kElements; + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorSimtDirectConv & add_tile_offset(TensorCoord const &tile_offset) { + + pointer_ += layout_({ + tile_offset.row() * Shape::kRow, + (tile_offset.column() * Shape::kColumn / int(AccessType::kElements)) + }); + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorSimtDirectConv & operator+=(TensorCoord const &tile_offset) { + + add_tile_offset(tile_offset); + + return *this; + } + + /// Store + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + // original vector stores + AccessType const *frag_ptr = reinterpret_cast(&frag); + AccessType * load_pointer_ = reinterpret_cast(reinterpret_cast(pointer_) + base_smem_address_); + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::kAccessesPerIteration; ++n) { + load_pointer_[n * Policy::MmaSimtPolicy::WarpShape::kColumn + pointer_offset / int(AccessType::kElements)] = frag_ptr[n]; + } + } + + /// Store + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } + + /// Load + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::kAccessesPerIteration; ++n) { + frag_ptr[n] = pointer_[n * Policy::MmaSimtPolicy::WarpShape::kColumn + pointer_offset / int(AccessType::kElements)]; + } + } + + /// Load + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Set smem base address + MCTLASS_HOST_DEVICE + void set_smem_base_address(Index address){ + base_smem_address_ = address; + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Template for reading and writing tiles of accumulators to shared memory +template +class TileIteratorSimtDirect2dConv { + public: + using WarpShape = WarpShape_; + using ThreadOutputShape = ThreadOutputShape_; + using ThreadBlockOutputShape = ThreadBlockOutputShape_; + using Operator = Operator_; + using Element = Element_; + using Layout = layout::RowMajor; + using MmaSimtPolicy = MmaSimtPolicy_; + + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + // Thread-level shape of a fragment + using ThreadShape = MatrixShape; + + static_assert(!(ThreadShape::kColumn % MmaSimtPolicy::LaneMmaShape::kN), + "Thread-level GEMM must be divisible by Policy::LaneMmaShape."); + + using ThreadTileCount = MatrixShape; + + using Iterations = + MatrixShape; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = typename Operator::FragmentC; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = AccumulatorTile; + + /// Padding quantity + using Padding = MatrixShape<0, 0>; + + private: + // Storage type for accessing memory + using AccessType = AlignedArray; + // + // Data members + // + + /// Internal pointer to memory + AccessType *pointer_; + + /// Internal layout object + Layout layout_; + + /// Base smem offset; + Index base_smem_address_; + + public: + /// Default constructor + MCTLASS_HOST_DEVICE + TileIteratorSimtDirect2dConv() : pointer_(nullptr) {} + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + TileIteratorSimtDirect2dConv(TensorRef const &ref, unsigned thread_id, unsigned lane_id) + : pointer_(reinterpret_cast(ref.data())), + layout_(ref.stride()[0] / AccessType::kElements) { + + auto lane_layout = MmaSimtPolicy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id); + + // Get base HW offset of current threads + const int threadgroup = thread_id / (ThreadBlockOutputShape::kC / ThreadOutputShape::kC); + const int base_p = (threadgroup / (ThreadTileCount::kColumn)) * ThreadOutputShape::kH; + const int base_q = (threadgroup % (ThreadTileCount::kColumn)) * ThreadOutputShape::kW; + + const int row_offset = base_p * ThreadBlockOutputShape::kW + base_q; + + pointer_ += layout_( + {row_offset, + lane_offset.column() * MmaSimtPolicy::LaneMmaShape::kN / int(AccessType::kElements)}); + } + + /// Adds a pointer offset + MCTLASS_HOST_DEVICE + TileIteratorSimtDirect2dConv &add_pointer_offset(Index pointer_offset) { + pointer_ += pointer_offset / AccessType::kElements; + return *this; + } + + /// Store + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + AccessType *storer_pointer_ = + reinterpret_cast(reinterpret_cast(pointer_) + base_smem_address_); + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int h = 0; h < ThreadOutputShape::kH; ++h) { + MCTLASS_PRAGMA_UNROLL + for (int w = 0; w < ThreadOutputShape::kW; ++w) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < Iterations::kColumn; ++col) { + int offset = (w + h * ThreadBlockOutputShape::kW) * + (ThreadBlockOutputShape::kC / AccessType::kElements) + + col; + storer_pointer_[offset + pointer_offset / int(AccessType::kElements)] = + frag_ptr[w + h * ThreadOutputShape::kW + col]; + } + } + } + } + + /// Store + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } + + /// Set smem base address + MCTLASS_HOST_DEVICE + void set_smem_base_address(Index address) { base_smem_address_ = address; } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template < + typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape) + typename Operator_, ///< matrix multiply operation (concept: arch::Mma) + typename Element_, ///< data type of element to be written + typename Layout_, ///< target shared memory layout + typename MmaSimtPolicy_ ///< policy defining lane arrangement (concept: MmaSimtPolicy) +> +class TileIteratorSimtCanonical { +public: + + using WarpShape = WarpShape_; + using Operator = Operator_; + using Element = Element_; + using Layout = Layout_; + + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + using Policy = SimtPolicy; + + /// Shape of the tile in memory + using Shape = MatrixShape< + Policy::kRowsPerIteration, + WarpShape::kN + >; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array< + typename Operator::ElementC, + Policy::kElementsPerIteration>; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = Array< + typename Operator::ElementC, + Policy::kAccumulatorElementCount>; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + + /// Padding quantity + using Padding = MatrixShape< + 0, + 4 * Policy::kElementsPerAccess + 1 + >; + +private: + + /// Storage type for accessing memory + using AccessType = AlignedArray< + Element, + 1 + >; + + // + // Data members + // + + /// Internal pointer to memory + AccessType *pointer_; + + /// Internal layout object + Layout layout_; + + /// Guard to indicate whether the shape is divisible + bool divisible_; + + /// Extent of the output tensor + MatrixCoord extent_; + + /// Thread offset + MatrixCoord thread_offset_; + +public: + + /// Default constructor + MCTLASS_HOST_DEVICE + TileIteratorSimtCanonical(): pointer_(nullptr) { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + TileIteratorSimtCanonical( + TensorRef const &ref, + unsigned lane_id + ): + pointer_(reinterpret_cast(ref.data())), + layout_(ref.stride()[0] / AccessType::kElements), + divisible_(true), + extent_(WarpShape::kM, WarpShape::kN) { + + auto lane_layout = Policy::MmaSimtPolicy::get_lane_layout(); + MatrixCoord lane_offset = lane_layout.inverse(lane_id); + + thread_offset_ = { + lane_offset.row() * Shape::kRow, + lane_offset.column() * Policy::kElementsPerAccess + }; + + pointer_ += layout_({ + lane_offset.row() * Shape::kRow, + lane_offset.column() * Policy::kElementsPerAccess / int(AccessType::kElements) + }); + } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + TileIteratorSimtCanonical( + TensorRef const &ref, + TensorCoord const &extent, + unsigned lane_id + ): + pointer_(reinterpret_cast(ref.data())), + layout_(ref.stride()[0] / AccessType::kElements), + divisible_(false), + extent_(extent) { + + auto lane_layout = Policy::MmaSimtPolicy::get_lane_layout(); + MatrixCoord lane_offset = lane_layout.inverse(lane_id); + + thread_offset_ = { + lane_offset.row() * Shape::kRow, + lane_offset.column() * Policy::kElementsPerAccess + }; + + pointer_ += layout_({ + lane_offset.row() * Shape::kRow, + lane_offset.column() * Policy::kElementsPerAccess / int(AccessType::kElements) + }); + } + + /// Adds a pointer offset + MCTLASS_HOST_DEVICE + TileIteratorSimtCanonical & add_pointer_offset(Index pointer_offset) { + pointer_ += pointer_offset / AccessType::kElements; + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorSimtCanonical & add_tile_offset(TensorCoord const &tile_offset) { + + MatrixCoord coord_offset( + tile_offset.row(), + tile_offset.column() * Shape::kColumn + ); + + thread_offset_ += coord_offset; + + pointer_ += layout_({ + coord_offset.row(), + coord_offset.column() + }); + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorSimtCanonical & operator+=(TensorCoord const &tile_offset) { + + add_tile_offset(tile_offset); + + return *this; + } + + /// Store + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + // de-vectorized stores + using ScalarAccessType = AlignedArray; + ScalarAccessType const *scalarFragPtr = reinterpret_cast(&frag); + ScalarAccessType *scalarPointer = reinterpret_cast(pointer_) + pointer_offset; + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::kAccessesPerIteration; ++n) { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::kElementsPerAccess; s++) { + + int ptr_idx = n * Policy::MmaSimtPolicy::WarpShape::kColumn * Policy::kElementsPerAccess + s; + int frag_idx = n * Policy::kElementsPerAccess + s; + + int col = thread_offset_.column() + ptr_idx; + + if (divisible_ || (thread_offset_.row() < extent_.row() && col < extent_.column())) { + scalarPointer[ptr_idx] = scalarFragPtr[frag_idx]; + } + } + } + } + + /// Store + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } + + /// Load + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + // de-vectorized loads + using ScalarAccessType = AlignedArray; + ScalarAccessType *scalarFragPtr = reinterpret_cast(&frag); + ScalarAccessType const *scalarPointer = reinterpret_cast(pointer_) + pointer_offset; + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::kAccessesPerIteration; ++n) { + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::kElementsPerAccess; s++) { + + int ptr_idx = n * Policy::MmaSimtPolicy::WarpShape::kColumn * Policy::kElementsPerAccess + s; + int frag_idx = n * Policy::kElementsPerAccess + s; + + int col = thread_offset_.column() + ptr_idx; + + if (divisible_ || (thread_offset_.row() < extent_.row() && col < extent_.column())) { + scalarFragPtr[frag_idx] = scalarPointer[ptr_idx]; + } + } + } + } + + /// Load + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + MCTLASS_HOST_DEVICE + TileIteratorSimtCanonical & operator++() { + return add_tile_offset({1, 0}); + } + + /// Set smem base address + MCTLASS_HOST_DEVICE + void set_smem_base_address(Index address) { + } +}; + + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_tensor_op.h new file mode 100644 index 0000000..e12a708 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_tensor_op.h @@ -0,0 +1,761 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" + +#include "mctlass/epilogue/warp/tensor_op_policy.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename OperatorShape, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename Element, ///< data type of element to be written + typename Layout ///< target shared memory layout +> +class TileIteratorTensorOp; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template < + typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape) + typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename Element_ ///< data type of element to be written +> +class TileIteratorTensorOp { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using Element = Element_; + using Layout = layout::RowMajor; + + using TensorLayout = Layout; + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + using Policy = TensorOpPolicy; + + /// Shape of the tile in memory + using Shape = MatrixShape< + Policy::kRowsPerIteration, + WarpShape::kN + >; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array< + Element, + Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>; + + /// This is the complete warp-level accumulator tile. + //using AccumulatorTile = typename Operator::FragmentC; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + + /// Number of times this iterator can be incremented + using TileIterations = typename Policy::TileIterations; + + // Internal constants + struct Detail { + static int const kLanesInQuad = 4; + }; + + /// Padding quantity + using Padding = MatrixShape< + 0, + Detail::kLanesInQuad * Policy::kElementsPerAccess>; + +private: + + /// Storage type for accessing memory + using AccessType = AlignedArray; + + // + // Data members + // + + /// Internal pointer to memory + AccessType *pointer_; + + /// Internal layout object + Layout layout_; + + /// Thread offset + MatrixCoord thread_offset_; + +public: + + /// Default constructor + MCTLASS_HOST_DEVICE + TileIteratorTensorOp(): pointer_(nullptr) { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + TileIteratorTensorOp( + TensorRef const &ref, + unsigned lane_id + ): + pointer_(reinterpret_cast(ref.data())), + layout_(ref.stride()[0] / Policy::kElementsPerAccess) { + + if ((platform::is_same::value && (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 16))|| platform::is_same::value == true) { + + //Only work correctly with m16n8k16_f16 && m16n8k32_int32 + int t_id = (((lane_id >> 4) << 5) + (lane_id & 0x7)) >> 1; + + int quad_id = (t_id / Detail::kLanesInQuad) - (lane_id >> 5) * 6; + if (platform::is_same::value == true && WarpShape::kM == 64) { + quad_id = ((lane_id >> 5) + (((lane_id >> 4) & 0x1) << 1)) << 1; + } + int lane_in_quad = (t_id % Detail::kLanesInQuad); + + pointer_ += layout_({quad_id, + lane_in_quad}); + + } + else { + int lane_in_quad0 = lane_id & 0x7; + int quad = (lane_id >> 4) & 0x3; + lane_id = lane_in_quad0 + quad * 8; + + int quad_id = (lane_id / Detail::kLanesInQuad); + int lane_in_quad = (lane_id % Detail::kLanesInQuad); + + // thread_offset_ = { + // quad_id, lane_in_quad * Policy::kElementsPerAccess + // }; + + // pointer_ += layout_({thread_offset_.row(), thread_offset_.column() / Policy::kElementsPerAccess}); + } + } + + /// Adds a pointer offset + MCTLASS_HOST_DEVICE + TileIteratorTensorOp & add_pointer_offset(Index pointer_offset) { + pointer_ += pointer_offset / Policy::kElementsPerAccess; + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorTensorOp & add_tile_offset(TensorCoord const &tile_offset) { + + MatrixCoord coord_offset( + tile_offset.row() * Shape::kRow, + tile_offset.column() * Shape::kColumn + ); + + thread_offset_ += coord_offset; + + pointer_ += layout_({ + coord_offset.row(), + coord_offset.column() / Policy::kElementsPerAccess + }); + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorTensorOp & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + /// Store + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { +#if defined(__MACA_ARCH__) + AccessType const *frag_ptr = reinterpret_cast(&frag); + + if (((platform::is_same::value) && (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 16)) || platform::is_same::value == true) { + + //Only work correctly with m16n8k16_f16 && m16n8k32_int32 + int lane_id = __lane_id(); + int idx_y = lane_id % 2; + int offset = layout_({1, 0}); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) { + pointer_[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess][idx_y] = frag_ptr[n][0]; + pointer_[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess + offset][idx_y] = frag_ptr[n][1]; + } + } + else if (platform::is_same>::value || platform::is_same::value) { + int lane_id = __lane_id(); + int t_id = (lane_id / 16) * 8 + (lane_id % 8); + int idx_x = t_id / 2; + int idx_y = t_id % 2; + int offset = layout_({1, 0}); + + int offset0 = (idx_x / 4) * offset + ((lane_id >> 4) & 0x01) * 3 * offset + (idx_x % 4); + int offset1 = offset0 + offset; + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) { + pointer_[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess + offset0].data()[idx_y] = frag_ptr[n].data()[0]; + pointer_[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess + offset1].data()[idx_y] = frag_ptr[n].data()[1]; + + } + } else if (platform::is_same::value && (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 8)) { + int t_laneid = threadIdx.x % 64; + int t_g = (t_laneid / 16); + int t_row0 = (t_g % 2) * 4 + (t_g / 2) * 2; + int t_row1 = t_row0 + 1; + + int idx0_y = t_laneid % 2; + int offset = layout_({1, 0}); + int offset0 = t_row0 * offset + (t_laneid % 8) / 2; + + + int idx1_y = idx0_y; + int offset1 = t_row1 * offset + (t_laneid % 8) / 2; + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) { + pointer_[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess + offset0].data()[idx0_y] = frag_ptr[n].data()[0]; + pointer_[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess + offset1].data()[idx1_y] = frag_ptr[n].data()[1]; + } + } + else { + + int t_laneid = threadIdx.x % 64; + int t_id = (t_laneid / 16) * 8 + (t_laneid % 8); + + int idx0_x = t_id / 2; + int idx0_y = t_id % 2; + int offset = layout_({1, 0}); + int offset0 = (idx0_x / 4) * offset + (idx0_x % 4); + + int idx1_x = (t_id / 2) + 16; + int idx1_y = t_id % 2; + int offset1 = (idx1_x / 4) * offset + (idx1_x % 4); + + // MCTLASS_PRAGMA_UNROLL + // for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) { + // pointer_[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess] = frag_ptr[n]; + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) { + pointer_[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess + offset0].data()[idx0_y] = frag_ptr[n].data()[0]; + pointer_[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess + offset1].data()[idx1_y] = frag_ptr[n].data()[1]; + } + } +#endif + } + + /// Store + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } + + /// Load + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) { + frag_ptr[n] = pointer_[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess]; + } + } + + /// Load + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + MCTLASS_HOST_DEVICE + TileIteratorTensorOp & operator++() { + return add_tile_offset({1, 0}); + } + + /// Set smem base address + MCTLASS_HOST_DEVICE + void set_smem_base_address(Index address) { + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template < + typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape) + typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename Element_, ///< data type of element to be written + int InterleavedK ///< number of interleaved k +> +class TileIteratorTensorOp > { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using Element = Element_; + using Layout = layout::ColumnMajorInterleaved; + using TensorLayout = Layout; ///< shared memory tensor ref layout + + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + using Policy = TensorOpPolicy; + + /// Shape of the tile in memory + using Shape = MatrixShape< +// Policy::kRowsPerIteration, + WarpShape::kM, + InterleavedK + >; + + /// This is the fragment size produced by one tile + using Fragment = Array< + Element, + Policy::OperatorCount::kRow * Policy::kIterationsPerInstruction + * Policy::kElementsPerIteration>; + + /// This is the fragment size produced by one iteration +// using Fragment = Array< +// Element, Policy::kElementsPerIteration >; + + /// This is the complete warp-level accumulator tile. + //using AccumulatorTile = typename Operator::FragmentC; + + /// Number of times this iterator can be incremented + using TileIterations = typename Policy::TileIterations; + + // Internal constants + struct Detail { + static int const kLanesInQuad = 4; + }; + + /// Padding quantity + using Padding = MatrixShape< + 0, + Detail::kLanesInQuad * Policy::kElementsPerIteration>; + +private: + + /// Storage type for accessing memory + using AccessType = AlignedArray; + + // + // Data members + // + + /// Internal pointer to memory + AccessType *pointer_; + + /// Internal layout object + TensorLayout layout_; + + /// Thread offset + MatrixCoord thread_offset_; + +public: + + /// Default constructor + MCTLASS_HOST_DEVICE + TileIteratorTensorOp(): pointer_(nullptr) { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + TileIteratorTensorOp( + TensorRef const &ref, + unsigned lane_id + ): + pointer_(reinterpret_cast(ref.data())), + layout_(ref.stride()[0]) { + + int quad_id = (lane_id / Detail::kLanesInQuad); + int lane_in_quad = (lane_id % Detail::kLanesInQuad); + + thread_offset_ = { + quad_id, lane_in_quad * Policy::kElementsPerIteration + }; + + pointer_ += (layout_({thread_offset_.row(), thread_offset_.column()}) / Policy::kElementsPerAccess); + } + + /// Adds a pointer offset + MCTLASS_HOST_DEVICE + TileIteratorTensorOp & add_pointer_offset(Index pointer_offset) { + pointer_ += pointer_offset / Policy::kElementsPerAccess; + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorTensorOp & add_tile_offset(TensorCoord const &tile_offset) { + + MatrixCoord coord_offset( + tile_offset.row() * Shape::kRow, + tile_offset.column() * Shape::kColumn + ); + + thread_offset_ += coord_offset; + + pointer_ += (layout_({ + coord_offset.row(), + coord_offset.column() + }) / Policy::kElementsPerAccess); + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorTensorOp & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + /// Store + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kRow * Policy::kIterationsPerInstruction; n++ ) { + + AccessType *ptr = pointer_ + layout_({n * Policy::kRowsPerIteration, 0}) / Policy::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int a = 0; a < Policy::kAccessPerIteration; ++a) { + ptr[a + pointer_offset / Policy::kElementsPerAccess] = frag_ptr[n * Policy::kAccessPerIteration + a]; + +// printf("store thread %d, address %p, bank %ld\n", threadIdx.x, pointer_+a+n*Detail::kLanesInQuad, +// ((long long)(pointer_+a+n*Detail::kLanesInQuad)>>2)&0x1f); + } + } + } + + /// Store + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } + + /// Load + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kRow * Policy::kIterationsPerInstruction; n++ ) { + + AccessType *ptr = pointer_ + layout_({n * Policy::kRowsPerIteration, 0}) / Policy::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int a = 0; a < Policy::kAccessPerIteration; ++a) { + frag_ptr[n * Policy::kAccessPerIteration + a] = ptr[a + pointer_offset / Policy::kElementsPerAccess]; + } + } + } + + /// Load + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + MCTLASS_HOST_DEVICE + TileIteratorTensorOp & operator++() { + return add_tile_offset({0, 1}); + } + + /// Set smem base address + MCTLASS_HOST_DEVICE + void set_smem_base_address(Index address) { + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template < + typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape) + typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename Element_, ///< data type of element to be written + typename Layout_ +> +class TileIteratorTensorOpCanonical { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using Element = Element_; + using Layout = Layout_; + + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + using Policy = TensorOpPolicy; + + static int const kAccessSize = 1; + static int const kAccessCount = Policy::kElementsPerAccess / kAccessSize; + + /// Shape of the tile in memory + using Shape = MatrixShape< + Policy::kRowsPerIteration, + WarpShape::kN + >; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array< + Element, + Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>; + + /// This is the complete warp-level accumulator tile. + //using AccumulatorTile = typename Operator::FragmentC; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + + // Internal constants + struct Detail { + static int const kLanesInQuad = 4; + }; + + /// Padding quantity + using Padding = MatrixShape< + 0, + Detail::kLanesInQuad * Policy::kElementsPerAccess>; + +private: + + /// Storage type for accessing memory + using AccessType = AlignedArray; + + // + // Data members + // + + /// Internal pointer to memory + AccessType *pointer_; + + /// Internal layout object + Layout layout_; + + /// Guard to indicate whether the shape is divisible + bool divisible_; + + /// Extent of the output tensor + MatrixCoord extent_; + + /// Thread offset + MatrixCoord thread_offset_; + +public: + + /// Default constructor + MCTLASS_HOST_DEVICE + TileIteratorTensorOpCanonical(): pointer_(nullptr) { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + TileIteratorTensorOpCanonical( + TensorRef const &ref, + unsigned lane_id + ): + pointer_(reinterpret_cast(ref.data())), + layout_(ref.stride()[0]), + divisible_(true), + extent_(WarpShape::kM, WarpShape::kN) { + + int quad_id = (lane_id / Detail::kLanesInQuad); + int lane_in_quad = (lane_id % Detail::kLanesInQuad); + + thread_offset_ = { + quad_id, lane_in_quad * Policy::kElementsPerAccess + }; + + pointer_ += layout_({thread_offset_.row(), thread_offset_.column()}); + } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + TileIteratorTensorOpCanonical( + TensorRef const &ref, + TensorCoord const &extent, + unsigned lane_id + ): + pointer_(reinterpret_cast(ref.data())), + layout_(ref.stride()[0]), + divisible_(false), + extent_(extent) { + + int quad_id = (lane_id / Detail::kLanesInQuad); + int lane_in_quad = (lane_id % Detail::kLanesInQuad); + + thread_offset_ = { + quad_id, lane_in_quad * Policy::kElementsPerAccess + }; + + pointer_ += layout_({thread_offset_.row(), thread_offset_.column()}); + } + + /// Adds a pointer offset + MCTLASS_HOST_DEVICE + TileIteratorTensorOpCanonical & add_pointer_offset(Index pointer_offset) { + pointer_ += pointer_offset; + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorTensorOpCanonical & add_tile_offset(TensorCoord const &tile_offset) { + + MatrixCoord coord_offset( + tile_offset.row() * Shape::kRow, + tile_offset.column() * Shape::kColumn + ); + + thread_offset_ += coord_offset; + + pointer_ += layout_({ + coord_offset.row(), + coord_offset.column() + }); + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorTensorOpCanonical & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + /// Store + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) { + MCTLASS_PRAGMA_UNROLL + for (int a = 0; a < kAccessCount; ++a) { + + int ptr_idx = n * Detail::kLanesInQuad * kAccessCount + pointer_offset + a; + int frag_idx = n * kAccessCount + a; + + int col = thread_offset_.column() + n * Detail::kLanesInQuad * Policy::kElementsPerAccess + a; + + if (divisible_ || (thread_offset_.row() < extent_.row() && col < extent_.column())) { + pointer_[ptr_idx] = frag_ptr[frag_idx]; + } + } + } + } + + /// Store + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } + + /// Load + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) { + MCTLASS_PRAGMA_UNROLL + for (int a = 0; a < kAccessCount; ++a) { + + int ptr_idx = n * Detail::kLanesInQuad * kAccessCount + pointer_offset + a; + int frag_idx = n * kAccessCount + a; + + int col = thread_offset_.column() + n * Detail::kLanesInQuad * Policy::kElementsPerAccess + a; + + if (divisible_ || (thread_offset_.row() < extent_.row() && col < extent_.column())) { + frag_ptr[frag_idx] = pointer_[ptr_idx]; + } + } + } + } + + /// Load + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + MCTLASS_HOST_DEVICE + TileIteratorTensorOpCanonical & operator++() { + return add_tile_offset({1, 0}); + } + + /// Set smem base address + MCTLASS_HOST_DEVICE + void set_smem_base_address(Index address) { + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_tensor_op_mixed.h b/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_tensor_op_mixed.h new file mode 100644 index 0000000..6942437 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_tensor_op_mixed.h @@ -0,0 +1,817 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/epilogue/warp/tensor_op_policy.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// This is an optimization available on CUDA 11.2 and beyond that eliminates branches in the epilogue. +//#define MCTLASS_EPILOGUE_WARP_TILE_ITERATOR_TENSOR_OP_MIXED_OPTIMIZATION_ENABLED ((__CUDACC_VER_MAJOR__ * 10 + __CUDACC_VER_MINOR__) >= 112) +#define MCTLASS_EPILOGUE_WARP_TILE_ITERATOR_TENSOR_OP_MIXED_OPTIMIZATION_ENABLED 0 + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Template for reading and writing tiles of accumulators to shared memory. This is optimized +/// for mixed-precision epilogues in which the accumulators are 32b in width, but the output +/// data type is smaller. +template < + typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape) + typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename Element_, ///< data type of accumulator element + int ElementSizeBits, ///< Size of accumulator element in bits + int OutputSizeBits, ///< Size of output element in bits + int OutputElementCount, ///< number of elements in output vector + int ContiguousLanes, ///< Number of consecutive lanes writing to contiguous memory + bool EightBitsOutputOrLess = (OutputSizeBits <= 8) +> +class TileIteratorTensorOpMixed { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kOutputElementCount = OutputElementCount; + + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + using Policy = TensorOpPolicy; + + /// Shape of the tile in memory + using Shape = MatrixShape< + Policy::kRowsPerIteration, + WarpShape::kN + >; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array< + Element, + Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>; + + /// This is the complete warp-level accumulator tile. + //using AccumulatorTile = typename Operator::FragmentC; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + + // Internal constants + struct Detail { + static int const kLanesInQuad = 4; + + /// Number of pointers needed to write accumulators + static int const kPointerCount = + (OutputElementCount * sizeof_bits::value) / (const_min(128, OutputElementCount * sizeof_bits::value)); + + static_assert(kPointerCount <= 4, "Can only accommodate four pointers at present."); + static_assert(sizeof(Element) == 4, "This can only be used with 32b accumulator data types (f32, s32)."); + }; + + /// Padding quantity + using Padding = MatrixShape< + 0, + Detail::kLanesInQuad * Policy::kElementsPerAccess>; + +private: + + /// Storage type for accessing memory + using AccessType = AlignedArray; + + // + // Data members + // + + /// Internal pointer to memory + AccessType *pointers_[Detail::kPointerCount]; + + /// Stride in units of AccessType + int stride_; + + /// Logical column in which warp tile is aligned + int warp_column_; + +public: + + /// Default constructor + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed() { + MCTLASS_PRAGMA_UNROLL + for (int64_t i = 0; i < Detail::kPointerCount; ++i) { + pointers_[i] = nullptr; + } + } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed( + TensorRef const &ref, + unsigned lane_id + ): + stride_(ref.stride()[0] / Policy::kElementsPerAccess), + warp_column_(0) { + + int quad_id = (lane_id / Detail::kLanesInQuad); + int lane_in_quad = (lane_id % Detail::kLanesInQuad); + + if (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 16 && platform::is_same::value == true) { + quad_id = ((lane_id >> 4) << 2) + ((lane_id >> 5) << 3); + lane_in_quad = lane_id & 0x7; + if (WarpShape::kM == 64 && WarpShape::kN == 64 && WarpShape::kK == 32) { //Only for Epilgoue_with_broadcast + quad_id = (lane_id / Detail::kLanesInQuad); + lane_in_quad = (lane_id % Detail::kLanesInQuad); + } + } + else if (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK ==8 && platform::is_same::value == true) { + quad_id = ((lane_id >> 5) + (((lane_id >> 4) & 0x1) << 1)) << 1; + lane_in_quad = lane_id & 0x7; + } + MCTLASS_PRAGMA_UNROLL + for (int64_t i = 0; i < Detail::kPointerCount; ++i) { + AccessType *ptr = reinterpret_cast(ref.data()) + quad_id * stride_; + int column_idx = (lane_in_quad % 2) + (((lane_in_quad / 2) + i) % Detail::kPointerCount) * 2; + + ptr += column_idx; + + if ((OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 16 && platform::is_same::value == true) || + (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 8 && platform::is_same::value == true)) { + int column_idx = (lane_in_quad ^ (4 * (i % Detail::kPointerCount))); + Element *element_ptr = ref.data() + quad_id * stride_ * Policy::kElementsPerAccess + column_idx; + ptr = reinterpret_cast(element_ptr); + } + + if (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 16 && platform::is_same::value == true) { + if (WarpShape::kM == 64 && WarpShape::kN == 64 && WarpShape::kK == 32) {//Only for Epilgoue_with_broadcast + ptr = reinterpret_cast(ref.data()) + quad_id * stride_; + int column_idx = (lane_in_quad % 2) + (((lane_in_quad / 2) + i) % Detail::kPointerCount) * 2; + ptr += column_idx; + } + } + + if (i == 0) { + pointers_[0 % Detail::kPointerCount] = ptr; + } + else if (i == 1) { + pointers_[1 % Detail::kPointerCount] = ptr; + } + else if (i == 2) { + pointers_[2 % Detail::kPointerCount] = ptr; + } + else if (i == 3) { + pointers_[3 % Detail::kPointerCount] = ptr; + } + } + } + + /// Adds a pointer offset + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed & add_pointer_offset(Index pointer_offset) { + + int offsetAccess = Policy::kElementsPerAccess; + if (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 16 && platform::is_same::value == true) { + offsetAccess = 16; + if (WarpShape::kM == 64 && WarpShape::kN == 64 && WarpShape::kK == 32) { + offsetAccess = Policy::kElementsPerAccess; + } + } + MCTLASS_PRAGMA_UNROLL + for (int64_t i = 0; i < Detail::kPointerCount; ++i) { + pointers_[i] += pointer_offset / offsetAccess; + } + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed & add_tile_offset(TensorCoord const &tile_offset) { + + MCTLASS_PRAGMA_UNROLL + for (int64_t i = 0; i < Detail::kPointerCount; ++i) { + pointers_[i] += tile_offset.row() * Shape::kRow * stride_ + + tile_offset.column() * Shape::kColumn / Policy::kElementsPerAccess; + } + + warp_column_ += tile_offset.column() * Shape::kColumn; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed & operator+=(TensorCoord const &tile_offset) { + return add_tile_offset(tile_offset); + } + + /// Store + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + AccessType *ptr = pointers_[0]; + +#if MCTLASS_EPILOGUE_WARP_TILE_ITERATOR_TENSOR_OP_MIXED_OPTIMIZATION_ENABLED + + // When the optimization is enabled, small tiles require separate logic. + bool kN32_optimization = (WarpShape::kN * Detail::kLanesInQuad * Policy::kElementsPerAccess * sizeof_bits::value) % 1024 == 0; + if (kN32_optimization) { + int ptr_idx = ((warp_column_ * sizeof_bits::value) / 1024) % Detail::kPointerCount; + if (ptr_idx == 0) { + ptr = pointers_[0]; + } else if (ptr_idx == 1) { + ptr = pointers_[1]; + } else if (ptr_idx == 2) { + ptr = pointers_[2]; + } else if (ptr_idx == 3) { + ptr = pointers_[3]; + } + } + +#endif + + MCTLASS_PRAGMA_UNROLL + for (int64_t n = 0; n < Policy::OperatorCount::kColumn; ++n) { + +#if MCTLASS_EPILOGUE_WARP_TILE_ITERATOR_TENSOR_OP_MIXED_OPTIMIZATION_ENABLED + + // + // When the optimization is enabled, this expression suffices to obtain the SMEM pointer. + // + if (WarpShape::kN == 64) { + ptr = pointers_[n / 4]; + } + else if (!kN32_optimization) +#endif + { + // This is the reference implementation + int column_idx = warp_column_ + n * Detail::kLanesInQuad * Policy::kElementsPerAccess; + int ptr_idx = ((column_idx * sizeof_bits::value) / 1024) % Detail::kPointerCount; + + if (ptr_idx == 0) { + ptr = pointers_[0 % Detail::kPointerCount]; + } + else if (ptr_idx == 1) { + ptr = pointers_[1 % Detail::kPointerCount]; + } + else if (ptr_idx == 2) { + ptr = pointers_[2 % Detail::kPointerCount]; + } + else if (ptr_idx == 3) { + ptr = pointers_[3 % Detail::kPointerCount]; + } + } + + int offset = n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess; + if (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 8 && platform::is_same::value == true) { + ptr[offset][0] = frag_ptr[n][0]; + ptr[offset + stride_][0] = frag_ptr[n][1]; + } + else if (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 16 && platform::is_same::value == true) { + if (WarpShape::kM == 64 && WarpShape::kN == 64 && WarpShape::kK == 32) { + ptr[offset] = frag_ptr[n]; + } + else { + ptr[offset][0] = frag_ptr[n][0]; + ptr[offset + stride_][0] = frag_ptr[n][1]; + } + } + else { + ptr[offset] = frag_ptr[n]; + } + } + } + + /// Store + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + if (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 16 && platform::is_same::value == true) { + if (WarpShape::kM == 64 && WarpShape::kN == 64 && WarpShape::kK == 32) { + if(__lane_id()>=32)return; + } + } + store_with_pointer_offset(frag, 0); + } + + /// Load + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int64_t n = 0; n < Policy::OperatorCount::kColumn; ++n) { + + int column_idx = warp_column_ + n * Detail::kLanesInQuad * Policy::kElementsPerAccess; + int ptr_idx = ((column_idx * sizeof_bits::value) / 1024) % Detail::kPointerCount; + + AccessType const *smem_ptr = pointers_[ptr_idx]; + frag_ptr[n] = smem_ptr[n * Detail::kLanesInQuad + pointer_offset / Policy::kElementsPerAccess]; + } + } + + /// Load + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Set smem base address + MCTLASS_HOST_DEVICE + void set_smem_base_address(Index address) { + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for int32_t x 16 => int8_t/int4b_t x 16 +template < + typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape) + typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape), + int OutputSizeBits ///< Size of output element in bits +> +class TileIteratorTensorOpMixed { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using Element = int32_t; + using Layout = layout::RowMajor; + static int const kOutputElementCount = 16; + + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + using Policy = TensorOpPolicy; + + /// Shape of the tile in memory + using Shape = MatrixShape< + Policy::kRowsPerIteration, + WarpShape::kN + >; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array< + Element, + Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>; + + /// This is the complete warp-level accumulator tile. + //using AccumulatorTile = typename Operator::FragmentC; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + + // Internal constants + struct Detail { + static int const kLanesInQuad = 4; + + /// Number of pointers needed to write accumulators + static int const kPointerCount = 2; + + /// Offsets added + static int const kOffsetCount = 4; + + static_assert(sizeof(Element) == 4, "This can only be used with 32b accumulator data types (f32, s32)."); + }; + + /// Padding quantity + using Padding = MatrixShape<0, Detail::kLanesInQuad * 2>; + +private: + + /// Storage type for accessing memory + using AccessType = AlignedArray; + + // + // Data members + // + + /// Internal pointer to memory + AccessType *pointers_[Detail::kPointerCount]; + + /// Stride in units of AccessType + int stride_; + + /// Uniform offset in bytes added to warp tile iterator + int uniform_offset_[Detail::kOffsetCount]; + +public: + + /// Default constructor + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed() { + MCTLASS_PRAGMA_UNROLL + for (int64_t i = 0; i < Detail::kPointerCount; ++i) { + pointers_[i] = nullptr; + } + } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed( + TensorRef const &ref, + unsigned lane_id + ): + stride_(ref.stride()[0] / AccessType::kElements) { + + int quad_id = (lane_id / Detail::kLanesInQuad); + int lane_in_quad = (lane_id % Detail::kLanesInQuad); + + if (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 32) { + quad_id = ((lane_id >> 4) << 2) - (lane_id >> 5) * 6; + lane_in_quad = (lane_id & 0x7); + } + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Detail::kPointerCount; ++i) { + AccessType *ptr = reinterpret_cast(ref.data()) + quad_id * stride_; + int column_idx = lane_in_quad ^ (i * 2); + + ptr += column_idx; + + if (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 32) { + column_idx = lane_in_quad ^ (i * 4); + Element *element_ptr = ref.data() + quad_id * stride_ * Policy::kElementsPerAccess + column_idx; + ptr = reinterpret_cast(element_ptr); + } + + if (i == 0) { + pointers_[0] = ptr; + } + else if (i == 1) { + pointers_[1] = ptr; + } + } + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Detail::kOffsetCount; ++i) { + uniform_offset_[i] = (i ^ 0) * 4 * sizeof(AccessType); + } + } + + /// Adds a pointer offset + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed & add_pointer_offset(Index pointer_offset) { + + MCTLASS_PRAGMA_UNROLL + for (int64_t i = 0; i < Detail::kPointerCount; ++i) { + pointers_[i] += pointer_offset / AccessType::kElements; + } + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed & add_tile_offset(TensorCoord const &tile_offset) { + + int ptr_offset = tile_offset.row() * Shape::kRow * stride_ + + tile_offset.column() * Shape::kColumn / AccessType::kElements; + + pointers_[0] += ptr_offset; + pointers_[1] += ptr_offset; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Detail::kOffsetCount; ++i) { + uniform_offset_[i] = (i ^ tile_offset.column()) * 4 * sizeof(AccessType); + } + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed & operator+=(TensorCoord const &tile_offset) { + return add_tile_offset(tile_offset); + } + + /// Store + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) { + + int ptr_idx = (n / 4); + int offset_idx = (n % 4); + + AccessType *ptr; + if (ptr_idx == 0) { + ptr = pointers_[0]; + } + else if (ptr_idx == 1) { + ptr = pointers_[1]; + } + + int offset = (n / 4) * 16 + pointer_offset / AccessType::kElements; + +#if 0 + // + // Using inline PTX to avoid generic memory + // + uint32_t smem_addr = arch::mctlass_get_smem_pointer(ptr); + uint32_t const *data = reinterpret_cast(frag_ptr + n); + uint32_t offset_in_bytes = offset * sizeof(AccessType) + uniform_offset_[offset_idx]; + + asm volatile( + "{ .reg .u32 smem_ptr; add.u32 smem_ptr, %0, %1; st.shared.v2.u32 [smem_ptr], {%2, %3}; }\n" + : : "r"(smem_addr), "r"(offset_in_bytes), "r"(data[0]), "r"(data[1]) + ); +#else + AccessType *smem_ptr = pointers_[ptr_idx]; + // + //When inline PTX is not used, offset needs uniform_offset_(not multipiled by sizeof(AccessType)) + // + int offset_without_bytes = offset + uniform_offset_[offset_idx] / sizeof(AccessType); + if (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 32) { + smem_ptr[offset_without_bytes][0] = frag_ptr[n][0]; + smem_ptr[offset_without_bytes + stride_][0] = frag_ptr[n][1]; + } + else { + smem_ptr[offset_without_bytes] = frag_ptr[n]; + } + +#endif + } + } + + /// Store + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } + + /// Set smem base address + MCTLASS_HOST_DEVICE + void set_smem_base_address(Index address) { + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for int32_t x 8 => int8_t/int4b_t x 8 +template < + typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape) + typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape) + int OutputSizeBits ///< Size of output element in bits +> +class TileIteratorTensorOpMixed { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using Element = int32_t; + using Layout = layout::RowMajor; + static int const kOutputElementCount = 8; + + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + using Policy = TensorOpPolicy; + + /// Shape of the tile in memory + using Shape = MatrixShape< + Policy::kRowsPerIteration, + WarpShape::kN + >; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array< + Element, + Policy::OperatorCount::kColumn * Policy::kElementsPerAccess>; + + /// This is the complete warp-level accumulator tile. + //using AccumulatorTile = typename Operator::FragmentC; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + + // Internal constants + struct Detail { + static int const kLanesInQuad = 4; + + /// Number of pointers needed to write accumulators + static int const kPointerCount = 2; + + static_assert(sizeof(Element) == 4, "This can only be used with 32b accumulator data types (f32, s32)."); + }; + + /// Padding quantity + using Padding = MatrixShape<0, Detail::kLanesInQuad * 2>; + +private: + + /// Storage type for accessing memory + using AccessType = AlignedArray; + + // + // Data members + // + + /// Internal pointer to memory + AccessType *pointers_[Detail::kPointerCount]; + + /// Stride in units of AccessType + int stride_; + +public: + + /// Default constructor + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed() { + MCTLASS_PRAGMA_UNROLL + for (int64_t i = 0; i < Detail::kPointerCount; ++i) { + pointers_[i] = nullptr; + } + } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed( + TensorRef const &ref, + unsigned lane_id + ): + stride_(ref.stride()[0] / AccessType::kElements) { + + int quad_id = (lane_id / Detail::kLanesInQuad); + int lane_in_quad = (lane_id % Detail::kLanesInQuad); + + if (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 32) { + quad_id = ((lane_id >> 4) << 2) - (lane_id >> 5) * 6; + lane_in_quad = (lane_id & 0x7); + } + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Detail::kPointerCount; ++i) { + AccessType *ptr = reinterpret_cast(ref.data()) + quad_id * stride_; + int column_idx = lane_in_quad ^ (i * 2); + + ptr += column_idx; + + if (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 32) { + column_idx = lane_in_quad ^ (i * 4); + Element *element_ptr = ref.data() + quad_id * stride_ * Policy::kElementsPerAccess + column_idx; + ptr = reinterpret_cast(element_ptr); + } + if (i == 0) { + pointers_[0] = ptr; + } + else if (i == 1) { + pointers_[1] = ptr; + } + } + } + + /// Adds a pointer offset + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed & add_pointer_offset(Index pointer_offset) { + + MCTLASS_PRAGMA_UNROLL + for (int64_t i = 0; i < Detail::kPointerCount; ++i) { + pointers_[i] += pointer_offset / AccessType::kElements; + } + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed & add_tile_offset(TensorCoord const &tile_offset) { + + int ptr_offset = tile_offset.row() * Shape::kRow * stride_ + + tile_offset.column() * Shape::kColumn / AccessType::kElements; + + pointers_[0] += ptr_offset; + pointers_[1] += ptr_offset; + + if (tile_offset.column() % 2) { + auto tmp = pointers_[0]; + pointers_[0] = pointers_[1]; + pointers_[1] = tmp; + } + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorTensorOpMixed & operator+=(TensorCoord const &tile_offset) { + return add_tile_offset(tile_offset); + } + + /// Store + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::OperatorCount::kColumn; ++n) { + + int ptr_idx = (n / 4); + + AccessType *ptr; + if (ptr_idx == 0) { + ptr = pointers_[0]; + } + else if (ptr_idx == 1) { + ptr = pointers_[1]; + } + + int offset = (n / 4) * 16 + pointer_offset / AccessType::kElements + (n % 4) * 4; + +#if 0 + // + // Using inline PTX to avoid generic memory + // + uint32_t smem_addr = arch::mctlass_get_smem_pointer(ptr); + uint32_t const *data = reinterpret_cast(frag_ptr + n); + uint32_t offset_in_bytes = offset * sizeof(AccessType); + asm volatile( + "{ .reg .u32 smem_ptr; add.u32 smem_ptr, %0, %1; st.shared.v2.u32 [smem_ptr], {%2, %3}; }\n" + : : "r"(smem_addr), "r"(offset_in_bytes), "r"(data[0]), "r"(data[1]) + ); +#else + AccessType *smem_ptr = pointers_[ptr_idx]; + if (OperatorShape::kM == 16 && OperatorShape::kN == 8 && OperatorShape::kK == 32) { + smem_ptr[offset][0] = frag_ptr[n][0]; + smem_ptr[offset + stride_][0] = frag_ptr[n][1]; + } + else { + smem_ptr[offset] = frag_ptr[n]; + } +#endif + } + } + + /// Store + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } + + /// Set smem base address + MCTLASS_HOST_DEVICE + void set_smem_base_address(Index address) { + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#undef MCTLASS_EPILOGUE_WARP_TILE_ITERATOR_TENSOR_OP_MIXED_OPTIMIZATION_ENABLED + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_volta_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_volta_tensor_op.h new file mode 100644 index 0000000..dcdd994 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_volta_tensor_op.h @@ -0,0 +1,440 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" + +#include "mctlass/epilogue/warp/tensor_op_policy.h" +#include "mctlass/epilogue/warp/volta_tensor_op_policy.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename InterleavedTileShape, ///< shape of indivisible instruction-level arrangement (concept: GemmShape) + typename ElementC, ///< Accumulator layout + typename Layout ///< target shared memory layout +> +struct TileIteratorVoltaTensorOp; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template < + typename WarpShape_ ///< shape of warp-level GEMM (concept: MatrixShape) +> +struct TileIteratorVoltaTensorOp, half_t, layout::RowMajor> { +public: + + using WarpShape = WarpShape_; + using InterleavedTileShape = gemm::GemmShape<32, 32, 4>; + using Element = half_t; + using Layout = layout::RowMajor; + + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + using Policy = VoltaTensorOpPolicy; + + /// Shape of the tile in memory + using Shape = MatrixShape< + Policy::kRowsPerIteration, + WarpShape::kN + >; + + /// Array type for aligned memory accesses + using AccessType = typename Policy::AccessType; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = typename Policy::Fragment; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = typename Policy::AccumulatorTile; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + + /// Number of elements per access + static int const kElementsPerAccess = Policy::kElementsPerAccess; + + // Internal constants + struct Detail { + static int const kLanesInQuad = 4; + static int const kRowsPerQuad = 4; + static int const kColumnsPerQuad = 8; + static int const kAccessesPerQuad = kColumnsPerQuad / Policy::kElementsPerAccess; + static int const kAccessQuadDelta = 16; + }; + + /// Padding quantity + using Padding = MatrixShape< + 0, + Policy::kElementsPerAccess>; + +private: + + // + // Data members + // + + /// Internal pointer to memory + AccessType *pointer_; + + /// Internal layout object + Layout layout_; + +public: + + /// Default constructor + MCTLASS_HOST_DEVICE + TileIteratorVoltaTensorOp(): pointer_(nullptr) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + TileIteratorVoltaTensorOp( + TensorRef const &ref, + unsigned lane_id + ): + pointer_(reinterpret_cast(ref.data())), + layout_(ref.stride()[0] / Policy::kElementsPerAccess) { + + int quad_id = lane_id / Detail::kLanesInQuad; + int lane_in_quad = (lane_id % Detail::kLanesInQuad); + + int quad_row_idx = ((quad_id & 4) >> 1) + (quad_id & 1); + int quad_col_idx = ((quad_id & 2) >> 1); + + int row = quad_row_idx * Detail::kRowsPerQuad + lane_in_quad; + int column = quad_col_idx * Detail::kColumnsPerQuad; + + pointer_ += layout_({row, column / kElementsPerAccess}); + } + + /// Adds a pointer offset + MCTLASS_HOST_DEVICE + TileIteratorVoltaTensorOp & add_pointer_offset(Index pointer_offset) { + pointer_ += pointer_offset / Policy::kElementsPerAccess; + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorVoltaTensorOp & add_tile_offset(TensorCoord const &tile_offset) { + + pointer_ += layout_({ + tile_offset.row() * Shape::kRow, + tile_offset.column() * Shape::kColumn / Policy::kElementsPerAccess}); + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorVoltaTensorOp & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + /// Store + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int tile_idx = 0; tile_idx < Policy::TileIterations::kColumn; ++tile_idx) { + + MCTLASS_PRAGMA_UNROLL + for (int access_idx = 0; access_idx < Policy::kAccessesPerInterleavedTile; ++access_idx) { + + int access_quad = access_idx / 2; + int access = access_idx % 2; + + int ptr_offset = tile_idx * InterleavedTileShape::kN / Policy::kElementsPerAccess + + access_quad * Detail::kAccessQuadDelta / Policy::kElementsPerAccess + + access + pointer_offset / Policy::kElementsPerAccess; + + int frag_idx = tile_idx * Policy::kAccessesPerInterleavedTile + access_idx; + + AccessType access_vector = frag_ptr[frag_idx]; + + pointer_[ptr_offset] = access_vector; + } + } + } + + /// Store + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } + + /// Load + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int tile_idx = 0; tile_idx < Policy::TileIterations::kColumn; ++tile_idx) { + + MCTLASS_PRAGMA_UNROLL + for (int access_idx = 0; access_idx < Policy::kAccessesPerInterleavedTile; ++access_idx) { + + int access_quad = access_idx / 2; + int access = access_idx % 2; + + int ptr_offset = tile_idx * Detail::kTileDelta + access_quad * Detail::kAccessQuadDelta + + access + pointer_offset / Policy::kElementsPerAccess; + + int frag_idx = tile_idx * Policy::kAccessesPerInterleavedTile + access_idx; + + frag_ptr[frag_idx] = pointer_[ptr_offset]; + } + } + } + + /// Load + MCTLASS_HOST_DEVICE + void load(Fragment const &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Set smem base address + MCTLASS_HOST_DEVICE + void set_smem_base_address(Index address) { + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template < + typename WarpShape_ ///< shape of warp-level GEMM (concept: MatrixShape) +> +struct TileIteratorVoltaTensorOp, float, layout::RowMajor> { +public: + + using WarpShape = WarpShape_; + using InterleavedTileShape = gemm::GemmShape<32, 32, 4>; + using Element = float; + using Layout = layout::RowMajor; + + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + using Policy = VoltaTensorOpPolicy; + + /// Shape of the tile in memory + using Shape = MatrixShape< + Policy::kRowsPerIteration, + WarpShape::kN + >; + + /// Array type for aligned memory accesses + using AccessType = typename Policy::AccessType; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = typename Policy::Fragment; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = typename Policy::AccumulatorTile; + + /// Number of times this iterator can be incremented + static int const kIterations = Policy::kIterations; + + /// Number of elements per access + static int const kElementsPerAccess = Policy::kElementsPerAccess; + + // Internal constants + struct Detail { + static int const kLanesInQuad = 4; + static int const kRowsPerQuad = 4; + static int const kColumnsPerQuad = 8; + static int const kAccessesPerQuad = kColumnsPerQuad / Policy::kElementsPerAccess; + static int const kAccessQuadDelta = 16; + }; + + /// Padding quantity + using Padding = MatrixShape< + 0, + Policy::kElementsPerAccess>; + +private: + + // + // Data members + // + + /// Internal pointer to memory + AccessType *pointer_; + + /// Internal layout object + Layout layout_; + +public: + + /// Default constructor + MCTLASS_HOST_DEVICE + TileIteratorVoltaTensorOp(): pointer_(nullptr) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + TileIteratorVoltaTensorOp( + TensorRef const &ref, + unsigned lane_id + ): + pointer_(reinterpret_cast(ref.data())), + layout_(ref.stride()[0] / Policy::kElementsPerAccess) { + + int quad_id = lane_id / Detail::kLanesInQuad; + int lane_in_quad = (lane_id % Detail::kLanesInQuad); + + int const kQuadRowDelta = 4; + int const kQuadColumnDelta = 2 * Policy::MmaIterations::kColumn; + + int quad_row_offset = ((quad_id & 4) / 2 + (quad_id & 1)) * kQuadRowDelta; + int quad_column_offset = (quad_id & 2) / 2 * kQuadColumnDelta; + + int thread_row_offset = (lane_in_quad & 1); + int thread_column_offset = (lane_in_quad & 2) / 2; + + int row = quad_row_offset + thread_row_offset; + int column = quad_column_offset + thread_column_offset; + + pointer_ += layout_({row, column}); + } + + /// Adds a pointer offset + MCTLASS_HOST_DEVICE + TileIteratorVoltaTensorOp & add_pointer_offset(Index pointer_offset) { + pointer_ += pointer_offset / Policy::kElementsPerAccess; + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorVoltaTensorOp & add_tile_offset(TensorCoord const &tile_offset) { + + pointer_ += layout_({ + tile_offset.row() * Shape::kRow, + tile_offset.column() * Shape::kColumn / Policy::kElementsPerAccess}); + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorVoltaTensorOp & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + /// Store + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + int const kAccessesPerRow = Policy::TileIterations::kColumn * Policy::MmaIterations::kColumn * 2; + + MCTLASS_PRAGMA_UNROLL + for (int row_idx = 0; row_idx < Policy::kRowsPerMmaTile; ++row_idx) { + + MCTLASS_PRAGMA_UNROLL + for (int access_idx = 0; access_idx < kAccessesPerRow; ++access_idx) { + + int frag_idx = row_idx * kAccessesPerRow + access_idx; + + int ptr_column_offset = (access_idx & 1) * 2 + + (access_idx & 2) * Policy::MmaIterations::kColumn * 2 + + (access_idx & 4) * Policy::MmaIterations::kColumn * 2; + + int ptr_row_offset = row_idx * 2; + + int ptr_offset = layout_({ptr_row_offset, ptr_column_offset}) + pointer_offset / Policy::kElementsPerAccess; + + pointer_[ptr_offset] = frag_ptr[frag_idx]; + } + } + } + + /// Store + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } + + /// Load + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + assert(0); + } + + /// Load + MCTLASS_HOST_DEVICE + void load(Fragment const &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Set smem base address + MCTLASS_HOST_DEVICE + void set_smem_base_address(Index address) { + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_wmma_tensor_op.h b/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_wmma_tensor_op.h new file mode 100644 index 0000000..a28ec97 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/tile_iterator_wmma_tensor_op.h @@ -0,0 +1,237 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief +*/ + +#pragma once + +//#if !(defined(__clang__) && defined(__MACA__)) +#if defined(__MACA__) + +#include "mctlass/mctlass.h" +#include "mctlass/wmma_array.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/tensor_ref.h" + +#include "mctlass/epilogue/warp/wmma_tensor_op_policy.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename OperatorShape, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename OperatorFragment, ///< wmma fragment to be written (concept: mxmaca::wmma::fragment) + typename Layout ///< target shared memory layout +> +class TileIteratorWmmaTensorOp; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template for reading and writing tiles of accumulators to shared memory +template < + typename WarpShape_, ///< shape of warp-level GEMM (concept: GemmShape) + typename OperatorShape_, ///< matrix multiply operation shape (concept: gemm::GemmShape) + typename OperatorFragment_ ///< wmma fragment to be written (concept: mxmaca::wmma::fragment) +> +class TileIteratorWmmaTensorOp { +public: + + using WarpShape = WarpShape_; + using OperatorShape = OperatorShape_; + using OperatorFragment = OperatorFragment_; + using Layout = layout::RowMajor; + + // + // Derived types + // + using WmmaDataType = typename OperatorFragment::element_type; + using Element = typename mctlass::arch::WmmaToMctlassDataType::Type; ///< Data Type of element stored in mxmaca::wmma::frament + using TensorRef = TensorRef; ///< Tensor Reference object + using TensorCoord = MatrixCoord; ///< Logical coordinate in referenced tensor + using Index = typename TensorRef::Index; + using LongIndex = typename TensorRef::LongIndex; + + using Policy = WmmaTensorOpPolicy; + + /// Shape of the tile in memory + using Shape = MatrixShape< + Policy::kRowsPerIteration, + WarpShape::kN + >; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = WmmaFragmentArray; + + + /// This is the complete warp-level accumulator tile. + //using AccumulatorTile = typename Operator::FragmentC; + + + /// Padding quantity + // (Epilogue shared memory padding for WMMA Gemm kernel is set to run optimaly on Turing) + using Padding = MatrixShape< + 0, + 4 * Policy::kElementsPerAccess + >; + +private: + + /// Storage type for accessing memory + //using AccessType = AlignedArray; + + // + // Data members + // + + /// Internal pointer to shared memory + TensorRef ref_; + + +public: + + /// Default constructor + MCTLASS_HOST_DEVICE + TileIteratorWmmaTensorOp(): ref_(nullptr) { + + } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + TileIteratorWmmaTensorOp( + TensorRef const &ref, + unsigned lane_id + ): ref_(ref) { + } + + /// Adds a pointer offset + MCTLASS_HOST_DEVICE + TileIteratorWmmaTensorOp & add_pointer_offset(Index pointer_offset) { + ref_.add_pointer_offset(pointer_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorWmmaTensorOp & add_tile_offset(TensorCoord const &tile_offset) { + ref_.add_coord_offset({tile_offset.row() * OperatorShape::kM, tile_offset.column() * WarpShape::kN}); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_HOST_DEVICE + TileIteratorWmmaTensorOp & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + /// Store + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + for(int n=0; n < Policy::OperatorCount::kColumn; n++) { + + WmmaDataType* ptr = reinterpret_cast (ref_.data() + ref_.offset({0, n * OperatorShape::kN}) + pointer_offset); + /*mxmaca::wmma::store_matrix_sync( + ptr, + frag[n], + ref_.stride()[0], + mxmaca::wmma::layout_t::mem_row_major + ); + */ + mxmaca::wmma::store_matrix_sync( + ptr, + frag[n], + ref_.stride()[0], + mxmaca::wmma::layout_t::mem_row_major + ); + } + } + + /// Store + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } + + /// Load + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + for(int n=0; n < Policy::OperatorCount::kColumn; n++) { + + WmmaDataType* ptr = reinterpret_cast (ref_.data() + ref_.offset({0, n * OperatorShape::kN}) + pointer_offset); + /* mxmaca::wmma::load_matrix_sync( + frag[n], + ptr, + ref_.stride()[0], + mxmaca::wmma::layout_t::mem_row_major + ); + */ + mxmaca::wmma::load_matrix_sync( + frag[n], + ptr, + ref_.stride()[0], + mxmaca::wmma::layout_t::mem_row_major + ); + } + } + + /// Load + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Set smem base address + MCTLASS_HOST_DEVICE + void set_smem_base_address(Index address) { + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#endif // !defined(__clang__) + diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/volta_tensor_op_policy.h b/csrc/mctlass/include/mctlass/epilogue/warp/volta_tensor_op_policy.h new file mode 100644 index 0000000..37a07c6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/volta_tensor_op_policy.h @@ -0,0 +1,195 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic structures needed for implementing the warp-scoped phase of the epilogue. + These quantities assume a 'column-major' arrangement of TensorOp instructions, of which + a row-oriented slice is visible per iteration. +*/ + +#pragma once + +#include "mctlass/matrix_shape.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/gemm/gemm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Policy details related to the epilogue +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename InterleavedTileShape, ///< shape of indivisible instruction-level arrangement (concept: GemmShape) + typename ElementC, ///< Accumulator layout + typename Layout ///< target shared memory layout +> +struct VoltaTensorOpPolicy; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for row-major +template < + typename WarpShape_ ///< shape of warp-level GEMM (concept: GemmShape) +> +struct VoltaTensorOpPolicy, half_t, layout::RowMajor> { + + using WarpShape = WarpShape_; + using InterleavedTileShape = gemm::GemmShape<32, 32, 4>; + using ElementC = half_t; + using Layout = layout::RowMajor; + + /// Shape of one warp-levelinstruction + using InstructionShape = gemm::GemmShape<16, 16, 4>; + + /// Number of mma operations performed for one 32x32x4 interleaved tile + using MmaIterations = MatrixShape< + InterleavedTileShape::kM / InstructionShape::kM, + InterleavedTileShape::kN / InstructionShape::kN + >; + + /// Number of 32x32x4 interleaved tiles performed to cover the warp-level GEMM shape + using TileIterations = MatrixShape< + WarpShape::kM / InterleavedTileShape::kM, + WarpShape::kN / InterleavedTileShape::kN + >; + + /// Number of accumulator elements owned by each thread per Mma + static int const kElementsPerMma = 8; + static int const kRowsPerIteration = 16; + + // + // Hard-coded constants regarding Tensor Operations + // + + /// Number of accumulator elements stored per memory instruction to shared memory + static int const kElementsPerAccess = 4; + + /// Number of accesses performed per interleaved tile + static int const kAccessesPerInterleavedTile = 4; + + /// Total number of iterations needed to cover the entire tile + static int const kIterations = TileIterations::kRow * 2; + + // + // Derived types + // + + /// Array type for aligned memory accesses + using AccessType = AlignedArray; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array< + ElementC, + kElementsPerAccess * kAccessesPerInterleavedTile * TileIterations::kColumn>; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = Array< + ElementC, + TileIterations::kCount * MmaIterations::kCount * kElementsPerMma>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for row-major +template < + typename WarpShape_ ///< shape of warp-level GEMM (concept: MatrixShape) +> +struct VoltaTensorOpPolicy, float, layout::RowMajor> { + + using WarpShape = WarpShape_; + using InterleavedTileShape = gemm::GemmShape<32, 32, 4>; + using ElementC = float; + using Layout = layout::RowMajor; + + /// Shape of one warp-levelinstruction + using InstructionShape = gemm::GemmShape<16, 16, 4>; + + /// Number of mma operations performed for one 32x32x4 interleaved tile + using MmaIterations = MatrixShape< + InterleavedTileShape::kM / InstructionShape::kM, + InterleavedTileShape::kN / InstructionShape::kN + >; + + /// Number of 32x32x4 interleaved tiles performed to cover the warp-level GEMM shape + using TileIterations = MatrixShape< + WarpShape::kM / InterleavedTileShape::kM, + WarpShape::kN / InterleavedTileShape::kN + >; + + /// Number of accumulator elements owned by each thread per Mma + static int const kElementsPerMma = 8; + static int const kRowsPerIteration = 16; + + // + // Hard-coded constants regarding Tensor Operations + // + + /// Number of accumulator elements stored per memory instruction to shared memory + static int const kElementsPerAccess = 2; + + /// Number of accesses performed per interleaved tile + static int const kAccessesPerInterleavedTile = 8; + + /// Number of rows per interleaved tile + static int const kRowsPerMmaTile = 2; + + /// Total number of iterations needed to cover the entire tile + static int const kIterations = TileIterations::kRow * MmaIterations::kRow; + + // + // Derived types + // + + /// Array type for aligned memory accesses + using AccessType = AlignedArray; + + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array< + ElementC, + kElementsPerAccess * kAccessesPerInterleavedTile * TileIterations::kColumn>; + + /// This is the complete warp-level accumulator tile. + using AccumulatorTile = Array< + ElementC, + TileIterations::kCount * MmaIterations::kCount * kElementsPerMma>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/epilogue/warp/wmma_tensor_op_policy.h b/csrc/mctlass/include/mctlass/epilogue/warp/wmma_tensor_op_policy.h new file mode 100644 index 0000000..5ebe1a5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/epilogue/warp/wmma_tensor_op_policy.h @@ -0,0 +1,101 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic structures needed for implementing the warp-scoped phase of the epilogue. + These quantities assume a 'column-major' arrangement of TensorOp instructions, of which + a row-oriented slice is visible per iteration. +*/ + +#pragma once + +#include "mctlass/arch/wmma.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/layout/matrix.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace epilogue { +namespace warp { + +//////////////////////////////////////////////////////////////////////////////// + +/// Policy details related to the epilogue +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename OperatorShape, ///< matrix multiply operation shape (concept: gemm:GemmShape) + typename Layout ///< target shared memory layout +> +struct WmmaTensorOpPolicy; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for row-major +template < + typename WarpShape, ///< shape of warp-level GEMM (concept: MatrixShape) + typename OperatorShape ///< matrix multiply operation shape (concept: gemm::GemmShape) +> +struct WmmaTensorOpPolicy { + + /// Number of operations + using OperatorCount = MatrixShape< + WarpShape::kM / OperatorShape::kM, + WarpShape::kN / OperatorShape::kN + >; + + // + // Hard-coded constants regarding Tensor Operations + // + static int const kElementsPerAccess = 2; + static int const kRowsPerIteration = OperatorShape::kM; + static int const kWmmaFragmentsPerAccess = 1; + + // + // Derived quantities + // + + // Number of externally visible iterations + static int const kIterations = OperatorCount::kRow; + +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace epilogue +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// + +#endif + diff --git a/csrc/mctlass/include/mctlass/fast_math.h b/csrc/mctlass/include/mctlass/fast_math.h new file mode 100644 index 0000000..985cc66 --- /dev/null +++ b/csrc/mctlass/include/mctlass/fast_math.h @@ -0,0 +1,971 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#include +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/uint128.h" +#include "mctlass/coord.h" +#include "mctlass/numeric_types.h" + +/** + * \file + * \brief Math utilities + */ + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +MCTLASS_HOST_DEVICE void swap(T &lhs, T &rhs) { + T tmp = lhs; + lhs = rhs; + rhs = tmp; +} + +/****************************************************************************** + * Static math utilities + ******************************************************************************/ + +/// Mixed precision dot product +template +MCTLASS_HOST_DEVICE LongIndex dot( + Coord const &coord, + Coord const &stride, + LongIndex acc = LongIndex()) { + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < N; ++n) { + acc += LongIndex(coord[n]) * stride[n]; + } + return acc; +} + +/** + * Statically determine if N is a power-of-two + */ +template +struct is_pow2 { + static bool const value = ((N & (N - 1)) == 0); +}; + +/** + * Statically determine log2(N), rounded down + */ +template +struct log2_down { + /// Static logarithm value + enum { value = log2_down> 1), Count + 1>::value }; +}; + +// Base case +template +struct log2_down { + enum { value = Count }; +}; + +/** + * Statically determine log2(N), rounded up + */ +template +struct log2_up { + /// Static logarithm value + enum { value = log2_up> 1), Count + 1>::value }; +}; + +// Base case +template +struct log2_up { + enum { value = ((1 << Count) < N) ? Count + 1 : Count }; +}; + +/** + * Statically estimate sqrt(N) to the nearest power-of-two + */ +template +struct sqrt_est { + enum { value = 1 << (log2_up::value / 2) }; +}; + +/** + * For performing a constant-division with a compile-time assertion that the + * Divisor evenly-divides the Dividend. + */ +template +struct divide_assert { + enum { value = Dividend / Divisor }; + + static_assert((Dividend % Divisor == 0), "Not an even multiple"); +}; + +/****************************************************************************** + * Rounding + ******************************************************************************/ + +/** + * Round dividend up to the nearest multiple of divisor + */ +template +MCTLASS_HOST_DEVICE dividend_t round_nearest(dividend_t dividend, divisor_t divisor) { + return ((dividend + divisor - 1) / divisor) * divisor; +} + +/** + * Greatest common divisor + */ +template +MCTLASS_HOST_DEVICE value_t gcd(value_t a, value_t b) { + for (;;) { + if (a == 0) return b; + b %= a; + if (b == 0) return a; + a %= b; + } +} + +/** + * Least common multiple + */ +template +MCTLASS_HOST_DEVICE value_t lcm(value_t a, value_t b) { + value_t temp = gcd(a, b); + + return temp ? (a / temp * b) : 0; +} + +/// Returns the smallest value in the half-open range [a, a+b) that is a multiple of b +MCTLASS_HOST_DEVICE +constexpr int round_up(int a, int b) { + return ((a + b - 1) / b) * b; +} + +/// Returns the ceiling of (a / b) +MCTLASS_HOST_DEVICE +constexpr int ceil_div(int a, int b) { + return (a + b - 1) / b; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/** + * log2 computation, what's the + * difference between the below codes and + * log2_up/down codes? + */ +template +MCTLASS_HOST_DEVICE value_t clz(value_t x) { + for (int i = 31; i >= 0; --i) { + if ((1 << i) & x) return 31 - i; + } + return 32; +} + +template +MCTLASS_HOST_DEVICE value_t find_log2(value_t x) { + int a = int(31 - clz(x)); + a += (x & (x - 1)) != 0; // Round up, add 1 if not a power of 2. + return a; +} + + +/** + * Find divisor, using find_log2 + */ +MCTLASS_HOST_DEVICE +void find_divisor(unsigned int& mul, unsigned int& shr, unsigned int denom) { + if (denom == 1) { + mul = 0; + shr = 0; + } else { + unsigned int p = 31 + find_log2(denom); + unsigned m = unsigned(((1ull << p) + unsigned(denom) - 1) / unsigned(denom)); + + mul = m; + shr = p - 32; + } +} + +/** + * Find quotient and remainder using device-side intrinsics + */ +MCTLASS_HOST_DEVICE +void fast_divmod(int& quo, int& rem, int src, int div, unsigned int mul, unsigned int shr) { + + #if defined(__MACA_ARCH__) + // Use IMUL.HI if div != 1, else simply copy the source. + quo = (div != 1) ? __umulhi(src, mul) >> shr : src; + #else + //quo = int((div != 1) ? int(((int64_t)src * mul) >> 32) >> shr : src); + if (div == 1) + { + quo = src; + } + else + { + int64_t tmp = ((int64_t)src * mul) >> 32; + tmp = tmp >> shr; + quo = int(tmp); + } + + #endif + // The remainder. + rem = src - (quo * div); +} + +// For long int input +MCTLASS_HOST_DEVICE +void fast_divmod(int& quo, int64_t& rem, int64_t src, int div, unsigned int mul, unsigned int shr) { + + #if defined(__MACA_ARCH__) + // Use IMUL.HI if div != 1, else simply copy the source. + quo = (div != 1) ? __umulhi(src, mul) >> shr : src; + #else + quo = int((div != 1) ? ((src * mul) >> 32) >> shr : src); + #endif + // The remainder. + rem = src - (quo * div); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Object to encapsulate the fast division+modulus operation. +/// +/// This object precomputes two values used to accelerate the computation and is best used +/// when the divisor is a grid-invariant. In this case, it may be computed in host code and +/// marshalled along other kernel arguments using the 'Params' pattern. +/// +/// Example: +/// +/// +/// int quotient, remainder, dividend, divisor; +/// +/// FastDivmod divmod(divisor); +/// +/// divmod(quotient, remainder, dividend); +/// +/// // quotient = (dividend / divisor) +/// // remainder = (dividend % divisor) +/// +struct FastDivmod { + + int divisor; + unsigned int multiplier; + unsigned int shift_right; + + /// Find quotient and remainder using device-side intrinsics + MCTLASS_HOST_DEVICE + void fast_divmod(int& quotient, int& remainder, int dividend) const { + +#if defined(__MACA_ARCH__) + // Use IMUL.HI if divisor != 1, else simply copy the source. + quotient = (divisor != 1) ? __umulhi(dividend, multiplier) >> shift_right : dividend; +#else + quotient = int((divisor != 1) ? int(((int64_t)dividend * multiplier) >> 32) >> shift_right : dividend); +#endif + + // The remainder. + remainder = dividend - (quotient * divisor); + } + + /// For long int input + MCTLASS_HOST_DEVICE + void fast_divmod(int& quotient, int64_t& remainder, int64_t dividend) const { + +#if defined(__MACA_ARCH__) + // Use IMUL.HI if divisor != 1, else simply copy the source. + quotient = (divisor != 1) ? __umulhi(dividend, multiplier) >> shift_right : dividend; +#else + quotient = int((divisor != 1) ? ((dividend * multiplier) >> 32) >> shift_right : dividend); +#endif + // The remainder. + remainder = dividend - (quotient * divisor); + } + + /// Construct the FastDivmod object, in host code ideally. + /// + /// This precomputes some values based on the divisor and is computationally expensive. + + MCTLASS_HOST_DEVICE + FastDivmod(): divisor(0), multiplier(0), shift_right(0) { } + + MCTLASS_HOST_DEVICE + FastDivmod(int divisor): divisor(divisor) { + + if (divisor != 1) { + unsigned int p = 31 + find_log2(divisor); + unsigned m = unsigned(((1ull << p) + unsigned(divisor) - 1) / unsigned(divisor)); + + multiplier = m; + shift_right = p - 32; + } else { + multiplier = 0; + shift_right = 0; + } + } + + /// Computes integer division and modulus using precomputed values. This is computationally + /// inexpensive. + MCTLASS_HOST_DEVICE + void operator()(int "ient, int &remainder, int dividend) const { + fast_divmod(quotient, remainder, dividend); + } + + /// Computes integer division using precomputed values. This is computationally + /// inexpensive. + MCTLASS_HOST_DEVICE + int div(int dividend) const { + int quotient, remainder; + fast_divmod(quotient, remainder, dividend); + return quotient; + } + + + /// Computes integer division and modulus using precomputed values. This is computationally + /// inexpensive. + /// + /// Simply returns the quotient + MCTLASS_HOST_DEVICE + int divmod(int &remainder, int dividend) const { + int quotient; + fast_divmod(quotient, remainder, dividend); + return quotient; + } + + /// Computes integer division and modulus using precomputed values. This is computationally + /// inexpensive. + MCTLASS_HOST_DEVICE + void operator()(int "ient, int64_t &remainder, int64_t dividend) const { + fast_divmod(quotient, remainder, dividend); + } + + /// Computes integer division and modulus using precomputed values. This is computationally + /// inexpensive. + MCTLASS_HOST_DEVICE + int divmod(int64_t &remainder, int64_t dividend) const { + int quotient; + fast_divmod(quotient, remainder, dividend); + return quotient; + } + + /// Returns the divisor when cast to integer + MCTLASS_HOST_DEVICE + operator int() const { return divisor; } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Object to encapsulate the fast division+modulus operation for 64b integer division. +/// +/// This object precomputes two values used to accelerate the computation and is best used +/// when the divisor is a grid-invariant. In this case, it may be computed in host code and +/// marshalled along other kernel arguments using the 'Params' pattern. +/// +/// Example: +/// +/// +/// uint64_t quotient, remainder, dividend, divisor; +/// +/// FastDivmodU64 divmod(divisor); +/// +/// divmod(quotient, remainder, dividend); +/// +/// // quotient = (dividend / divisor) +/// // remainder = (dividend % divisor) +/// +struct FastDivmodU64 { + + uint64_t divisor; + uint64_t multiplier; + unsigned int shift_right; + unsigned int round_up; + + // + // Static methods + // + + /// Computes b, where 2^b is the greatest power of two that is less than or equal to x + MCTLASS_HOST_DEVICE + static uint32_t integer_log2(uint64_t x) { + uint32_t n = 0; + while (x >>= 1) { + ++n; + } + return n; + } + + /// Default ctor + MCTLASS_HOST_DEVICE + FastDivmodU64(): divisor(0), multiplier(0), shift_right(0), round_up(0) { } + + /// Construct the FastDivmod object, in host code ideally. + /// + /// This precomputes some values based on the divisor and is computationally expensive. + MCTLASS_HOST_DEVICE + FastDivmodU64(uint64_t divisor_): divisor(divisor_), multiplier(1), shift_right(0), round_up(0) { + + if (divisor) { + shift_right = integer_log2(divisor); + + if ((divisor & (divisor - 1)) == 0) { + multiplier = 0; + } + else { + uint64_t power_of_two = (uint64_t(1) << shift_right); + uint64_t multiplier_lo = uint128_t(0, power_of_two) / divisor; + multiplier = uint128_t(power_of_two, power_of_two) / divisor; + round_up = (multiplier_lo == multiplier ? 1 : 0); + } + } + } + + /// Returns the quotient of floor(dividend / divisor) + MCTLASS_HOST_DEVICE + uint64_t divide(uint64_t dividend) const { + uint64_t quotient = 0; + + #if defined(__MACA_ARCH__) + uint64_t x = dividend; + if (multiplier) { + x = __umul64hi(dividend + round_up, multiplier); + } + quotient = (x >> shift_right); + #else + quotient = dividend / divisor; + #endif + + return quotient; + } + + /// Computes the remainder given a computed quotient and dividend + MCTLASS_HOST_DEVICE + uint64_t modulus(uint64_t quotient, uint64_t dividend) const { + return uint32_t(dividend - quotient * divisor); + } + + /// Returns the quotient of floor(dividend / divisor) and computes the remainder + MCTLASS_HOST_DEVICE + uint64_t divmod(uint64_t &remainder, uint64_t dividend) const { + uint64_t quotient = divide(dividend); + remainder = modulus(quotient, dividend); + return quotient; + } + + /// Computes integer division and modulus using precomputed values. This is computationally + /// inexpensive. + MCTLASS_HOST_DEVICE + void operator()(uint64_t "ient, uint64_t &remainder, uint64_t dividend) const { + quotient = divmod(remainder, dividend); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Computes the coordinate decomposition from a linear index (64-bit linear index => coord) +/// +/// This decomposition is accelerated by the FastDivmodU64 object. It is assumed that +/// a coordinate of indices can be decomposed by div/mod operations. +/// Note, is assumed that element divmod[0] divides by extent[1]. +/// +/// For example, assume 4-D coordinate (n, p, q, c) is mapped to a linear index `npqc`. This +/// can be decomposed via three divide and modulus operations: +/// +/// c = npqc % C; | divmod[2] = FastDivmodU64(C) +/// npq = npqc / C; | coord[3] = c +/// +/// q = npq % Q; | divmod[1] = FastDivmodU64(Q) +/// np = npq / Q; | coord[2] = q +/// +/// p = np % P; | divmod[0] = FastDivmodU64(P) +/// n = np / P; | coord[1] = p +/// +/// | coord[0] = n +/// +template +MCTLASS_HOST_DEVICE Coord CoordinateDecomposition( + uint64_t linear_idx, ///< Linear index to decompose + FastDivmodU64 const *divmod) { ///< Pointer to array of Rank-1 FastDivmodU64 objects + + static_assert(Rank > 0, "CoordinateDecomposition requires Rank=1 or greater."); + + Coord coord; + + MCTLASS_PRAGMA_UNROLL + for (int i = Rank; i > 1; --i) { + uint64_t remainder; + linear_idx = divmod[i - 2].divmod(remainder, linear_idx); + coord[i - 1] = int(remainder); + } + + coord[0] = int(linear_idx); + + return coord; +} + +/// Computes the coordinate decomposition from a linear index (32-bit linear index => coord) +template +MCTLASS_HOST_DEVICE Coord CoordinateDecomposition( + int linear_idx, ///< Linear index to decompose + FastDivmod const *divmod) { ///< Pointer to array of Rank-1 FastDivmodU64 objects + + static_assert(Rank > 0, "CoordinateDecomposition requires Rank=1 or greater."); + + Coord coord; + + MCTLASS_PRAGMA_UNROLL + for (int i = Rank; i > 1; --i) { + int remainder; + linear_idx = divmod[i - 2].divmod(remainder, linear_idx); + coord[i - 1] = int(remainder); + } + + coord[0] = int(linear_idx); + + return coord; +} + +template +MCTLASS_HOST_DEVICE Coord CoordinateDecompositionLittleEndian( + uint64_t linear_idx, ///< Linear index to decompose + FastDivmodU64 const *divmod) { ///< Pointer to array of Rank-1 FastDivmodU64 objects + + static_assert(Rank > 0, "CoordinateDecomposition requires Rank=1 or greater."); + + Coord coord; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Rank - 1; ++i) { + uint64_t remainder; + linear_idx = divmod[i].divmod(remainder, linear_idx); + coord[i] = int(remainder); + } + + coord[Rank - 1] = int(linear_idx); + + return coord; +} + +/// Computes the coordinate decomposition from a linear index (32-bit linear index => coord) +template +MCTLASS_HOST_DEVICE Coord CoordinateDecompositionLittleEndian( + int linear_idx, ///< Linear index to decompose + FastDivmod const *divmod) { ///< Pointer to array of Rank-1 FastDivmodU64 objects + + static_assert(Rank > 0, "CoordinateDecomposition requires Rank=1 or greater."); + + Coord coord; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Rank - 1; ++i) { + int remainder; + linear_idx = divmod[i].divmod(remainder, linear_idx); + coord[i] = int(remainder); + } + + coord[Rank - 1] = int(linear_idx); + + return coord; +} + +/// Safely computes the offset of a linear index in bytes for all types +template +MCTLASS_HOST_DEVICE int64_t OffsetBytes(int64_t index) { + + static_assert( + (sizeof_bits::value >= 8 && !(sizeof_bits::value % 8)) || + (sizeof_bits::value < 8 && !(8 % sizeof_bits::value)), + "Size of numeric type in bits must either be divisible by 8 bits, or 8 bits must be divisible by the size."); + + if (sizeof_bits::value >= 8) { + return index * (sizeof_bits::value / 8); + } + else { + int const kElementsPerByte = ((8 / sizeof_bits::value) + ((sizeof_bits::value >= 8) ? 1 : 0)); + return index / kElementsPerByte; + } +} + +MCTLASS_HOST_DEVICE int64_t OffsetBytes(int64_t index, int64_t element_sizeof_bits) { + if (element_sizeof_bits >= 8) { + return index * (element_sizeof_bits / 8); + } + else { + int64_t const kElementsPerByte = ((8 / element_sizeof_bits) + ((element_sizeof_bits >= 8) ? 1 : 0)); + return index / kElementsPerByte; + } +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Min/Max +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct Min { + static int const kValue = (A < B) ? A : B; +}; + +template +struct Max { + static int const kValue = (A > B) ? A : B; +}; + +MCTLASS_HOST_DEVICE +constexpr int const_min(int a, int b) { + return (b < a ? b : a); +} + +MCTLASS_HOST_DEVICE +constexpr int const_max(int a, int b) { + return (b > a ? b : a); +} + +template +MCTLASS_HOST_DEVICE +T fast_min(T a, T b) { + return (b < a ? b : a); +} + +template <> +MCTLASS_HOST_DEVICE +float fast_min(float a, float b) { + return fminf(a, b); +} + +template +MCTLASS_HOST_DEVICE +T fast_max(T a, T b) { + return (a < b ? b : a); +} + +template <> +MCTLASS_HOST_DEVICE +float fast_max(float a, float b) { + return fmaxf(a, b); +} + +MCTLASS_HOST_DEVICE +float fast_cos(float theta) { + #if defined(__MACA_ARCH__) + return ::cosf(theta); + #else + return std::cos(theta); + #endif +} + +MCTLASS_HOST_DEVICE +double fast_cos(double theta) { + #if defined(__MACA_ARCH__) + return ::cos(theta); + #else + return std::cos(theta); + #endif +} + +MCTLASS_HOST_DEVICE +float fast_sin(float theta) { + #if defined(__MACA_ARCH__) + return ::sinf(theta); + #else + return std::sin(theta); + #endif +} + +MCTLASS_HOST_DEVICE +double fast_sin(double theta) { + #if defined(__MACA_ARCH__) + return ::sin(theta); + #else + return std::sin(theta); + #endif +} + +MCTLASS_HOST_DEVICE +float fast_acos(float theta) { + #if defined(__MACA_ARCH__) + return ::acosf(theta); + #else + return std::acos(theta); + #endif +} + +MCTLASS_HOST_DEVICE +double fast_acos(double theta) { + #if defined(__MACA_ARCH__) + return ::acos(theta); + #else + return std::acos(theta); + #endif +} + +MCTLASS_HOST_DEVICE +float fast_asin(float theta) { + #if defined(__MACA_ARCH__) + return ::asinf(theta); + #else + return std::asin(theta); + #endif +} + +MCTLASS_HOST_DEVICE +double fast_asin(double theta) { + #if defined(__MACA_ARCH__) + return ::asin(theta); + #else + return std::asin(theta); + #endif +} + +MCTLASS_HOST_DEVICE +float fast_sqrt(float theta) { + #if defined(__MACA_ARCH__) + return ::sqrtf(theta); + #else + return std::sqrt(theta); + #endif +} + +MCTLASS_HOST_DEVICE +double fast_sqrt(double theta) { + #if defined(__MACA_ARCH__) + return ::sqrt(theta); + #else + return std::sqrt(theta); + #endif +} + +MCTLASS_HOST_DEVICE +float fast_exp(float x) { + #if defined(__MACA_ARCH__) + return ::expf(x); + #else + return std::exp(x); + #endif +} + +MCTLASS_HOST_DEVICE +double fast_exp(double x) { + #if defined(__MACA_ARCH__) + return ::exp(x); + #else + return std::exp(x); + #endif +} + +MCTLASS_HOST_DEVICE +half_t fast_exp(half_t x) { + #if defined(__MACA_ARCH__) + return (half_t)(::hexp(x.to_half())); + #else + return (half_t)(fast_exp(float(x))); + #endif +} + +MCTLASS_HOST_DEVICE +float fast_log(float x) { + #if defined(__MACA_ARCH__) + return ::logf(x); + #else + return std::log(x); + #endif +} + +MCTLASS_HOST_DEVICE +double fast_log(double x) { + #if defined(__MACA_ARCH__) + return ::log(x); + #else + return std::log(x); + #endif +} + +MCTLASS_HOST_DEVICE +float fast_tanh(float x) { + #if defined(__MACA_ARCH__) + return ::tanhf(x); + #else + return std::tanh(x); + #endif +} + +MCTLASS_HOST_DEVICE +double fast_tanh(double x) { + #if defined(__MACA_ARCH__) + return ::tanh(x); + #else + return std::tanh(x); + #endif +} + +MCTLASS_HOST_DEVICE +half_t fast_tanh(half_t x) { + return half_t(fast_tanh(float(x))); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct fast_exp_op { + MCTLASS_HOST_DEVICE + T operator()(T const &rhs) const { + return fast_exp(rhs); + } +}; + +#if defined(__MACA_ARCH__) +template +struct fast_exp_op> { + MCTLASS_DEVICE + Array operator()(Array const &rhs) const { + + Array result; + + // use x2 specialization + __half2 const *in = reinterpret_cast<__half2 const *>(&rhs); + __half2 *out = reinterpret_cast<__half2 *>(&result); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + out[i] = ::h2exp(in[i]); + } + + // residual + if (N % 2) { + half_t last = rhs[N - 1]; + result[N - 1] = half_t(::hexp(last.to_half())); + } + + return result; + } +}; +#endif // #if defined(__MACA_ARCH__) + +template +struct fast_exp_op> { + MCTLASS_HOST_DEVICE + Array operator()(Array const &rhs) const { + + fast_exp_op fast_op; + Array y; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + y[i] = fast_op(rhs[i]); + } + + return y; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct fast_tanh_op { + MCTLASS_HOST_DEVICE + T operator()(T const &rhs) const { + return fast_tanh(rhs); + } +}; + +//#if defined(__MACA_ARCH__) +#if 0 +template +struct fast_tanh_op> { + MCTLASS_DEVICE + Array operator()(Array const &rhs) const { + + Array result; + + // use x2 specialization + uint32_t const *in = reinterpret_cast(&rhs); + uint32_t *out = reinterpret_cast(&result); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + asm volatile ("tanh.approx.f16x2 %0, %1;" : "=r"(out[i]) : "r"(in[i])); + } + + // residual + if (N % 2) { + uint16_t const *in = reinterpret_cast(&rhs); + uint16_t *out = reinterpret_cast(&result); + asm volatile ("tanh.approx.f16 %0, %1;" : "=h"(out[N - 1]) : "h"(in[N - 1])); + } + + return result; + } +}; +#endif // #if defined(__MACA_ARCH__) + +template +struct fast_tanh_op> { + MCTLASS_HOST_DEVICE + Array operator()(Array const &rhs) const { + + fast_tanh_op fast_op; + Array y; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + y[i] = fast_op(rhs[i]); + } + + return y; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Absolute value function +template +MCTLASS_HOST_DEVICE +T absolute_value(T x) { + if (x < T()) { + return -x; + } + return x; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/float8.h b/csrc/mctlass/include/mctlass/float8.h new file mode 100644 index 0000000..370c048 --- /dev/null +++ b/csrc/mctlass/include/mctlass/float8.h @@ -0,0 +1,1207 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Defines a class for using IEEE half-precision floating-point types in host or + device code. +*/ +#pragma once + +// FP8 types are available starting MACA 11.8+ +#if 0 +#define MACA_FP8_ENABLED 1 +#endif + +#if defined(__MACA_ARCH__) +#if 0 +# define MACA_PTX_FP8_CVT_ENABLED 1 +# endif +#endif // defined(__MACA_ARCH__) + +#ifdef __GNUC__ +// Ignore checks on reinterpret-casts that are being used for bitcasts. +#pragma GCC diagnostic ignored "-Wstrict-aliasing" +#endif + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(__MACACC_RTC__) + +#include "mctlass/floating_point_nvrtc.h" + +#else +// +// Standard Library headers belong here to avoid conflicts with NVRTC. +// +#include +#include +#include +#include +#endif + +#ifdef MACA_FP8_ENABLED +#include +#endif +#include + +#include "mctlass/mctlass.h" + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// FP8 Has 2 encodings possible : E4M3 and E5M2 +// +// E4M3 : 7 | 6 5 4 3 | 2 1 0 +// E5M2 : 7 | 6 5 4 3 2 | 1 0 +// +/////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class FloatEncoding { + E4M3, + E5M2 +}; + +template +struct alignas(1) float8_base { + + static constexpr bool IS_E4M3 = (T == FloatEncoding::E4M3); + static constexpr bool IS_E5M2 = (T == FloatEncoding::E5M2); + + // Number of Bits representing mantissa and exponents + static constexpr int FP32_NUM_BITS = 32; + static constexpr int FP32_NUM_EXPONENT_BITS = 8; + static constexpr int FP32_NUM_MANTISSA_BITS = 23; + static constexpr uint32_t FP32_NAN = 0x7fffffff; + static constexpr uint32_t FP32_INFINITY_MASK = 0x7f800000; + static constexpr int FP32_MAX_EXPONENT = 127; + static constexpr int FP32_MIN_EXPONENT = -126; + static constexpr int FP32_EXPONENT_BIAS = 127; + + static constexpr int FP16_NUM_BITS = 16; + static constexpr int FP16_NUM_EXPONENT_BITS = 5; + static constexpr int FP16_NUM_MANTISSA_BITS = 10; + static constexpr uint16_t FP16_NAN = 0x7fff; + static constexpr uint16_t FP16_INFINITY_MASK = 0x7c00; + static constexpr int FP16_MAX_EXPONENT = 15; + static constexpr int FP16_MIN_EXPONENT = -14; + static constexpr int FP16_EXPONENT_BIAS = 15; + + static constexpr int FP8_NUM_BITS = 8; + static constexpr int FP8_NUM_EXPONENT_BITS = IS_E4M3 ? 4 : 5; + static constexpr int FP8_NUM_MANTISSA_BITS = IS_E4M3 ? 3 : 2; + static constexpr uint8_t FP8_NAN = 0x7f; // Also F8_INF + static constexpr uint8_t FP8_INFINITY_MASK = IS_E4M3 ? 0x78 : 0x7c; + static constexpr int FP8_MAX_EXPONENT = IS_E4M3 ? 7 : 15; + static constexpr int FP8_MIN_EXPONENT = IS_E4M3 ? -6 : -14; + static constexpr int FP8_EXPONENT_BIAS = IS_E4M3 ? 7 : 15; + + static constexpr uint8_t FP8_EXPONENT_MASK = (1 << FP8_NUM_EXPONENT_BITS) - 1; + static constexpr uint8_t FP8_MANTISSA_MASK = (1 << FP8_NUM_MANTISSA_BITS) - 1; + + static constexpr uint8_t FP8_MAX_FLT = (IS_E4M3 ? 0x7e : 0x7b); + + // 256 in float + static constexpr uint32_t FP8_SAT_VAL_FP32 = 0x43800000; + + // + // Data members + // + + /// Data container + uint8_t storage; + + /// Ctors. + MCTLASS_HOST_DEVICE + float8_base() : storage(0) { } + + /// Is finite implementation + MCTLASS_HOST_DEVICE + static bool isfinite(float flt) { + uint32_t s; + + #if defined(__MACA_ARCH__) + s = reinterpret_cast(flt); + #else + std::memcpy(&s, &flt, sizeof(s)); + #endif + + return (s & 0x7f800000) < 0x7f800000; + } + + /// Is NaN implementation + MCTLASS_HOST_DEVICE + static bool isnan(float flt) { + uint32_t s; + + #if defined(__MACA_ARCH__) + s = reinterpret_cast(flt); + #else + std::memcpy(&s, &flt, sizeof(s)); + #endif + + return (s & 0x7fffffff) > 0x7f800000; + } + + /// Is infinite implementation + MCTLASS_HOST_DEVICE + static bool isinf(float flt) { + uint32_t s; + + #if defined(__MACA_ARCH__) + s = reinterpret_cast(flt); + #else + std::memcpy(&s, &flt, sizeof(s)); + #endif + + // Sign = 0 for +inf, 1 for -inf + // Exponent = all ones + // Mantissa = all zeros + return (s == 0x7f800000) || (s == 0xff800000); + } + + /// FP32 -> FP8 conversion - rounds to nearest even + MCTLASS_HOST_DEVICE + static uint8_t convert_float_to_fp8(float const& flt) { + + // software implementation rounds toward nearest even + uint32_t s; + + #if defined(__MACA_ARCH__) + s = reinterpret_cast(flt); + #else + std::memcpy(&s, &flt, sizeof(s)); + #endif + + // Extract the bits in the FP32 type + uint8_t sign = uint8_t((s >> 24 & 0x80)); + int8_t exp = uint8_t(((s >> FP32_NUM_MANTISSA_BITS) & 0xff) - FP32_EXPONENT_BIAS); + int mantissa = s & 0x7fffff; + uint8_t u = 0; + + uint8_t const kF8_NaN = 0x7f; + + // NaN => NaN + if (isnan(flt)) { + return kF8_NaN; + } + + // Inf => MAX_FLT (satfinite) + if (isinf(flt)) { + return sign | FP8_MAX_FLT; + } + + // Special handling + if ( exp == -128 ) { + // int8 range is from -128 to 127 + // So 255(inf) - 127(bias) = 128 - will show up as -128 + + // satfinite + return (sign | FP8_MAX_FLT); + } + + int sticky_bit = 0; + + bool skip_sign = false; + bool may_be_nan = false; + + if ( (exp >= FP8_MIN_EXPONENT) && (exp <= FP8_MAX_EXPONENT) ) { + // normal fp32 to normal fp8 + exp = uint8_t(exp + uint8_t(FP8_EXPONENT_BIAS)); + u = uint8_t(((exp & FP8_EXPONENT_MASK) << FP8_NUM_MANTISSA_BITS)); + u = uint8_t(u | (mantissa >> (FP32_NUM_MANTISSA_BITS - FP8_NUM_MANTISSA_BITS))); + } else if(exp < FP8_MIN_EXPONENT) { + // normal single-precision to subnormal float8-precision representation + int rshift = (FP8_MIN_EXPONENT - exp); + if (rshift < FP32_NUM_BITS) { + mantissa |= (1 << FP32_NUM_MANTISSA_BITS); + + sticky_bit = ((mantissa & ((1 << rshift) - 1)) != 0); + + mantissa = (mantissa >> rshift); + u = (uint8_t(mantissa >> (FP32_NUM_MANTISSA_BITS- FP8_NUM_MANTISSA_BITS)) & FP8_MANTISSA_MASK); + } else { + mantissa = 0; + u = 0; + } + // Exponent > FP8_MAX_EXPONENT - this is a special case done to match HW + // 0x4380_0000 to 0x43e0_0000 - maps from 256 to 448, and does not saturate / inf. + } else { + if( exp == (FP8_MAX_EXPONENT + 1) ) { + uint8_t mantissa_tmp = uint8_t(mantissa >> (FP32_NUM_MANTISSA_BITS - FP8_NUM_MANTISSA_BITS)); + if( mantissa_tmp < FP8_MANTISSA_MASK) { + exp = uint8_t(exp + uint8_t(FP8_EXPONENT_BIAS)); + u = uint8_t(exp << FP8_NUM_MANTISSA_BITS) | mantissa_tmp; + may_be_nan = (mantissa_tmp == (FP8_MANTISSA_MASK-1)); + } else { + // satfinite + return (sign | FP8_MAX_FLT); + } + } else{ + // satfinite + return (sign | FP8_MAX_FLT); + } + } + + // round to nearest even + int NUM_BITS_SHIFT = FP32_NUM_MANTISSA_BITS - (FP8_NUM_MANTISSA_BITS + 1); + int round_bit = ((mantissa >> NUM_BITS_SHIFT) & 1); + sticky_bit |= ((mantissa & ((1 << NUM_BITS_SHIFT) - 1)) != 0); + + if ((round_bit && sticky_bit) || (round_bit && (u & 1))) { + u = uint8_t(u + 1); + if( may_be_nan ) { + skip_sign = true; + } + } + + if (u > FP8_MAX_FLT) { + // satfinite + u = (sign | FP8_MAX_FLT); + } + + if( ! skip_sign ) { + u |= sign; + } + + return u; + } + + + /// Converts a fp8 value stored as a uint8_t to a float + MCTLASS_HOST_DEVICE + static float convert_fp8_to_float(uint8_t const& x) { + + uint32_t constexpr kF32_NaN = 0x7fffffff; + + uint8_t const &f8 = x; + int sign = (f8 >> (FP8_NUM_BITS - 1)) & 1; + int exp = (f8 >> FP8_NUM_MANTISSA_BITS) & FP8_EXPONENT_MASK; + int mantissa = f8 & FP8_MANTISSA_MASK; + unsigned f = (sign << (FP32_NUM_BITS-1)); + + if (IS_E4M3 && exp == 15 && mantissa == 0x7) { + f = kF32_NaN; + } + else if (exp > 0 && (IS_E4M3 || exp < (FP8_MAX_EXPONENT + FP8_EXPONENT_BIAS + 1))) { + // normal + exp += (FP32_EXPONENT_BIAS - FP8_EXPONENT_BIAS); + f = f | + (exp << FP32_NUM_MANTISSA_BITS) | + (mantissa << (FP32_NUM_MANTISSA_BITS-FP8_NUM_MANTISSA_BITS)); + } else if (exp == 0) { + if (mantissa) { + // subnormal + exp += (FP32_EXPONENT_BIAS - FP8_EXPONENT_BIAS) + 1; + while ((mantissa & (1 << FP8_NUM_MANTISSA_BITS)) == 0) { + mantissa <<= 1; + exp--; + } + mantissa &= FP8_MANTISSA_MASK; + f = f | + (exp << FP32_NUM_MANTISSA_BITS) | + (mantissa << (FP32_NUM_MANTISSA_BITS-FP8_NUM_MANTISSA_BITS)); + } else { + // sign-preserving zero + } + } else { + if(mantissa == 0){ + // Sign-preserving infinity + f = (f | 0x7f800000); + } else { + // Canonical NaN + f = kF32_NaN; + } + } + + #if defined(__MACA_ARCH__) + return reinterpret_cast(f); + #else + float flt; + std::memcpy(&flt, &f, sizeof(flt)); + return flt; + #endif + } +}; + + +// Forward declaration of float_e5m2_t to define float_e4m3_t <=> float_e5m2_t +// conversions in class float_e4m3_t +struct float_e5m2_t; + + +/////////////////////////////////////////////////////////////// +/// +/// floating-point 8 type : E4M3 +/// +/////////////////////////////////////////////////////////////// +struct alignas(1) float_e4m3_t : float8_base { + + using Base = float8_base; + + static constexpr int MAX_EXPONENT = Base::FP8_MAX_EXPONENT; + + // + // Static conversion operators + // + + /// Constructs from an uint8_t + MCTLASS_HOST_DEVICE + static float_e4m3_t bitcast(uint8_t x) { + float_e4m3_t f; + f.storage = x; + return f; + } + + /// FP32 -> FP8 conversion - rounds to nearest even + MCTLASS_HOST_DEVICE + static float_e4m3_t from_float(float const& flt) { + #if defined(MACA_PTX_FP8_CVT_ENABLED) + uint16_t tmp; + float y = float(); + asm volatile("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(tmp) : "f"(y), "f"(flt)); + + return *reinterpret_cast(&tmp); + #else + return bitcast(Base::convert_float_to_fp8(flt)); + #endif + } + + /// FP16 -> E5M2 conversion - rounds to nearest even + MCTLASS_HOST_DEVICE + static float_e4m3_t from_half(half const& flt) { + #if defined(MACA_PTX_FP8_CVT_ENABLED) + uint16_t tmp = 0; + uint32_t bits = reinterpret_cast(flt); + asm volatile("cvt.rn.satfinite.e4m3x2.f16x2 %0, %1;" : "=h"(tmp) : "r"(bits)); + + return *reinterpret_cast(&tmp); + #else + return bitcast(Base::convert_float_to_fp8(__half2float(flt))); + #endif + } + + // E4M3 -> half + MCTLASS_HOST_DEVICE + static half to_half(float_e4m3_t const& x) { + #if defined(MACA_PTX_FP8_CVT_ENABLED) + uint16_t bits = x.storage; + uint32_t packed; + asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;\n" : "=r"(packed) : "h"(bits)); + + return reinterpret_cast(packed).x; + #else + return __float2half(Base::convert_fp8_to_float(x.storage)); + #endif + } + + // E4M3 -> Float + MCTLASS_HOST_DEVICE + static float to_float(float_e4m3_t const& x) { + #if defined(MACA_PTX_FP8_CVT_ENABLED) + uint16_t bits = x.storage; + uint32_t packed; + asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;\n" : "=r"(packed) : "h"(bits)); + + return __half2float(reinterpret_cast(packed).x); + #else + return Base::convert_fp8_to_float(x.storage); + #endif + } + + // + // Methods + // + + /// Constructor inheritance + using Base::Base; + +#ifdef MACA_FP8_ENABLED + /// Conversion from MACA's FP8 type + MCTLASS_HOST_DEVICE + explicit float_e4m3_t(__nv_fp8_e4m3 x) { + storage = x.__x; + } +#endif + + /// Floating point conversion + MCTLASS_HOST_DEVICE + explicit float_e4m3_t(float x) { + storage = from_float(x).storage; + } + + MCTLASS_HOST_DEVICE + explicit float_e4m3_t(half x) { + storage = from_half(x).storage; + } + + /// Floating point conversion + MCTLASS_HOST_DEVICE + explicit float_e4m3_t(double x): float_e4m3_t(float(x)) { + } + + /// Integer conversion + MCTLASS_HOST_DEVICE + explicit float_e4m3_t(int x): float_e4m3_t(float(x)) { + } + + /// E5M2 conversion. Defined after float_e5m2_t is defined. + MCTLASS_HOST_DEVICE + explicit float_e4m3_t(float_e5m2_t x); + +#ifdef MACA_FP8_ENABLED + /// Assignment from MACA's FP8 type + MCTLASS_HOST_DEVICE + float_e4m3_t & operator=(__nv_fp8_e4m3 x) { + storage = x.__x; + return *this; + } +#endif + + /// Converts to float + MCTLASS_HOST_DEVICE + operator float() const { + return to_float(*this); + } + + /// Converts to half + MCTLASS_HOST_DEVICE + operator half() const { + return to_half(*this); + } + + /// Converts to float + MCTLASS_HOST_DEVICE + explicit operator double() const { + return double(to_float(*this)); + } + + /// Converts to int + MCTLASS_HOST_DEVICE + explicit operator int() const { + #if defined(__MACA_ARCH__) + return __half2int_rn(to_half(*this)); + #else + return int(to_float(*this)); + #endif + } + + /// Casts to bool + MCTLASS_HOST_DEVICE + explicit operator bool() const { + #if defined(__MACA_ARCH__) + return bool(__half2int_rn(to_half(*this))); + #else + return bool(int(to_float(*this))); + #endif + } + + /// Accesses raw internal state + MCTLASS_HOST_DEVICE + uint8_t& raw() { + return storage; + } + + /// Accesses raw internal state + MCTLASS_HOST_DEVICE + uint8_t raw() const { + return storage; + } + + /// Returns the sign bit + MCTLASS_HOST_DEVICE + bool signbit() const { + return ((storage & (1 << (Base::FP8_NUM_BITS - 1))) != 0); + } + + /// Returns the biased exponent + MCTLASS_HOST_DEVICE + int exponent_biased() const { + return int((storage >> FP8_NUM_MANTISSA_BITS) & Base::FP8_EXPONENT_MASK); + } + + /// Returns the unbiased exponent + MCTLASS_HOST_DEVICE + int exponent() const { + return exponent_biased() - 15; + } + + /// Returns the mantissa + MCTLASS_HOST_DEVICE + int mantissa() const { + return int(storage & Base::FP8_MANTISSA_MASK); + } +}; +/////////////////////////////////////////////////////////////// +/// +/// floating-point 8 type : E5M2 +/// +/////////////////////////////////////////////////////////////// +struct alignas(1) float_e5m2_t : float8_base { + + using Base = float8_base; + + static constexpr int MAX_EXPONENT = Base::FP8_MAX_EXPONENT; + + // + // Static conversion operators + // + + /// Constructs from an uint8_t + MCTLASS_HOST_DEVICE + static float_e5m2_t bitcast(uint8_t x) { + float_e5m2_t f; + f.storage = x; + return f; + } + + /// FP32 -> FP8 conversion - rounds to nearest even + MCTLASS_HOST_DEVICE + static float_e5m2_t from_float(float const& flt) { + #if defined(MACA_PTX_FP8_CVT_ENABLED) + uint16_t tmp; + float y = float(); + asm volatile("cvt.rn.satfinite.e5m2x2.f32 %0, %1, %2;" : "=h"(tmp) : "f"(y), "f"(flt)); + + return *reinterpret_cast(&tmp); + #else + return bitcast(Base::convert_float_to_fp8(flt)); + #endif + } + + /// FP16 -> E5M2 conversion - rounds to nearest even + MCTLASS_HOST_DEVICE + static float_e5m2_t from_half(half const& flt) { + #if defined(MACA_PTX_FP8_CVT_ENABLED) + uint16_t tmp = 0; + uint32_t bits = reinterpret_cast(flt); + asm volatile("cvt.rn.satfinite.e5m2x2.f16x2 %0, %1;" : "=h"(tmp) : "r"(bits)); + + return *reinterpret_cast(&tmp); + #else + return bitcast(Base::convert_float_to_fp8(__half2float(flt))); + #endif + } + + // E5M2 -> half + MCTLASS_HOST_DEVICE + static half to_half(float_e5m2_t const& x) { + #if defined(MACA_PTX_FP8_CVT_ENABLED) + uint16_t bits = x.storage; + uint32_t packed; + asm volatile("cvt.rn.f16x2.e5m2x2 %0, %1;\n" : "=r"(packed) : "h"(bits)); + + return reinterpret_cast(packed).x; + #else + return __float2half(Base::convert_fp8_to_float(x.storage)); + #endif + } + + // E5M2 -> Float + MCTLASS_HOST_DEVICE + static float to_float(float_e5m2_t const& x) { + #if defined(MACA_PTX_FP8_CVT_ENABLED) + uint16_t bits = x.storage; + uint32_t packed; + asm volatile("cvt.rn.f16x2.e5m2x2 %0, %1;\n" : "=r"(packed) : "h"(bits)); + + return __half2float(reinterpret_cast(packed).x); + #else + return Base::convert_fp8_to_float(x.storage); + #endif + } + + // + // Methods + // + + /// Constructor inheritance + using Base::Base; + +#ifdef MACA_FP8_ENABLED + /// Conversion from MACA's FP8 type + MCTLASS_HOST_DEVICE + explicit float_e5m2_t(__nv_fp8_e5m2 x) { + storage = x.__x; + } +#endif + + /// Floating point conversion + MCTLASS_HOST_DEVICE + explicit float_e5m2_t(float x) { + storage = from_float(x).storage; + } + + MCTLASS_HOST_DEVICE + explicit float_e5m2_t(half x) { + storage = from_half(x).storage; + } + + /// Floating point conversion + MCTLASS_HOST_DEVICE + explicit float_e5m2_t(double x): float_e5m2_t(float(x)) { + } + + /// Integer conversion + MCTLASS_HOST_DEVICE + explicit float_e5m2_t(int x): float_e5m2_t(float(x)) { + } + + /// E4M3 conversion + MCTLASS_HOST_DEVICE + explicit float_e5m2_t(float_e4m3_t x); + +#ifdef MACA_FP8_ENABLED + /// Assignment from MACA's FP8 type + MCTLASS_HOST_DEVICE + float_e5m2_t & operator=(__nv_fp8_e5m2 x) { + storage = x.__x; + return *this; + } +#endif + + /// Converts to float + MCTLASS_HOST_DEVICE + operator float() const { + return to_float(*this); + } + + /// Converts to half + MCTLASS_HOST_DEVICE + operator half() const { + return to_half(*this); + } + + /// Converts to float + MCTLASS_HOST_DEVICE + explicit operator double() const { + return double(to_float(*this)); + } + + /// Converts to int + MCTLASS_HOST_DEVICE + explicit operator int() const { + #if defined(__MACA_ARCH__) + return __half2int_rn(to_half(*this)); + #else + return int(to_float(*this)); + #endif + } + + /// Casts to bool + MCTLASS_HOST_DEVICE + explicit operator bool() const { + #if defined(__MACA_ARCH__) + return bool(__half2int_rn(to_half(*this))); + #else + return bool(int(to_float(*this))); + #endif + } + + /// Accesses raw internal state + MCTLASS_HOST_DEVICE + uint8_t& raw() { + return storage; + } + + /// Accesses raw internal state + MCTLASS_HOST_DEVICE + uint8_t raw() const { + return storage; + } + + /// Returns the sign bit + MCTLASS_HOST_DEVICE + bool signbit() const { + return ((storage & (1 << (Base::FP8_NUM_BITS - 1))) != 0); + } + + /// Returns the biased exponent + MCTLASS_HOST_DEVICE + int exponent_biased() const { + return int((storage >> FP8_NUM_MANTISSA_BITS) & Base::FP8_EXPONENT_MASK); + } + + /// Returns the unbiased exponent + MCTLASS_HOST_DEVICE + int exponent() const { + return exponent_biased() - 15; + } + + /// Returns the mantissa + MCTLASS_HOST_DEVICE + int mantissa() const { + return int(storage & Base::FP8_MANTISSA_MASK); + } +}; +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Arithmetic operators +// +/////////////////////////////////////////////////////////////////////////////////////////////////// + +MCTLASS_HOST_DEVICE +bool operator==(float_e4m3_t const& lhs, float_e4m3_t const& rhs) { + return float(lhs) == float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator!=(float_e4m3_t const& lhs, float_e4m3_t const& rhs) { + return float(lhs) != float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator<(float_e4m3_t const& lhs, float_e4m3_t const& rhs) { + return float(lhs) < float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator<=(float_e4m3_t const& lhs, float_e4m3_t const& rhs) { + return float(lhs) <= float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator>(float_e4m3_t const& lhs, float_e4m3_t const& rhs) { + return float(lhs) > float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator>=(float_e4m3_t const& lhs, float_e4m3_t const& rhs) { + return float(lhs) >= float(rhs); +} + +MCTLASS_HOST_DEVICE +float_e4m3_t operator+(float_e4m3_t const& lhs, float_e4m3_t const& rhs) { + return float_e4m3_t(float(lhs) + float(rhs)); +} + +MCTLASS_HOST_DEVICE +float_e4m3_t operator-(float_e4m3_t const& lhs) { + return float_e4m3_t(-float(lhs)); +} + +MCTLASS_HOST_DEVICE +float_e4m3_t operator-(float_e4m3_t const& lhs, float_e4m3_t const& rhs) { + return float_e4m3_t(float(lhs) - float(rhs)); +} + +MCTLASS_HOST_DEVICE +float_e4m3_t operator*(float_e4m3_t const& lhs, float_e4m3_t const& rhs) { + return float_e4m3_t(float(lhs) * float(rhs)); +} + +MCTLASS_HOST_DEVICE +float_e4m3_t operator/(float_e4m3_t const& lhs, float_e4m3_t const& rhs) { + return float_e4m3_t(float(lhs) / float(rhs)); +} + +MCTLASS_HOST_DEVICE +float_e4m3_t& operator+=(float_e4m3_t & lhs, float_e4m3_t const& rhs) { + lhs = float_e4m3_t(float(lhs) + float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +float_e4m3_t& operator-=(float_e4m3_t & lhs, float_e4m3_t const& rhs) { + lhs = float_e4m3_t(float(lhs) - float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +float_e4m3_t& operator*=(float_e4m3_t & lhs, float_e4m3_t const& rhs) { + lhs = float_e4m3_t(float(lhs) * float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +float_e4m3_t& operator/=(float_e4m3_t & lhs, float_e4m3_t const& rhs) { + lhs = float_e4m3_t(float(lhs) / float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +float_e4m3_t& operator++(float_e4m3_t & lhs) { + float tmp(lhs); + ++tmp; + lhs = float_e4m3_t(tmp); + return lhs; +} + +MCTLASS_HOST_DEVICE +float_e4m3_t& operator--(float_e4m3_t & lhs) { + float tmp(lhs); + --tmp; + lhs = float_e4m3_t(tmp); + return lhs; +} + +MCTLASS_HOST_DEVICE +float_e4m3_t operator++(float_e4m3_t & lhs, int) { + float_e4m3_t ret(lhs); + float tmp(lhs); + tmp++; + lhs = float_e4m3_t(tmp); + return ret; +} + +MCTLASS_HOST_DEVICE +float_e4m3_t operator--(float_e4m3_t & lhs, int) { + float_e4m3_t ret(lhs); + float tmp(lhs); + tmp--; + lhs = float_e4m3_t(tmp); + return ret; +} + +MCTLASS_HOST_DEVICE +bool operator==(float_e5m2_t const& lhs, float_e5m2_t const& rhs) { + return float(lhs) == float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator!=(float_e5m2_t const& lhs, float_e5m2_t const& rhs) { + return float(lhs) != float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator<(float_e5m2_t const& lhs, float_e5m2_t const& rhs) { + return float(lhs) < float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator<=(float_e5m2_t const& lhs, float_e5m2_t const& rhs) { + return float(lhs) <= float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator>(float_e5m2_t const& lhs, float_e5m2_t const& rhs) { + return float(lhs) > float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator>=(float_e5m2_t const& lhs, float_e5m2_t const& rhs) { + return float(lhs) >= float(rhs); +} + +MCTLASS_HOST_DEVICE +float_e5m2_t operator+(float_e5m2_t const& lhs, float_e5m2_t const& rhs) { + return float_e5m2_t(float(lhs) + float(rhs)); +} + +MCTLASS_HOST_DEVICE +float_e5m2_t operator-(float_e5m2_t const& lhs) { + return float_e5m2_t(-float(lhs)); +} + +MCTLASS_HOST_DEVICE +float_e5m2_t operator-(float_e5m2_t const& lhs, float_e5m2_t const& rhs) { + return float_e5m2_t(float(lhs) - float(rhs)); +} + +MCTLASS_HOST_DEVICE +float_e5m2_t operator*(float_e5m2_t const& lhs, float_e5m2_t const& rhs) { + return float_e5m2_t(float(lhs) * float(rhs)); +} + +MCTLASS_HOST_DEVICE +float_e5m2_t operator/(float_e5m2_t const& lhs, float_e5m2_t const& rhs) { + return float_e5m2_t(float(lhs) / float(rhs)); +} + +MCTLASS_HOST_DEVICE +float_e5m2_t& operator+=(float_e5m2_t & lhs, float_e5m2_t const& rhs) { + lhs = float_e5m2_t(float(lhs) + float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +float_e5m2_t& operator-=(float_e5m2_t & lhs, float_e5m2_t const& rhs) { + lhs = float_e5m2_t(float(lhs) - float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +float_e5m2_t& operator*=(float_e5m2_t & lhs, float_e5m2_t const& rhs) { + lhs = float_e5m2_t(float(lhs) * float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +float_e5m2_t& operator/=(float_e5m2_t & lhs, float_e5m2_t const& rhs) { + lhs = float_e5m2_t(float(lhs) / float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +float_e5m2_t& operator++(float_e5m2_t & lhs) { + float tmp(lhs); + ++tmp; + lhs = float_e5m2_t(tmp); + return lhs; +} + +MCTLASS_HOST_DEVICE +float_e5m2_t& operator--(float_e5m2_t & lhs) { + float tmp(lhs); + --tmp; + lhs = float_e5m2_t(tmp); + return lhs; +} + +MCTLASS_HOST_DEVICE +float_e5m2_t operator++(float_e5m2_t & lhs, int) { + float_e5m2_t ret(lhs); + float tmp(lhs); + tmp++; + lhs = float_e5m2_t(tmp); + return ret; +} + +MCTLASS_HOST_DEVICE +float_e5m2_t operator--(float_e5m2_t & lhs, int) { + float_e5m2_t ret(lhs); + float tmp(lhs); + tmp--; + lhs = float_e5m2_t(tmp); + return ret; +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// float_e4m3_t <=> float_e5m2_t conversions +// +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// float_e4m3_t <= float_e5m2_t +MCTLASS_HOST_DEVICE +float_e4m3_t::float_e4m3_t(float_e5m2_t x) { + storage = from_float(float_e5m2_t::to_float(x)).storage; +} + +/// float_e5m2_t <= float_e4m3_t +MCTLASS_HOST_DEVICE +float_e5m2_t::float_e5m2_t(float_e4m3_t x) { + storage = from_float(float_e4m3_t::to_float(x)).storage; +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Standard Library operations and definitions +// +/////////////////////////////////////////////////////////////////////////////////////////////////// + +#if !defined(__MACACC_RTC__) +namespace std { + +/// Numeric limits common to all float8 types +template +struct float8_base_numeric_limits { +private: + using F8Type = T; +public: + static bool const is_specialized = true; + static bool const is_signed = true; + static bool const is_integer = false; + static bool const is_exact = false; + static bool const has_quiet_NaN = true; + static bool const has_signaling_NaN = false; + static std::float_denorm_style const has_denorm = std::denorm_present; + static bool const has_denorm_loss = true; + static std::float_round_style const round_style = std::round_to_nearest; + static bool const is_iec559 = false; + static bool const is_bounded = true; + static bool const is_modulo = false; + static int const digits = F8Type::FP8_NUM_MANTISSA_BITS; + + /// Least positive value + static F8Type min() { return F8Type::bitcast(0x01); } + + /// Maximum finite value + static F8Type max() { return F8Type::bitcast(F8Type::FP8_MAX_FLT); } + + /// Returns maximum rounding error + static F8Type round_error() { return F8Type(0.5f); } + + /// Returns positive infinity value + static F8Type infinity() { return F8Type::bitcast(F8Type::FP8_INFINITY_MASK); } + + /// Returns quiet NaN value + static F8Type quiet_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); } + + /// Returns signaling NaN value + static F8Type signaling_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); } + + /// Returns smallest positive subnormal value + static F8Type denorm_min() { return F8Type::bitcast(0x01); } +}; + +/// Numeric limits for float_e4m3_t +template <> +struct numeric_limits : + public float8_base_numeric_limits { + static bool const has_infinity = false; + + /// Minimum finite value + static mctlass::float_e4m3_t lowest() { return mctlass::float_e4m3_t::bitcast(0xfe); } + + /// Returns smallest finite value + static mctlass::float_e4m3_t epsilon() { return mctlass::float_e4m3_t::bitcast(0x20); } +}; + +/// Numeric limits for float_e5m2_t +template <> +struct numeric_limits : + public float8_base_numeric_limits { + static bool const has_infinity = true; + + /// Minimum finite value + static mctlass::float_e5m2_t lowest() { return mctlass::float_e5m2_t::bitcast(0xfb); } + + /// Returns smallest finite value + static mctlass::float_e5m2_t epsilon() { return mctlass::float_e5m2_t::bitcast(0x34); } +}; + +} // namespace std +#endif + +namespace platform { + +/// Numeric limits common to all float8 types +template +struct float8_base_numeric_limits { +private: + using F8Type = T; +public: + static bool const is_specialized = true; + static bool const is_signed = true; + static bool const is_integer = false; + static bool const is_exact = false; + static bool const has_quiet_NaN = true; + static bool const has_signaling_NaN = false; +#if !defined(__MACACC_RTC__) + static std::float_denorm_style const has_denorm = std::denorm_present; +#endif + static bool const has_denorm_loss = true; +#if !defined(__MACACC_RTC__) + static std::float_round_style const round_style = std::round_to_nearest; +#endif + static bool const is_iec559 = false; + static bool const is_bounded = true; + static bool const is_modulo = false; + static int const digits = F8Type::FP8_NUM_MANTISSA_BITS; + + /// Least positive value + static F8Type min() { return F8Type::bitcast(0x01); } + + /// Maximum finite value + static F8Type max() { return F8Type::bitcast(F8Type::FP8_MAX_FLT); } + + /// Returns maximum rounding error + static F8Type round_error() { return F8Type(0.5f); } + + /// Returns positive infinity value + static F8Type infinity() { return F8Type::bitcast(F8Type::FP8_INFINITY_MASK); } + + /// Returns quiet NaN value + static F8Type quiet_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); } + + /// Returns signaling NaN value + static F8Type signaling_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); } + + /// Returns smallest positive subnormal value + static F8Type denorm_min() { return F8Type::bitcast(0x01); } +}; + +/// std::numeric_limits +template +struct numeric_limits; + +/// Numeric limits for float_e4m3_t +template <> +struct numeric_limits : + public float8_base_numeric_limits { + static bool const has_infinity = false; + + /// Minimum finite value + static mctlass::float_e4m3_t lowest() { return mctlass::float_e4m3_t::bitcast(0xfe); } + + /// Returns smallest finite value + static mctlass::float_e4m3_t epsilon() { return mctlass::float_e4m3_t::bitcast(0x20); } +}; + +/// Numeric limits for float_e5m2_t +template <> +struct numeric_limits : + public float8_base_numeric_limits { + static bool const has_infinity = true; + + /// Minimum finite value + static mctlass::float_e5m2_t lowest() { return mctlass::float_e5m2_t::bitcast(0xfb); } + + /// Returns smallest finite value + static mctlass::float_e5m2_t epsilon() { return mctlass::float_e5m2_t::bitcast(0x34); } +}; + +} // namespace platform + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// +// User-defined literals +// + +MCTLASS_HOST_DEVICE +mctlass::float_e4m3_t operator "" _fe4m3(long double x) { + return mctlass::float_e4m3_t(float(x)); +} + +MCTLASS_HOST_DEVICE +mctlass::float_e4m3_t operator "" _fe4m3(unsigned long long int x) { + return mctlass::float_e4m3_t(int(x)); +} + +MCTLASS_HOST_DEVICE +mctlass::float_e5m2_t operator "" _fe5m2(long double x) { + return mctlass::float_e5m2_t(float(x)); +} + +MCTLASS_HOST_DEVICE +mctlass::float_e5m2_t operator "" _fe5m2(unsigned long long int x) { + return mctlass::float_e5m2_t(int(x)); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/floating_point_nvrtc.h b/csrc/mctlass/include/mctlass/floating_point_nvrtc.h new file mode 100644 index 0000000..117fd95 --- /dev/null +++ b/csrc/mctlass/include/mctlass/floating_point_nvrtc.h @@ -0,0 +1,65 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Defines categories for floating point numbers for use in NVRTC-compiled code +*/ + +#pragma once + +namespace mctlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// All floating-point numbers can be put in one of these categories. +enum { + FP_NAN = +# define FP_NAN 0 + FP_NAN, + FP_INFINITE = +# define FP_INFINITE 1 + FP_INFINITE, + FP_ZERO = +# define FP_ZERO 2 + FP_ZERO, + FP_SUBNORMAL = +# define FP_SUBNORMAL 3 + FP_SUBNORMAL, + FP_NORMAL = +# define FP_NORMAL 4 + FP_NORMAL +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +/////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/functional.h b/csrc/mctlass/include/mctlass/functional.h new file mode 100644 index 0000000..a6a29f7 --- /dev/null +++ b/csrc/mctlass/include/mctlass/functional.h @@ -0,0 +1,533 @@ + /*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Define basic numeric operators + + This is inspired by the Standard Library's header. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/half.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +// #include +# include "mctlass/arch/mma.h" +#endif // defined(MCTLASS_ARCH_WMMA_ENABLED) + +#ifdef _MSC_VER +// Provides support for alternate operators such as 'and', 'or', ... +#include +#endif // _MSC_VER + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct absolute_value_op { + MCTLASS_HOST_DEVICE + T operator()(T lhs) const { + return abs(lhs); + } +}; + +template <> +struct absolute_value_op { + MCTLASS_HOST_DEVICE + float operator()(float lhs) const { return fabs(lhs); } +}; + +template +struct plus { + MCTLASS_HOST_DEVICE + T operator()(T lhs, T const &rhs) const { + lhs += rhs; + return lhs; + } +}; + +template +struct minus { + MCTLASS_HOST_DEVICE + T operator()(T lhs, T const &rhs) const { + lhs -= rhs; + return lhs; + } +}; + +template +struct multiplies { + MCTLASS_HOST_DEVICE + T operator()(T lhs, T const &rhs) const { + lhs *= rhs; + return lhs; + } +}; + +#if defined(__MACA_ARCH__) +/// Partial specializations needed when __MACA_NO_HALF2_OPERATORS__ is set +template<> +struct plus<__half2> { + MCTLASS_HOST_DEVICE + __half2 operator()(__half2 lhs, __half2 const &rhs) const { + return __hadd2(lhs, rhs); + } +}; + +template<> +struct minus<__half2> { + MCTLASS_HOST_DEVICE + __half2 operator()(__half2 lhs, __half2 const &rhs) const { + return __hsub2(lhs, rhs); + } +}; + +template<> +struct multiplies<__half2> { + MCTLASS_HOST_DEVICE + __half2 operator()(__half2 lhs, __half2 const &rhs) const { + return __hmul2(lhs, rhs); + } +}; + +/// Partial specializations needed when __MACA_NO_HALF_OPERATORS__ is set +template<> +struct plus<__half> { + MCTLASS_HOST_DEVICE + __half operator()(__half lhs, __half const &rhs) const { + return __hadd(lhs, rhs); + } +}; + +template<> +struct minus<__half> { + MCTLASS_HOST_DEVICE + __half operator()(__half lhs, __half const &rhs) const { + return __hsub(lhs, rhs); + } +}; + +template<> +struct multiplies<__half> { + MCTLASS_HOST_DEVICE + __half operator()(__half lhs, __half const &rhs) const { + return __hmul(lhs, rhs); + } +}; +#endif // defined(__MACA_ARCH__) + + +// Maximum with nan propogation +// To propgate the NANs, the "max" of a two element that contains NaNs should also return a NaN +template +struct maximum_with_nan_propogation { + MCTLASS_HOST_DEVICE + T operator()(T const &lhs, T const &rhs) const { +#if defined(__MACA_ARCH__) + return lhs > rhs or isnan(lhs) ? lhs : rhs; +#else + return lhs > rhs or std::isnan(lhs) ? lhs : rhs; +#endif + } +}; + +template <> +struct maximum_with_nan_propogation { + MCTLASS_HOST_DEVICE + float operator()(float const lhs, float const rhs) const { + float res; +//#if defined(__MACA_ARCH__) && (__MACA_ARCH__ >= 800) +// asm volatile("max.NaN.f32 %0, %1, %2;\n" : "=f"(res) : "f"(lhs), "f"(rhs)); +//#elif defined(__MACA_ARCH__) +#if defined(__MACA_ARCH__) + res = lhs > rhs or isnan(lhs) ? lhs : rhs; +#else + res = lhs > rhs or std::isnan(lhs) ? lhs : rhs; +#endif + return res; + } +}; + +/// Squares with optional conversion +template +struct square { + MCTLASS_HOST_DEVICE + Output operator()(T lhs) const { + multiplies mul_op; + + Output y = Output(lhs); + return mul_op(y, y); + } +}; + +/// Returns the magnitude squared of an element. +template +struct magnitude_squared { + MCTLASS_HOST_DEVICE + Output operator()(T lhs) const { + multiplies mul_op; + + Output y = Output(lhs); + return mul_op(y, y); + } +}; + +/// Computes the square of a difference with optional conversion +template +struct square_difference { + MCTLASS_HOST_DEVICE + Output operator()(T lhs, T rhs) const { + multiplies mul_op; + + Output y = Output(lhs) - Output(rhs); + return mul_op(y, y); + } +}; + +/// Computes the square of a difference with optional conversion +template +struct magnitude_squared_difference { + MCTLASS_HOST_DEVICE + Output operator()(T lhs, T rhs) const { + multiplies mul_op; + + Output y = Output(lhs) - Output(rhs); + return mul_op(y, y); + } +}; + +template +struct divides { + MCTLASS_HOST_DEVICE + T operator()(T lhs, T const &rhs) const { + lhs /= rhs; + return lhs; + } +}; + +/// Negate +template +struct negate { + MCTLASS_HOST_DEVICE + T operator()(T lhs) const { + return -lhs; + } +}; + +/// Greater equal +template +struct greater_equal { + MCTLASS_HOST_DEVICE + bool operator()(T const &lhs, T const &rhs) const { + return (lhs >= rhs); + } +}; + +/// Greater +template +struct greater { + MCTLASS_HOST_DEVICE + bool operator()(T const &lhs, T const &rhs) const { + return (lhs > rhs); + } +}; + +/// Less equal +template +struct less_equal { + MCTLASS_HOST_DEVICE + bool operator()(T const &lhs, T const &rhs) const { + return (lhs <= rhs); + } +}; + +/// Less +template +struct less { + MCTLASS_HOST_DEVICE + bool operator()(T const &lhs, T const &rhs) const { + return (lhs < rhs); + } +}; + +template +struct maximum { + + MCTLASS_HOST_DEVICE + T operator()(T const &lhs, T const &rhs) const { + return (lhs < rhs ? rhs : lhs); + } +}; + +template <> +struct maximum { + MCTLASS_HOST_DEVICE + float operator()(float const &lhs, float const &rhs) const { + return fmaxf(lhs, rhs); + } +}; + +template +struct minimum { + + MCTLASS_HOST_DEVICE + T operator()(T const &lhs, T const &rhs) const { + return (rhs < lhs ? rhs : lhs); + } +}; + +template <> +struct minimum { + MCTLASS_HOST_DEVICE + float operator()(float const &lhs, float const &rhs) const { + return fminf(lhs, rhs); + } +}; + +/// Fused multiply-add +template +struct multiply_add { + MCTLASS_HOST_DEVICE + C operator()(A const &a, B const &b, C const &c) const { + return C(a) * C(b) + c; + } +}; + +/// Fused multiply-add +template +struct multiply_add_relu0 { + MCTLASS_HOST_DEVICE + C operator()(A const &a, B const &b, C const &c) const { + maximum mx; + return mx(C(a) * C(b) + c, C(0)); + } +}; + +/// Fused multiply-add +template +struct and_add { + MCTLASS_HOST_DEVICE + T operator()(T const &a, T const &b, T const &c) const { + return ((a & b) + c); + } +}; + + +/// Fused multiply-add +template +struct xor_add { + MCTLASS_HOST_DEVICE + T operator()(T const &a, T const &b, T const &c) const { + return ((a ^ b) + c); + } +}; + +template +struct conjugate { + MCTLASS_HOST_DEVICE + T operator()(T const &a) const { + return a; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct logical_and { + MCTLASS_HOST_DEVICE + T operator()(T const &a, T const &b) const { + return ((a && b) ? T(1) : T()); + } +}; + +template +struct logical_or { + MCTLASS_HOST_DEVICE + T operator()(T const &a, T const &b) const { + return ((a || b) ? T(1) : T()); + } +}; + +template +struct logical_not { + MCTLASS_HOST_DEVICE + T operator()(T const &a) const { + return T(!(a)); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct bit_and { + MCTLASS_HOST_DEVICE + T operator()(T const &a, T const &b) const { + return a & b; + } +}; + +template +struct bit_or { + MCTLASS_HOST_DEVICE + T operator()(T const &a, T const &b) const { + return a | b; + } +}; + +template +struct bit_not { + MCTLASS_HOST_DEVICE + T operator()(T const &a) const { + return ~a; + } +}; + +template +struct bit_xor { + MCTLASS_HOST_DEVICE + T operator()(T const &a, T const &b) const { + return a ^ b; + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Reduces value into the data pointed to by ptr +template +struct red +{ + MCTLASS_DEVICE + void operator()(T *ptr, const T &data) + { + atomicAdd(ptr, data); + } +}; + + +/// Reduces value into the data pointed to by ptr (double specialization) +template<> +struct red +{ + MCTLASS_DEVICE + void operator()(double *ptr, const double &data) + { +#if !defined(__MACA_ARCH__) + MCTLASS_UNUSED(ptr); + MCTLASS_UNUSED(data); +#elif (__MACA_ARCH__ >= 600) + + atomicAdd(ptr, data); + +#else + + // Use CAS loop + unsigned long long int* ptr_int = reinterpret_cast(ptr); + unsigned long long int old_int = *ptr_int; + unsigned long long int assumed_int; + + do { + double update = data + __longlong_as_double(old_int); + assumed_int = old_int; + old_int = atomicCAS(ptr_int, assumed_int, __double_as_longlong(update)); + } while (assumed_int != old_int); + +#endif // (__MACA_ARCH__ >= 600) + } +}; + + +/// Reduces value into the data pointed to by ptr (half2 specialization) +template<> +struct red +{ + MCTLASS_DEVICE + void operator()(half2 *ptr, const half2 &data) + { +#if defined(__MACA_ARCH__) + MCTLASS_UNUSED(ptr); + MCTLASS_UNUSED(data); +#else + + // Vector-2 atomic reduction requires .target sm_60 or higher + uint32_t word = reinterpret_cast(data); + asm volatile ("red.gpu.global.add.noftz.f16x2 [%0], %1;\n" : : "l"(ptr), "r"(word)); + +#endif // (__MACA_ARCH__ >= 600) + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations for mxmaca::wmma::fragment +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) + +template +struct plus> +{ + using Fragment = mxmaca::wmma::fragment; + using ElementType = typename Fragment::element_type; + + MCTLASS_HOST_DEVICE + Fragment operator()(Fragment const &lhs, Fragment const &rhs) const + { + Fragment result; + plus scalar_op; + + ElementType *result_elts = reinterpret_cast(&result); + const ElementType *lhs_elts = reinterpret_cast(&lhs); + const ElementType *rhs_elts = reinterpret_cast(&rhs); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Fragment::num_elements; i++) { + result_elts[i] = scalar_op(lhs_elts[i], rhs_elts[i]); + } + + return result; + } +}; + +#endif // defined(MCTLASS_ARCH_WMMA_ENABLED) + + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/collective/builders/sm90_gmma_builder.inl b/csrc/mctlass/include/mctlass/gemm/collective/builders/sm90_gmma_builder.inl new file mode 100644 index 0000000..d0af4b2 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/collective/builders/sm90_gmma_builder.inl @@ -0,0 +1,832 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/arch/mma.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/dispatch_policy.hpp" +#include "cute/atom/mma_traits_sm90_gmma.hpp" +#include "cute/atom/copy_traits_sm90_tma.hpp" + +// SM90 Collective Builders should be used only starting CUDA 12.0 +#if (__CUDACC_VER_MAJOR__ >= 12) +#define MCTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +// +// Some named constants +// +constexpr int tma_alignment_bytes = 16; +constexpr int cp_async_min_alignment_bytes = 4; +constexpr int sm90_smem_capacity_bytes = 232448; + +// Maps 2.x A matrix layout tag to respective GMMA major mode enum +template +constexpr cute::GMMA::Major +gmma_ss_tag_to_major_A() { + // MN major mode is only valid for non-TF32, non-int + if constexpr (mctlass::gemm::detail::is_mn_major_A() && + not cute::is_same_v && + sizeof(ElementA) != 1) { + return cute::GMMA::Major::MN; + } + else { + return cute::GMMA::Major::K; + } +} + +// Maps 2.x B matrix layout tag to respective GMMA major mode enum +template +constexpr cute::GMMA::Major +gmma_ss_tag_to_major_B() { + // MN major mode is only valid for non-TF32, non-int + if constexpr (mctlass::gemm::detail::is_mn_major_B() && + not cute::is_same_v && + sizeof(ElementB) != 1) { + return cute::GMMA::Major::MN; + } + else { + return cute::GMMA::Major::K; + } +} + +template +constexpr cute::GMMA::Major +gmma_rs_tag_to_major_A() { + // MN major mode is only valid for non-TF32 and non-int MMAs + if constexpr (mctlass::gemm::detail::is_mn_major_A()) { + return cute::GMMA::Major::MN; + } + else { + return cute::GMMA::Major::K; + } +} + +template +constexpr cute::GMMA::Major +gmma_rs_tag_to_major_B() { + // MN major mode is only valid for non-TF32 and non-int MMAs + if constexpr (mctlass::gemm::detail::is_mn_major_B()) { + return cute::GMMA::Major::MN; + } + else { + return cute::GMMA::Major::K; + } +} +// Maps a rank-1 cute::Shape<> representing the cluster shape on to the TMA atom that should be used with it +template +constexpr auto +sm90_cluster_shape_to_tma_atom(UnimodalClusterShape unimodal_cluster_shape) { + static_assert(cute::rank(unimodal_cluster_shape) == 1, + "Use this function to figure out TMA for each mode individually."); + + if constexpr (cute::size(unimodal_cluster_shape) == 1) { + return cute::SM90_TMA_LOAD{}; + } + else { + return cute::SM90_TMA_LOAD_MULTICAST{}; + } +} + +// Generates the most efficient possible TiledCopy with cp.async copy atom given a set of parameters. +template +constexpr auto +make_cp_async_gmem_tiled_copy() { + using AlignmentType = cute::uint_byte_t(sizeof(Element)) * Alignment>; + constexpr int TileSizeMN = cute::size(TileMN{}); + constexpr int TileSizeK = cute::size(TileK{}); + + // Maximize the number of threads along the gmem major mode to promote coalesced reads + // While making sure our thread layout tiles the threadblock tile evenly + + if constexpr (mctlass::gemm::detail::is_k_major()) { + // K major thread layout for K major gmem + constexpr int threads_major = TileSizeK / Alignment; + constexpr int threads_minor = ThreadCount / threads_major; + static_assert(threads_major > 0); + static_assert(ThreadCount % threads_major == 0); + static_assert(threads_minor == 0 || (TileSizeMN % threads_minor == 0)); + return make_tiled_copy( + Copy_Atom, Element>{}, + Layout,Int>, + Stride, _1>>{}, + Layout>>{}); + } + else if constexpr (mctlass::gemm::detail::is_mn_major()) { + // MN major thread layout for MN major gmem + constexpr int threads_major = TileSizeMN / Alignment; + constexpr int threads_minor = ThreadCount / threads_major; + static_assert(threads_major > 0); + static_assert(ThreadCount % threads_major == 0); + static_assert(threads_minor == 0 || (TileSizeK % threads_minor == 0)); + return make_tiled_copy( + Copy_Atom, Element>{}, + Layout,Int>, + Stride< _1,Int>>{}, + Layout,_1>>{}); + } + else { + static_assert(cute::is_void_v, "Unsupported gmem layout for automatic gmem tiled copy builder."); + } +} + + +// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count. +template +constexpr int +compute_stage_count_or_override(StageCount stage_count) { + return stages; +} + +// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count. +template +constexpr int +compute_stage_count_or_override(cute::integral_constant stage_count) { + return stages; +} + +// Returns the maximum number of smem tiles that can be used with a given smem capacity, or overrides with manual count. +template +constexpr int +compute_stage_count_or_override(StageCountAutoCarveout stage_count) { + // 32 bytes to account for barriers etc. + constexpr int stage_barrier_bytes = 32; + constexpr int a_bytes = static_cast(sizeof(ElementA)); + constexpr int b_bytes = static_cast(sizeof(ElementB)); + constexpr int stage_bytes = + (a_bytes * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + + (b_bytes * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) + + stage_barrier_bytes; + + return (CapacityBytes - carveout_bytes) / stage_bytes; +} + +// Helper for SS GMMA smem selection that considers a tensor TileShape: +// (BLK_MN, BLK_K) +// or hierarchically +// ((BLK_MN0,BLK_MN1,...),(BLK_K0,BLK_K1,...)) +// and returns the optimal GMMA::Layout that fits BLK_MN0 and BLK_K0 +template +constexpr auto +rs_smem_selector() { + auto BLK_MN0 = size<0>(BLK_MN{}); + auto BLK_K0 = size<0>(BLK_K{}); + + static_assert(BLK_MN0 % 8 == 0, "BLK_MN0 must be a multiple of 8."); + static_assert(BLK_K0 % 8 == 0, "BLK_K0 must be a multiple of 8."); + if constexpr (major == GMMA::Major::MN) { + if constexpr (sizeof(ElementType) == 4){ + if constexpr (is_ws_transposed_B) { + // only optimized transpositionB(SW32 and SW128 for tf32) can be used, but prefer SW32 due to free bank conflict + if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_SW32_Atom{}) == 0) { + return GMMA::Layout_MN_SW32_Atom{}; + } + else { + static_assert(BLK_MN0 % size<0>(GMMA::Layout_MN_SW32_Atom{}) == 0, + "BLK_MN0 must be a multiple of size<0>(GMMA::Layout_MN_SW32_Atom{})"); + } + } + else { + // Fall into SW32 due to free bank conflict + if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_SW32_Atom{}) == 0) { + return GMMA::Layout_MN_SW32_Atom{}; + } + else if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_INTER_Atom{}) == 0) { + return GMMA::Layout_MN_INTER_Atom{}; + } + else { + static_assert(BLK_MN0 % size<0>(GMMA::Layout_MN_INTER_Atom{}) == 0, + "BLK_MN0 must be a multiple of size<0>(GMMA::Layout_MN_INTER_Atom{})"); + } + } + } + // Used for int8, fp16 and bf16 I/O kernels + else if constexpr (sizeof(ElementType) == 1 || sizeof(ElementType) == 2) { + if constexpr (sizeof(ElementType) == 1 && is_ws_transposed_B) { + // Only optimized transpositionB (SW32 for int8) can be used + if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_SW128_Atom{}) == 0) { + return GMMA::Layout_MN_SW128_Atom{}; + } + else { + static_assert(BLK_MN0 % size<0>(GMMA::Layout_MN_SW128_Atom{}) == 0, + "BLK_MN0 must be a multiple of size<0>(GMMA::Layout_MN_128_Atom{})"); + } + } + else { + if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_SW128_Atom{}) == 0) { + return GMMA::Layout_MN_SW128_Atom{}; + } + else if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_SW64_Atom{}) == 0) { + return GMMA::Layout_MN_SW64_Atom{}; + } + else if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_SW32_Atom{}) == 0) { + return GMMA::Layout_MN_SW32_Atom{}; + } + else if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_INTER_Atom{}) == 0) { + return GMMA::Layout_MN_INTER_Atom{}; + } + else { + static_assert(BLK_MN0 % size<0>(GMMA::Layout_MN_INTER_Atom{}) == 0, + "BLK_MN0 must be a multiple of size<0>(GMMA::Layout_MN_INTER_Atom{})"); + } + } + } + else { + static_assert(mctlass::detail::dependent_false, "Smem selector does not support this element type"); + } + } + else if constexpr (major == GMMA::Major::K) { + if constexpr (BLK_K0 % size<1>(GMMA::Layout_K_SW128_Atom{}) == 0) { + return GMMA::Layout_K_SW128_Atom{}; + } + else if constexpr (BLK_K0 % size<1>(GMMA::Layout_K_SW64_Atom{}) == 0) { + return GMMA::Layout_K_SW64_Atom{}; + } + else if constexpr (BLK_K0 % size<1>(GMMA::Layout_K_SW32_Atom{}) == 0) { + return GMMA::Layout_K_SW32_Atom{}; + } + else if constexpr (BLK_K0 % size<1>(GMMA::Layout_K_INTER_Atom{}) == 0) { + return GMMA::Layout_K_INTER_Atom{}; + } + else { + static_assert(BLK_K0 % size<1>(GMMA::Layout_K_INTER_Atom{}) == 0, + "BLK_K0 must be a multiple of size<1>(GMMA::Layout_K_INTER_Atom{})"); + } + } +} + +// Helper for SS GMMA smem selection that considers a tensor TileShape: +// (BLK_MN, BLK_K) +// or hierarchically +// ((BLK_MN0,BLK_MN1,...),(BLK_K0,BLK_K1,...)) +// and returns the largest GMMA::Layout that fits BLK_MN0 and BLK_K0 +template +CUTE_HOST_DEVICE constexpr +auto +ss_smem_selector() +{ + auto BLK_MN0 = size<0>(BLK_MN{}); + auto BLK_K0 = size<0>(BLK_K{}); + + static_assert(BLK_MN0 % 8 == 0, "BLK_MN0 must be a multiple of 8."); + static_assert(BLK_K0 % 8 == 0, "BLK_K0 must be a multiple of 8."); + + + if constexpr (major == GMMA::Major::MN) { + if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_SW128_Atom{}) == 0) { + return GMMA::Layout_MN_SW128_Atom{}; + } + else if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_SW64_Atom{}) == 0) { + return GMMA::Layout_MN_SW64_Atom{}; + } + else if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_SW32_Atom{}) == 0) { + return GMMA::Layout_MN_SW32_Atom{}; + } + else if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_INTER_Atom{}) == 0) { + return GMMA::Layout_MN_INTER_Atom{}; + } + else { + static_assert(BLK_MN0 % size<0>(GMMA::Layout_MN_INTER_Atom{}) == 0, + "BLK_MN0 must be a multiple of size<0>(GMMA::Layout_MN_INTER_Atom{})"); + } + } + else if constexpr (major == GMMA::Major::K) { + if constexpr (BLK_K0 % size<1>(GMMA::Layout_K_SW128_Atom{}) == 0) { + return GMMA::Layout_K_SW128_Atom{}; + } + else if constexpr (BLK_K0 % size<1>(GMMA::Layout_K_SW64_Atom{}) == 0) { + return GMMA::Layout_K_SW64_Atom{}; + } + else if constexpr (BLK_K0 % size<1>(GMMA::Layout_K_SW32_Atom{}) == 0) { + return GMMA::Layout_K_SW32_Atom{}; + } + else if constexpr (BLK_K0 % size<1>(GMMA::Layout_K_INTER_Atom{}) == 0) { + return GMMA::Layout_K_INTER_Atom{}; + } + else { + static_assert(BLK_K0 % size<1>(GMMA::Layout_K_INTER_Atom{}) == 0, + "BLK_K0 must be a multiple of size<1>(GMMA::Layout_K_INTER_Atom{})"); + } + } +} + +template +constexpr bool +is_input_size_two_bytes() { + return (sizeof(ElementA) == 2 && sizeof(ElementB) == 2); +} + +template +constexpr bool +is_use_rmem_A() { + constexpr bool IsInputSizeTwoBytes = is_input_size_two_bytes(); + constexpr bool IsLayoutAkBk = mctlass::gemm::detail::is_k_major_A() && + mctlass::gemm::detail::is_k_major_B(); + constexpr bool IsUseRmemA = !IsInputSizeTwoBytes && !IsLayoutAkBk; + return IsUseRmemA; +} + +template +constexpr bool +is_swapAB(){ + constexpr bool IsInputSizeTwoBytes = is_input_size_two_bytes(); + constexpr bool IsLayoutAkBmn = mctlass::gemm::detail::is_k_major_A() && + mctlass::gemm::detail::is_mn_major_B(); + constexpr bool SwapAB = !IsInputSizeTwoBytes && IsLayoutAkBmn; + return SwapAB; +} + +template +constexpr bool +is_aligned() { + return ((sizeof(ElementA) * AlignmentA) % RequiredAlignment == 0) && + ((sizeof(ElementB) * AlignmentB) % RequiredAlignment == 0); +} + +template +constexpr bool +is_warpspecialized_transpose_B(){ + constexpr bool IsInputSizeTwoBytes = is_input_size_two_bytes(); + constexpr bool IsLayoutAmnBmn = mctlass::gemm::detail::is_mn_major_A() && + mctlass::gemm::detail::is_mn_major_B(); + constexpr bool IsWarpSpecialized = cute::is_base_of_v || + cute::is_base_of_v || + cute::is_base_of_v; + constexpr bool IsWarpSpecializedTransposeB = !IsInputSizeTwoBytes && IsLayoutAmnBmn && IsWarpSpecialized; + return IsWarpSpecializedTransposeB; +} + +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA_TMA_WS_SS +template < + class ElementA, + class GmemLayoutA, + int AlignmentA, + class ElementB, + class GmemLayoutB, + int AlignmentB, + class ElementAccumulator, + class TileShape_MNK, + class ClusterShape_MNK, + class StageCountType, + class KernelScheduleType +> +struct CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + ElementA, + GmemLayoutA, + AlignmentA, + ElementB, + GmemLayoutB, + AlignmentB, + ElementAccumulator, + TileShape_MNK, + ClusterShape_MNK, + StageCountType, + KernelScheduleType, + cute::enable_if_t< + (cute::is_same_v || + cute::is_same_v || + cute::is_same_v) && + not detail::is_use_rmem_A()> +> { + static_assert(is_static::value); + static_assert(is_static::value); +#ifndef MCTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED + static_assert(mctlass::detail::dependent_false == 0, "Unsupported Toolkit for SM90 Collective Builder\n"); +#endif + static_assert(detail::is_aligned(), + "Should meet TMA alignment requirement\n"); + + // For fp32 types, map to tf32 MMA value type + using MmaElementA = cute::conditional_t, tfloat32_t, ElementA>; + using MmaElementB = cute::conditional_t, tfloat32_t, ElementB>; + + static constexpr cute::GMMA::Major GmmaMajorA = detail::gmma_ss_tag_to_major_A(); + static constexpr cute::GMMA::Major GmmaMajorB = detail::gmma_ss_tag_to_major_B(); + + using AtomLayoutMNK = cute::conditional_t, + Layout>, Layout>>; + + using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::ss_op_selector< + MmaElementA, MmaElementB, ElementAccumulator, TileShape_MNK, GmmaMajorA, GmmaMajorB>(), AtomLayoutMNK{})); + + using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<1>(ClusterShape_MNK{}))); + using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<0>(ClusterShape_MNK{}))); + + using SmemLayoutAtomA = decltype(detail::ss_smem_selector< + GmmaMajorA, MmaElementA, decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>()); + using SmemLayoutAtomB = decltype(detail::ss_smem_selector< + GmmaMajorB, MmaElementB, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>()); + + static constexpr int PipelineStages = detail::compute_stage_count_or_override(StageCountType{}); + using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecialized< + PipelineStages, ClusterShape_MNK, KernelScheduleType>; + + using SmemCopyAtomA = void; + using SmemCopyAtomB = void; + + using CollectiveOp = CollectiveMma< + DispatchPolicy, + TileShape_MNK, + ElementA, + TagToStrideA_t, + ElementB, + TagToStrideB_t, + TiledMma, + GmemTiledCopyA, + SmemLayoutAtomA, + SmemCopyAtomA, + cute::identity, + GmemTiledCopyB, + SmemLayoutAtomB, + SmemCopyAtomB, + cute::identity + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA_TMA_WS_RS +template < + class ElementA, + class GmemLayoutA, + int AlignmentA, + class ElementB, + class GmemLayoutB, + int AlignmentB, + class ElementAccumulator, + class TileShape_MNK, + class ClusterShape_MNK, + class StageCountType, + class KernelScheduleType +> +struct CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + ElementA, + GmemLayoutA, + AlignmentA, + ElementB, + GmemLayoutB, + AlignmentB, + ElementAccumulator, + TileShape_MNK, + ClusterShape_MNK, + StageCountType, + KernelScheduleType, + cute::enable_if_t< + (cute::is_same_v || + cute::is_same_v || + cute::is_same_v) && + detail::is_use_rmem_A()> +> { + static_assert(is_static::value); + static_assert(is_static::value); + static_assert(detail::is_aligned(), + "Should meet TMA alignment requirement\n"); +#ifndef MCTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED + static_assert(mctlass::detail::dependent_false == 0, "Unsupported Toolkit for SM90 Collective Builder\n"); +#endif + static constexpr cute::GMMA::Major GmmaMajorA = detail::gmma_rs_tag_to_major_A(); + static constexpr cute::GMMA::Major GmmaMajorB = detail::gmma_rs_tag_to_major_B(); + static constexpr bool SwapAB = detail::is_swapAB(); + static constexpr bool IsWarpSpecializedTransposeB = detail::is_warpspecialized_transpose_B< + ElementA, GmemLayoutA, ElementB, GmemLayoutB, KernelScheduleType>(); + + // For fp32 types, map to tf32 MMA value type + using MmaElementA = cute::conditional_t, tfloat32_t, ElementA>; + using MmaElementB = cute::conditional_t, tfloat32_t, ElementB>; + + using AtomLayoutMNK = cute::conditional_t, + Layout>, Layout>>; + + using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::rs_op_selector< + MmaElementA, MmaElementB, ElementAccumulator, TileShape_MNK, GMMA::Major::K, GMMA::Major::K>(), AtomLayoutMNK{})); + + using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<1>(ClusterShape_MNK{}))); + using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<0>(ClusterShape_MNK{}))); + + using SmemLayoutAtomA = decltype(detail::rs_smem_selector(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{})), IsWarpSpecializedTransposeB>()); + using SmemLayoutAtomB = decltype(detail::rs_smem_selector(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{})), IsWarpSpecializedTransposeB>()); + + static constexpr int PipelineStages = detail::compute_stage_count_or_override(StageCountType{}); + + using DispatchPolicy = MainloopSm90TmaGmmaRmemAWarpSpecialized< + PipelineStages, ClusterShape_MNK, KernelScheduleType>; + + using SmemCopyAtomA = cute::conditional_t>; + using SmemCopyAtomB = cute::conditional_t, void>; + + using CollectiveOp = CollectiveMma< + DispatchPolicy, + TileShape_MNK, + ElementA, + TagToStrideA_t, + ElementB, + TagToStrideB_t, + TiledMma, + GmemTiledCopyA, + SmemLayoutAtomA, + SmemCopyAtomA, + cute::identity, + GmemTiledCopyB, + SmemLayoutAtomB, + SmemCopyAtomB, + cute::identity + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA_TMA_SS +template < + class ElementA, + class GmemLayoutA, + int AlignmentA, + class ElementB, + class GmemLayoutB, + int AlignmentB, + class ElementAccumulator, + class TileShape_MNK, + class ClusterShape_MNK, + class StageCountType, + class KernelScheduleType +> +struct CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + ElementA, + GmemLayoutA, + AlignmentA, + ElementB, + GmemLayoutB, + AlignmentB, + ElementAccumulator, + TileShape_MNK, + ClusterShape_MNK, + StageCountType, + KernelScheduleType, + cute::enable_if_t && + not detail::is_use_rmem_A()> +> { + static_assert(is_static::value); + static_assert(is_static::value); + static_assert(detail::is_aligned(), + "Should meet TMA alignment requirement\n"); +#ifndef MCTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED + static_assert(mctlass::detail::dependent_false == 0, "Unsupported Toolkit for SM90 Collective Builder\n"); +#endif + + // For fp32 types, map to tf32 MMA value type + using MmaElementA = cute::conditional_t, tfloat32_t, ElementA>; + using MmaElementB = cute::conditional_t, tfloat32_t, ElementB>; + + static constexpr cute::GMMA::Major GmmaMajorA = detail::gmma_ss_tag_to_major_A(); + static constexpr cute::GMMA::Major GmmaMajorB = detail::gmma_ss_tag_to_major_B(); + + using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::ss_op_selector< + MmaElementA, MmaElementB, ElementAccumulator, TileShape_MNK, GmmaMajorA, GmmaMajorB>())); + + using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<1>(ClusterShape_MNK{}))); + using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<0>(ClusterShape_MNK{}))); + + using SmemLayoutAtomA = decltype(detail::ss_smem_selector< + GmmaMajorA, MmaElementA, decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>()); + using SmemLayoutAtomB = decltype(detail::ss_smem_selector< + GmmaMajorB, MmaElementB, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>()); + + static constexpr int PipelineStages = detail::compute_stage_count_or_override(StageCountType{}); + using DispatchPolicy = MainloopSm90TmaGmma; + + using SmemCopyAtomA = void; + using SmemCopyAtomB = void; + + using CollectiveOp = CollectiveMma< + DispatchPolicy, + TileShape_MNK, + ElementA, + TagToStrideA_t, + ElementB, + TagToStrideB_t, + TiledMma, + GmemTiledCopyA, + SmemLayoutAtomA, + SmemCopyAtomA, + cute::identity, + GmemTiledCopyB, + SmemLayoutAtomB, + SmemCopyAtomB, + cute::identity + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA_CpAsync +template < + class ElementA, + class GmemLayoutA, + int AlignmentA, + class ElementB, + class GmemLayoutB, + int AlignmentB, + class ElementAccumulator, + class TileShape_MNK, + class ClusterShape_MNK, + class StageCountType, + class KernelScheduleType +> +struct CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + ElementA, + GmemLayoutA, + AlignmentA, + ElementB, + GmemLayoutB, + AlignmentB, + ElementAccumulator, + TileShape_MNK, + ClusterShape_MNK, + StageCountType, + KernelScheduleType, + cute::enable_if_t< + cute::is_same_v> +> { + static_assert(is_static::value); + static_assert(is_static::value); +#ifndef MCTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED + static_assert(mctlass::detail::dependent_false == 0, "Unsupported Toolkit for SM90 Collective Builder\n"); +#endif + + // For fp32 types, map to tf32 MMA value type + using MmaElementA = cute::conditional_t, tfloat32_t, ElementA>; + using MmaElementB = cute::conditional_t, tfloat32_t, ElementB>; + + static_assert(detail::is_aligned(), + "Minimum alignment required for cp.async is 4B."); + + static constexpr cute::GMMA::Major GmmaMajorA = detail::gmma_ss_tag_to_major_A(); + static constexpr cute::GMMA::Major GmmaMajorB = detail::gmma_ss_tag_to_major_B(); + + using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::ss_op_selector< + MmaElementA, MmaElementB, ElementAccumulator, TileShape_MNK, GmmaMajorA, GmmaMajorB>())); + + using GmemTiledCopyA = decltype(detail::make_cp_async_gmem_tiled_copy< + 128, ElementA, AlignmentA, TagToStrideA_t, + decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>()); + using GmemTiledCopyB = decltype(detail::make_cp_async_gmem_tiled_copy< + 128, ElementB, AlignmentB, TagToStrideB_t, + decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>()); + + using SmemLayoutAtomA = decltype(detail::ss_smem_selector< + GmmaMajorA, MmaElementA, decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>()); + using SmemLayoutAtomB = decltype(detail::ss_smem_selector< + GmmaMajorB, MmaElementB, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>()); + + static constexpr int PipelineStages = detail::compute_stage_count_or_override< + detail::sm90_smem_capacity_bytes, MmaElementA, MmaElementB, TileShape_MNK>(StageCountType{}); + + using CollectiveOp = CollectiveMma< + MainloopSm90CpAsyncGmma, + TileShape_MNK, + ElementA, + TagToStrideA_t, + ElementB, + TagToStrideB_t, + TiledMma, + GmemTiledCopyA, + SmemLayoutAtomA, + void, + cute::identity, + GmemTiledCopyB, + SmemLayoutAtomB, + void, + cute::identity + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// GMMA auto kernel schedule +template < + class ElementA, + class GmemLayoutA, + int AlignmentA, + class ElementB, + class GmemLayoutB, + int AlignmentB, + class ElementAccumulator, + class TileShape_MNK, + class ClusterShape_MNK, + class StageCountType, + class KernelScheduleType +> +struct CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + ElementA, + GmemLayoutA, + AlignmentA, + ElementB, + GmemLayoutB, + AlignmentB, + ElementAccumulator, + TileShape_MNK, + ClusterShape_MNK, + StageCountType, + KernelScheduleType, + cute::enable_if_t> +> { + static_assert(is_static::value); + static_assert(is_static::value); +#ifndef MCTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED + static_assert(mctlass::detail::dependent_false == 0, "Unsupported Toolkit for SM90 Collective Builder\n"); +#endif + +static constexpr bool IsTmaWarpSpecialized = detail::is_aligned< + ElementA, AlignmentA, ElementB, AlignmentB, detail::tma_alignment_bytes>(); + +#if ((__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 1))) + // Cooperative schedule performs best for CUDA Toolkits with version >= 12.1 + + // For TileShape_M == 64, choosing KernelTmaWarpSpecialized as the KernelSchedule + // Since KernelTmaWarpSpecializedCooperative requires TileShape_M to be at least 128 + using KernelWarpSpecializedSchedule = cute::conditional_t(TileShape_MNK{}) == Int<64>{}, + KernelTmaWarpSpecialized, KernelTmaWarpSpecializedCooperative>; +#else + using KernelWarpSpecializedSchedule = KernelTmaWarpSpecialized; +#endif + + using CollectiveOp = typename CollectiveBuilder< + arch::Sm90, + arch::OpClassTensorOp, + ElementA, + GmemLayoutA, + AlignmentA, + ElementB, + GmemLayoutB, + AlignmentB, + ElementAccumulator, + TileShape_MNK, + ClusterShape_MNK, + StageCountType, + cute::conditional_t + >::CollectiveOp; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/collective/collective_builder.hpp b/csrc/mctlass/include/mctlass/gemm/collective/collective_builder.hpp new file mode 100644 index 0000000..66261eb --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/collective/collective_builder.hpp @@ -0,0 +1,82 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +///////////////////////////////////////////////////////////////////////////////////////////////// +#include "collective_mma.hpp" + +namespace mctlass::gemm::collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Used to specify stage counts or dispatch to automatic computation of stage count +template +struct StageCount { static constexpr int value = num_stages; }; + +template +struct StageCountAutoCarveout { static constexpr int bytes = carveout_bytes; }; + +using StageCountAuto = StageCountAutoCarveout<0>; + +// Used to automatically let the builder pick the kernel schedule. +// Can be overridden with kernel schedule tags in mctlass/gemm/dispatch_policy.hpp +struct KernelScheduleAuto {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + class ArchTag, + class OpClass, + class ElementA, + class GmemLayoutA, + int AlignmentA, + class ElementB, + class GmemLayoutB, + int AlignmentB, + class ElementAccumulator, + class TileShape_MNK, + class ClusterShape_MNK, + class StageCountType, + class KernelScheduleType, + class Enable = void +> +struct CollectiveBuilder { + static_assert(sizeof(ElementA) == 0, "Could not build a collective for given parameters."); +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "builders/sm90_gmma_builder.inl" +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/collective/collective_mma.hpp b/csrc/mctlass/include/mctlass/gemm/collective/collective_mma.hpp new file mode 100644 index 0000000..4c35f83 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/collective/collective_mma.hpp @@ -0,0 +1,74 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/detail/dependent_false.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + class DispatchPolicy, + class TileShape, + class ElementA, + class StrideA, + class ElementB, + class StrideB, + class TiledMma, + class GmemTiledCopyA, + class SmemLayoutAtomA, + class SmemCopyAtomA, + class TransformA, + class GmemTiledCopyB, + class SmemLayoutAtomB, + class SmemCopyAtomB, + class TransformB +> +struct CollectiveMma { + static_assert(mctlass::detail::dependent_false == 0, "Could not find a mainloop specialization."); +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// #include "sm70_mma_twostage.hpp" +#include "sm80_mma_multistage.hpp" +// #include "sm90_mma_multistage_gmma_ss.hpp" +// #include "sm90_mma_tma_gmma_ss.hpp" +// #include "sm90_mma_tma_gmma_rs_warpspecialized.hpp" +// #include "sm90_mma_tma_gmma_ss_warpspecialized.hpp" +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/collective/sm70_mma_twostage.hpp b/csrc/mctlass/include/mctlass/gemm/collective/sm70_mma_twostage.hpp new file mode 100644 index 0000000..4184341 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/collective/sm70_mma_twostage.hpp @@ -0,0 +1,596 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/dispatch_policy.hpp" + +#include "cute/algorithm/functional.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/algorithm/gemm.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/tensor_predicate.hpp" + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::collective { +using namespace cute; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + class TileShape_, + class ElementA_, + class StrideA_, + class ElementB_, + class StrideB_, + class TiledMma_, + class GmemTiledCopyA_, + class SmemLayoutAtomA_, + class SmemCopyAtomA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomB_, + class SmemCopyAtomB_, + class TransformB_> +struct CollectiveMma< + MainloopSm70TwoStageUnpredicated, + TileShape_, + ElementA_, + StrideA_, + ElementB_, + StrideB_, + TiledMma_, + GmemTiledCopyA_, + SmemLayoutAtomA_, + SmemCopyAtomA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomB_, + SmemCopyAtomB_, + TransformB_> +{ + // + // Type Aliases + // + using DispatchPolicy = MainloopSm70TwoStageUnpredicated; + using TileShape = TileShape_; + using ElementA = ElementA_; + using StrideA = StrideA_; + using ElementB = ElementB_; + using StrideB = StrideB_; + using TiledMma = TiledMma_; + using ElementAccumulator = typename TiledMma::ValTypeC; + using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomA = SmemLayoutAtomA_; + using SmemLayoutAtomB = SmemLayoutAtomB_; + using SmemCopyAtomA = SmemCopyAtomA_; + using SmemCopyAtomB = SmemCopyAtomB_; + using TransformA = TransformA_; + using TransformB = TransformB_; + using ArchTag = typename DispatchPolicy::ArchTag; + + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + using SmemLayoutA = decltype(tile_to_shape( + SmemLayoutAtomA{}, + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})))); + using SmemLayoutB = decltype(tile_to_shape( + SmemLayoutAtomB{}, + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})))); + + struct SharedStorage + { + cute::array_aligned> smem_a; + cute::array_aligned> smem_b; + }; + + // Host side kernel arguments + struct Arguments { + ElementA const* ptr_A; + StrideA dA; + ElementB const* ptr_B; + StrideB dB; + }; + + // Device side kernel params + using Params = Arguments; + + // + // Methods + // + + CollectiveMma() = default; + + template + static constexpr Params + to_underlying_arguments(ProblemShape const& _, Arguments const& args, void* workspace) { + (void) workspace; + return args; + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + template < + class FrgTensorD, + class TensorA, + class TensorB, + class FrgTensorC, + class KTileIterator, + class ResidueMNK + > + MCTLASS_DEVICE void + operator() ( + FrgTensorD &accum, + TensorA gA, + TensorB gB, + FrgTensorC const &src_accum, + KTileIterator k_tile_iter, int k_tile_count, + ResidueMNK residue_mnk, + int thread_idx, + char *smem_buf) + { + using namespace cute; + + (void)residue_mnk; + + static_assert(is_rmem::value, "D tensor must be rmem resident."); + static_assert(is_gmem::value, "A tensor must be gmem resident."); + static_assert(is_gmem::value, "B tensor must be gmem resident."); + static_assert(is_rmem::value, "C tensor must be rmem resident."); + static_assert(rank(SmemLayoutA{}) == 2, + "MainloopTwoStage must not have a smem shape with a pipeline mode."); + static_assert(rank(SmemLayoutB{}) == 2, + "MainloopTwoStage must not have a smem shape with a pipeline mode."); + + // Construct shared memory tiles + SharedStorage& storage = *reinterpret_cast(smem_buf); + Tensor sA = make_tensor(make_smem_ptr(storage.smem_a.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(storage.smem_b.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + // Partition the copying of A and B tiles across the threads + GmemTiledCopyA gmem_tiled_copy_a; + GmemTiledCopyB gmem_tiled_copy_b; + auto copy_a_thr = gmem_tiled_copy_a.get_slice(thread_idx); + auto copy_b_thr = gmem_tiled_copy_b.get_slice(thread_idx); + + Tensor tAgA = copy_a_thr.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k) + Tensor tAsA = copy_a_thr.partition_D(sA); // (ACPY,ACPY_M,ACPY_K) + Tensor tBgB = copy_b_thr.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k) + Tensor tBsB = copy_b_thr.partition_D(sB); // (BCPY,BCPY_N,BCPY_K) + + // Allocate the register tiles for double buffering -- same shape as partitioned data + Tensor tArA = make_fragment_like(tAsA); // (ACPY,ACPY_M,ACPY_K) + Tensor tBrB = make_fragment_like(tBsB); // (BCPY,BCPY_N,BCPY_K) + + // Tile MMA compute thread partitions and allocate accumulators + TiledMma tiled_mma; + auto thr_mma = tiled_mma.get_thread_slice(thread_idx); + Tensor tCrA = thr_mma.partition_fragment_A(sA); // (MMA,MMA_M,MMA_K) + Tensor tCrB = thr_mma.partition_fragment_B(sB); // (MMA,MMA_M,MMA_K) + + CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(accum)); // MMA_M + CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(src_accum)); // MMA_M + CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(accum)); // MMA_N + CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(src_accum)); // MMA_N + CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K + + // + // Copy Atom retiling + // + + auto thr_copy_A = make_tiled_copy_A(SmemCopyAtomA{}, tiled_mma).get_thread_slice(thread_idx); + Tensor tCsA = thr_copy_A.partition_S(sA); + Tensor tCrA_copy_view = thr_copy_A.retile_D(tCrA); + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // M + + auto thr_copy_B = make_tiled_copy_B(SmemCopyAtomB{}, tiled_mma).get_thread_slice(thread_idx); + Tensor tCsB = thr_copy_B.partition_S(sB); + Tensor tCrB_copy_view = thr_copy_B.retile_D(tCrB); + CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // N + + // + // Prologue + // + + // Copy gmem to rmem for the first k_tile + copy(gmem_tiled_copy_a, tAgA(_,_,_,*k_tile_iter), tArA); + copy(gmem_tiled_copy_b, tBgB(_,_,_,*k_tile_iter), tBrB); + if (--k_tile_count > 0) ++k_tile_iter; + // Copy rmem to smem + copy(tArA, tAsA); + copy(tBrB, tBsB); + // Clear accumulators + __syncthreads(); + + // Load A, B smem->rmem for k=0 + copy(tCsA(_,_,0), tCrA_copy_view(_,_,0)); + copy(tCsB(_,_,0), tCrB_copy_view(_,_,0)); + // + // Mainloop + // + + // Size of the k-tiles's outer product mode (k) + auto K_BLOCK_MAX = size<2>(tCrA); + + MCTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > -1) + { + // Pipeline the outer products with a static for loop + for_each(make_int_sequence{}, [&] (auto k_block) + { + if (k_block == K_BLOCK_MAX - 1) + { + __syncthreads(); + + // Copy rmem to smem + copy(tArA, tAsA); + copy(tBrB, tBsB); + __syncthreads(); + } + + // Load A, B smem->rmem for k+1 + int k_block_next = (k_block + Int<1>{}) % K_BLOCK_MAX; // static + copy(tCsA(_,_,k_block_next), tCrA_copy_view(_,_,k_block_next)); + copy(tCsB(_,_,k_block_next), tCrB_copy_view(_,_,k_block_next)); + if (k_block == 0) + { + // Copy gmem to rmem + copy(gmem_tiled_copy_a, tAgA(_,_,_,*k_tile_iter), tArA); + copy(gmem_tiled_copy_b, tBgB(_,_,_,*k_tile_iter), tBrB); + if (--k_tile_count > 0) ++k_tile_iter; + } + + // transform before compute + cute::transform(tCrA(_,_,k_block), TransformA{}); + cute::transform(tCrB(_,_,k_block), TransformB{}); + + // Thread-level register gemm for k + // disambiguate gemm (shared with the namespace name) + cute::gemm(tiled_mma, accum, tCrA(_,_,k_block), tCrB(_,_,k_block), src_accum); + }); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + class TileShape_, + class ElementA_, + class StrideA_, + class ElementB_, + class StrideB_, + class TiledMma_, + class GmemTiledCopyA_, + class SmemLayoutAtomA_, + class SmemCopyAtomA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomB_, + class SmemCopyAtomB_, + class TransformB_> +struct CollectiveMma< + MainloopSm70TwoStage, + TileShape_, + ElementA_, + StrideA_, + ElementB_, + StrideB_, + TiledMma_, + GmemTiledCopyA_, + SmemLayoutAtomA_, + SmemCopyAtomA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomB_, + SmemCopyAtomB_, + TransformB_> +{ + // + // Type Aliases + // + using DispatchPolicy = MainloopSm70TwoStage; + using TileShape = TileShape_; + using ElementA = ElementA_; + using StrideA = StrideA_; + using ElementB = ElementB_; + using StrideB = StrideB_; + using TiledMma = TiledMma_; + using ElementAccumulator = typename TiledMma::ValTypeC; + using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomA = SmemLayoutAtomA_; + using SmemLayoutAtomB = SmemLayoutAtomB_; + using SmemCopyAtomA = SmemCopyAtomA_; + using SmemCopyAtomB = SmemCopyAtomB_; + using TransformA = TransformA_; + using TransformB = TransformB_; + using ArchTag = typename DispatchPolicy::ArchTag; + + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + using SmemLayoutA = decltype(tile_to_shape( + SmemLayoutAtomA{}, + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})))); + using SmemLayoutB = decltype(tile_to_shape( + SmemLayoutAtomB{}, + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})))); + + struct SharedStorage + { + cute::array_aligned> smem_a; + cute::array_aligned> smem_b; + }; + + // Host side kernel arguments + struct Arguments { + ElementA const* ptr_A; + StrideA dA; + ElementB const* ptr_B; + StrideB dB; + }; + + // Device side kernel params + using Params = Arguments; + + // + // Methods + // + + CollectiveMma() = default; + + template + static constexpr Params + to_underlying_arguments(ProblemShape const& _, Arguments const& args, void* workspace) { + (void) workspace; + return args; + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + template < + class FrgTensorD, + class TensorA, + class TensorB, + class FrgTensorC, + class KTileIterator, + class ResidueMNK + > + MCTLASS_DEVICE void + operator() ( + FrgTensorD &accum, + TensorA gA, + TensorB gB, + FrgTensorC const &src_accum, + KTileIterator k_tile_iter, int k_tile_count, + ResidueMNK residue_mnk, + int thread_idx, + char *smem_buf) + { + using namespace cute; + + static_assert(is_rmem::value, "D tensor must be rmem resident."); + static_assert(is_gmem::value, "A tensor must be gmem resident."); + static_assert(is_gmem::value, "B tensor must be gmem resident."); + static_assert(is_rmem::value, "C tensor must be rmem resident."); + static_assert(rank(SmemLayoutA{}) == 2, + "MainloopTwoStage must not have a smem shape with a pipeline mode."); + static_assert(rank(SmemLayoutB{}) == 2, + "MainloopTwoStage must not have a smem shape with a pipeline mode."); + + // Construct shared memory tiles + SharedStorage& storage = *reinterpret_cast(smem_buf); + Tensor sA = make_tensor(make_smem_ptr(storage.smem_a.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(storage.smem_b.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + // Shift tensor so residue_k is at origin (Can't read any k_coord < residue_k) + // This aligns the tensor with BLK_K for all but the 0th k_tile + gA.data() = &gA(0, get<2>(residue_mnk), 0); + gB.data() = &gB(0, get<2>(residue_mnk), 0); + + // Partition the copying of A and B tiles across the threads + GmemTiledCopyA gmem_tiled_copy_a; + GmemTiledCopyB gmem_tiled_copy_b; + auto gmem_thr_copy_a = gmem_tiled_copy_a.get_slice(thread_idx); + auto gmem_thr_copy_b = gmem_tiled_copy_b.get_slice(thread_idx); + + Tensor tAgA = gmem_thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k) + Tensor tAsA = gmem_thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE) + Tensor tBgB = gmem_thr_copy_b.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k) + Tensor tBsB = gmem_thr_copy_b.partition_D(sB); // (BCPY,BCPY_N,BCPY_K,PIPE) + + // Allocate the register tiles for double buffering -- same shape as partitioned data + Tensor tArA = make_fragment_like(tAsA); // (ACPY,ACPY_M,ACPY_K) + Tensor tBrB = make_fragment_like(tBsB); // (BCPY,BCPY_N,BCPY_K) + + // + // PREDICATES + // + + // Allocate predicate tensors for m and n + Tensor tApA = make_tensor(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{}); + Tensor tBpB = make_tensor(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{}); + + // Construct identity layout for sA and sB + Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + + // Repeat the partitioning with identity layouts + Tensor tAcA = gmem_thr_copy_a.partition_S(cA); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tBcB = gmem_thr_copy_b.partition_S(cB); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + + // Set predicates for m bounds + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < size<0>(tApA); ++m) { + tApA(m,0) = get<0>(tAcA(0,m,0)) < get<0>(residue_mnk); // blk_m coord < residue_m + } + // Set predicates for n bounds + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < size<0>(tBpB); ++n) { + tBpB(n,0) = get<0>(tBcB(0,n,0)) < get<1>(residue_mnk); // blk_n coord < residue_n + } + + // + // PREFETCH + // + + // Clear the rmem tiles to account for predicated off loads + clear(tArA); + clear(tBrB); + + // Start async loads for 0th k-tile, where we take care of the k residue + { + Tensor tAgAk = tAgA(_,_,_,*k_tile_iter); + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < size<2>(tArA); ++k) { + if (get<1>(tAcA(0,0,k)) >= -get<2>(residue_mnk)) { // blk_k coord < residue_k (gA shifted) + copy_if(gmem_tiled_copy_a, tApA(_,k), tAgAk(_,_,k), tArA(_,_,k)); + } + } + Tensor tBgBk = tBgB(_,_,_,*k_tile_iter); + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < size<2>(tBrB); ++k) { + if (get<1>(tBcB(0,0,k)) >= -get<2>(residue_mnk)) { // blk_k coord < residue_k (gB shifted) + copy_if(gmem_tiled_copy_b, tBpB(_,k), tBgBk(_,_,k), tBrB(_,_,k)); + } + } + ++k_tile_iter; + --k_tile_count; + } + + // Tile MMA compute thread partitions and allocate accumulators + TiledMma tiled_mma; + auto thr_mma = tiled_mma.get_thread_slice(thread_idx); + Tensor tCrA = thr_mma.make_fragment_A(thr_mma.partition_A(sA)); // (MMA,MMA_M,MMA_K) + Tensor tCrB = thr_mma.make_fragment_B(thr_mma.partition_B(sB)); // (MMA,MMA_M,MMA_K) + + CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(accum)); // MMA_M + CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(src_accum)); // MMA_M + CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(accum)); // MMA_N + CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(src_accum)); // MMA_N + CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K + + // + // Copy Atom retiling + // + + auto thr_copy_A = make_tiled_copy_A(SmemCopyAtomA{}, tiled_mma).get_thread_slice(thread_idx); + Tensor tCsA = thr_copy_A.partition_S(sA); + Tensor tCrA_copy_view = thr_copy_A.retile_D(tCrA); + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // M + + auto thr_copy_B = make_tiled_copy_B(SmemCopyAtomB{}, tiled_mma).get_thread_slice(thread_idx); + Tensor tCsB = thr_copy_B.partition_S(sB); + Tensor tCrB_copy_view = thr_copy_B.retile_D(tCrB); + CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // N + + // + // Prologue + // + + // Copy rmem to smem + copy(tArA, tAsA); + copy(tBrB, tBsB); + // Clear accumulators + __syncthreads(); + + // Load A, B smem->rmem for k=0 + copy(tCsA(_,_,0), tCrA_copy_view(_,_,0)); + copy(tCsB(_,_,0), tCrB_copy_view(_,_,0)); + // + // Mainloop + // + + // Size of the k-tiles's outer product mode (k) + auto K_BLOCK_MAX = size<2>(tCrA); + + MCTLASS_PRAGMA_NO_UNROLL + while (k_tile_count > -1) + { + // Pipeline the outer products with a static for loop + for_each(make_int_sequence{}, [&] (auto k_block) + { + if (k_block == K_BLOCK_MAX - 1) + { + __syncthreads(); + + // Copy rmem to smem + copy(tArA, tAsA); + copy(tBrB, tBsB); + __syncthreads(); + } + + // Load A, B smem->rmem for k+1 + int k_block_next = (k_block + Int<1>{}) % K_BLOCK_MAX; // static + copy(tCsA(_,_,k_block_next), tCrA_copy_view(_,_,k_block_next)); + copy(tCsB(_,_,k_block_next), tCrB_copy_view(_,_,k_block_next)); + if (k_block == 0) + { + if (k_tile_count <= 0) { + clear(tApA); + clear(tBpB); + } + copy_if(gmem_tiled_copy_a, tApA, tAgA(_,_,_,*k_tile_iter), tArA); + copy_if(gmem_tiled_copy_b, tBpB, tBgB(_,_,_,*k_tile_iter), tBrB); + ++k_tile_iter; + --k_tile_count; + } + + // transform before compute + cute::transform(tCrA(_,_,k_block), TransformA{}); + cute::transform(tCrB(_,_,k_block), TransformB{}); + + // Thread-level register gemm for k + // disambiguate gemm (shared with the namespace name) + cute::gemm(tiled_mma, accum, tCrA(_,_,k_block), tCrB(_,_,k_block), src_accum); + }); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/collective/sm80_mma_multistage.hpp b/csrc/mctlass/include/mctlass/gemm/collective/sm80_mma_multistage.hpp new file mode 100644 index 0000000..b2ca175 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/collective/sm80_mma_multistage.hpp @@ -0,0 +1,688 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/dispatch_policy.hpp" + +#include "cute/algorithm/functional.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/algorithm/gemm.hpp" +#include "cute/tensor_predicate.hpp" +#include "cute/numeric/arithmetic_tuple.hpp" + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::collective { +using namespace cute; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int Stages, + class TileShape_, + class ElementA_, + class StrideA_, + class ElementB_, + class StrideB_, + class TiledMma_, + class GmemTiledCopyA_, + class SmemLayoutAtomA_, + class SmemCopyAtomA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomB_, + class SmemCopyAtomB_, + class TransformB_> +struct CollectiveMma< + MainloopSm80CpAsyncUnpredicated, + TileShape_, + ElementA_, + StrideA_, + ElementB_, + StrideB_, + TiledMma_, + GmemTiledCopyA_, + SmemLayoutAtomA_, + SmemCopyAtomA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomB_, + SmemCopyAtomB_, + TransformB_> +{ + // + // Type Aliases + // + using DispatchPolicy = MainloopSm80CpAsyncUnpredicated; + using TileShape = TileShape_; + using ElementA = ElementA_; + using StrideA = StrideA_; + using ElementB = ElementB_; + using StrideB = StrideB_; + using TiledMma = TiledMma_; + using ElementAccumulator = typename TiledMma::ValTypeC; + using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomA = SmemLayoutAtomA_; + using SmemLayoutAtomB = SmemLayoutAtomB_; + using SmemCopyAtomA = SmemCopyAtomA_; + using SmemCopyAtomB = SmemCopyAtomB_; + using TransformA = TransformA_; + using TransformB = TransformB_; + using ArchTag = typename DispatchPolicy::ArchTag; + + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + using SmemLayoutA = decltype(tile_to_shape( + SmemLayoutAtomA{}, + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int{}))); + using SmemLayoutB = decltype(tile_to_shape( + SmemLayoutAtomB{}, + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}))); + + static_assert(DispatchPolicy::Stages >= 2, "CpAsync mainloop must have at least 2 stages in the pipeline."); + + struct SharedStorage + { + cute::array_aligned> smem_a; + cute::array_aligned> smem_b; + }; + + // Host side kernel arguments + struct Arguments { + ElementA const* ptr_A; + StrideA dA; + ElementB const* ptr_B; + StrideB dB; + }; + + // Device side kernel params + using Params = Arguments; + + // + // Methods + // + + CollectiveMma() = default; + + template + static constexpr Params + to_underlying_arguments(ProblemShape const& _, Arguments const& args, void* workspace) { + (void) workspace; + return args; + } + + /// Perform a collective-scoped matrix multiply-accumulate + template < + class FrgTensorD, + class TensorA, + class TensorB, + class FrgTensorC, + class KTileIterator, + class ResidueMNK + > + MCTLASS_DEVICE void + operator() ( + FrgTensorD &accum, + TensorA gA, + TensorB gB, + FrgTensorC const &src_accum, + KTileIterator k_tile_iter, int k_tile_count, + ResidueMNK residue_mnk, + int thread_idx, + char *smem_buf) + { + using namespace cute; + + static_assert(is_rmem::value, "D tensor must be rmem resident."); + static_assert(is_gmem::value, "A tensor must be gmem resident."); + static_assert(is_gmem::value, "B tensor must be gmem resident."); + static_assert(is_rmem::value, "C tensor must be rmem resident."); + static_assert(rank(SmemLayoutA{}) == 3, + "MainloopSm80CpAsync must have a pipeline mode in the smem layout."); + static_assert(rank(SmemLayoutB{}) == 3, + "MainloopSm80CpAsync must have a pipeline mode in the smem layout."); + + // Construct shared memory tiles + SharedStorage& storage = *reinterpret_cast(smem_buf); + Tensor sA = make_tensor(make_smem_ptr(storage.smem_a.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(storage.smem_b.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + CUTE_STATIC_ASSERT_V(size<0>(gA) == size<0>(sA)); // BLK_M + CUTE_STATIC_ASSERT_V(size<1>(gA) == size<1>(sA)); // BLK_K + CUTE_STATIC_ASSERT_V(size<0>(gB) == size<0>(sB)); // BLK_N + CUTE_STATIC_ASSERT_V(size<1>(gB) == size<1>(sB)); // BLK_K + CUTE_STATIC_ASSERT_V(size<1>(sA) == size<1>(sB)); // BLK_K + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sA)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sB)); // PIPE + + // Partition the copying of A and B tiles across the threads + GmemTiledCopyA gmem_tiled_copy_A; + GmemTiledCopyB gmem_tiled_copy_B; + auto gmem_thr_copy_A = gmem_tiled_copy_A.get_slice(thread_idx); + auto gmem_thr_copy_B = gmem_tiled_copy_B.get_slice(thread_idx); + + Tensor tAgA = gmem_thr_copy_A.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k) + Tensor tAsA = gmem_thr_copy_A.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE) + Tensor tBgB = gmem_thr_copy_B.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k) + Tensor tBsB = gmem_thr_copy_B.partition_D(sB); // (BCPY,BCPY_N,BCPY_K,PIPE) + + // + // PREDICATES + // + + (void) residue_mnk; + //assert(residue_mnk == make_tuple(0,0,0)); + + // + // PREFETCH + // + + // Start async loads for all pipes but the last + MCTLASS_PRAGMA_UNROLL + for (int k_pipe = 0; k_pipe < DispatchPolicy::Stages-1; ++k_pipe) { + copy(gmem_tiled_copy_A, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,k_pipe)); + copy(gmem_tiled_copy_B, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,k_pipe)); + cp_async_fence(); + --k_tile_count; + if (k_tile_count > 0) { ++k_tile_iter; } + } + + // + // MMA Atom partitioning + // + + // Tile MMA compute thread partitions and allocate accumulators + TiledMma tiled_mma; + auto thr_mma = tiled_mma.get_thread_slice(thread_idx); + Tensor tCrA = thr_mma.partition_fragment_A(sA(_,_,0)); // (MMA,MMA_M,MMA_K) + Tensor tCrB = thr_mma.partition_fragment_B(sB(_,_,0)); // (MMA,MMA_N,MMA_K) + + CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(accum)); // MMA_M + CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(src_accum)); // MMA_M + CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(accum)); // MMA_N + CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(src_accum)); // MMA_N + CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K + CUTE_STATIC_ASSERT_V(size(gmem_tiled_copy_A) == size(tiled_mma)); + CUTE_STATIC_ASSERT_V(size(gmem_tiled_copy_B) == size(tiled_mma)); + + // + // Copy Atom retiling + // + + auto smem_tiled_copy_A = make_tiled_copy_A(SmemCopyAtomA{}, tiled_mma); + auto smem_thr_copy_A = smem_tiled_copy_A.get_thread_slice(thread_idx); + Tensor tCsA = smem_thr_copy_A.partition_S(sA); // (CPY,CPY_M,CPY_K,PIPE) + Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA); // (CPY,CPY_M,CPY_K) + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // CPY_M + CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrA_copy_view)); // CPY_K + + auto smem_tiled_copy_B = make_tiled_copy_B(SmemCopyAtomB{}, tiled_mma); + auto smem_thr_copy_B = smem_tiled_copy_B.get_thread_slice(thread_idx); + Tensor tCsB = smem_thr_copy_B.partition_S(sB); // (CPY,CPY_N,CPY_K,PIPE) + Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB); // (CPY,CPY_N,CPY_K) + CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // CPY_N + CUTE_STATIC_ASSERT_V(size<2>(tCsB) == size<2>(tCrB_copy_view)); // CPY_K + + // + // PIPELINED MAIN LOOP + // + + // Current pipe index in smem to read from + int smem_pipe_read = 0; + // Current pipe index in smem to write to + int smem_pipe_write = DispatchPolicy::Stages-1; + + Tensor tCsA_p = tCsA(_,_,_,smem_pipe_read); + Tensor tCsB_p = tCsB(_,_,_,smem_pipe_read); + + // Size of the register pipeline + auto K_BLOCK_MAX = size<2>(tCrA); + + // PREFETCH register pipeline + if (K_BLOCK_MAX > 1) { + // Wait until our first prefetched tile is loaded in + cp_async_wait(); + __syncthreads(); + + // Prefetch the first rmem from the first k-tile + copy(smem_tiled_copy_A, tCsA_p(_,_,Int<0>{}), tCrA_copy_view(_,_,Int<0>{})); + copy(smem_tiled_copy_B, tCsB_p(_,_,Int<0>{}), tCrB_copy_view(_,_,Int<0>{})); + } + + + MCTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > -(DispatchPolicy::Stages-1); --k_tile_count) + { + // Pipeline the outer products with a static for loop. + // + // Note, the for_each() function is required here to ensure `k_block` is of type Int. + for_each(make_int_sequence{}, [&] (auto k_block) + { + if (k_block == K_BLOCK_MAX - 1) + { + // Slice the smem_pipe_read smem + tCsA_p = tCsA(_,_,_,smem_pipe_read); + tCsB_p = tCsB(_,_,_,smem_pipe_read); + + // Commit the smem for smem_pipe_read + cp_async_wait(); + __syncthreads(); + } + + // Load A, B shmem->regs for k_block+1 + auto k_block_next = (k_block + Int<1>{}) % K_BLOCK_MAX; // static + copy(smem_tiled_copy_A, tCsA_p(_,_,k_block_next), tCrA_copy_view(_,_,k_block_next)); + copy(smem_tiled_copy_B, tCsB_p(_,_,k_block_next), tCrB_copy_view(_,_,k_block_next)); + // Copy gmem to smem before computing gemm on each k-pipe + if (k_block == 0) + { + copy(gmem_tiled_copy_A, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,smem_pipe_write)); + copy(gmem_tiled_copy_B, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,smem_pipe_write)); + cp_async_fence(); + if (k_tile_count > 0) { ++k_tile_iter; } + + // Advance the pipe -- Doing it here accounts for K_BLOCK_MAX = 1 (no rmem pipe) + smem_pipe_write = smem_pipe_read; + ++smem_pipe_read; + smem_pipe_read = (smem_pipe_read == DispatchPolicy::Stages) ? 0 : smem_pipe_read; + } + + // Transform before compute + cute::transform(tCrA(_,_,k_block), TransformA{}); + cute::transform(tCrB(_,_,k_block), TransformB{}); + // Thread-level register gemm for k_block + cute::gemm(tiled_mma, accum, tCrA(_,_,k_block), tCrB(_,_,k_block), src_accum); + }); + + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int Stages, + class TileShape_, + class ElementA_, + class StrideA_, + class ElementB_, + class StrideB_, + class TiledMma_, + class GmemTiledCopyA_, + class SmemLayoutAtomA_, + class SmemCopyAtomA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomB_, + class SmemCopyAtomB_, + class TransformB_> +struct CollectiveMma< + MainloopSm80CpAsync, + TileShape_, + ElementA_, + StrideA_, + ElementB_, + StrideB_, + TiledMma_, + GmemTiledCopyA_, + SmemLayoutAtomA_, + SmemCopyAtomA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomB_, + SmemCopyAtomB_, + TransformB_> +{ + // + // Type Aliases + // + using DispatchPolicy = MainloopSm80CpAsync; + using TileShape = TileShape_; + using ElementA = ElementA_; + using StrideA = StrideA_; + using ElementB = ElementB_; + using StrideB = StrideB_; + using TiledMma = TiledMma_; + using ElementAccumulator = typename TiledMma::ValTypeC; using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomA = SmemLayoutAtomA_; + using SmemLayoutAtomB = SmemLayoutAtomB_; + using SmemCopyAtomA = SmemCopyAtomA_; + using SmemCopyAtomB = SmemCopyAtomB_; + using TransformA = TransformA_; + using TransformB = TransformB_; + using ArchTag = typename DispatchPolicy::ArchTag; + + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + using SmemLayoutA = decltype(tile_to_shape( + SmemLayoutAtomA{}, + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int{}))); + using SmemLayoutB = decltype(tile_to_shape( + SmemLayoutAtomB{}, + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}))); + + static_assert(DispatchPolicy::Stages >= 2, "CpAsync mainloop must have at least 2 stages in the pipeline."); + + struct SharedStorage + { + cute::array_aligned> smem_a; + cute::array_aligned> smem_b; + }; + + // Host side kernel arguments + struct Arguments { + ElementA const* ptr_A; + StrideA dA; + ElementB const* ptr_B; + StrideB dB; + }; + + // Device side kernel params + using Params = Arguments; + + // + // Methods + // + + CollectiveMma() = default; + + template + static constexpr Params + to_underlying_arguments(ProblemShape const& _, Arguments const& args, void* workspace) { + (void) workspace; + return args; + } + + /// Perform a collective-scoped matrix multiply-accumulate + template < + class FrgTensorD, + class TensorA, + class TensorB, + class FrgTensorC, + class KTileIterator, + class ResidueMNK + > + MCTLASS_DEVICE void + operator() ( + FrgTensorD &accum, + TensorA gA, // (BLK_M, BLK_K, K_TILES) + TensorB gB, // (BLK_N, BLK_K, K_TILES) + FrgTensorC const &src_accum, + KTileIterator k_tile_iter, int k_tile_count, + ResidueMNK residue_mnk, + int thread_idx, + char *smem_buf) + { + using namespace cute; + + static_assert(is_rmem::value, "D tensor must be rmem resident."); + static_assert(is_gmem::value, "A tensor must be gmem resident."); + static_assert(is_gmem::value, "B tensor must be gmem resident."); + static_assert(is_rmem::value, "C tensor must be rmem resident."); + static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3."); + static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3."); + + // Construct shared memory tiles + SharedStorage& storage = *reinterpret_cast(smem_buf); + Tensor sA = make_tensor(make_smem_ptr(storage.smem_a.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(storage.smem_b.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + CUTE_STATIC_ASSERT_V(size<0>(gA) == size<0>(sA)); // BLK_M + CUTE_STATIC_ASSERT_V(size<1>(gA) == size<1>(sA)); // BLK_K + CUTE_STATIC_ASSERT_V(size<0>(gB) == size<0>(sB)); // BLK_N + CUTE_STATIC_ASSERT_V(size<1>(gB) == size<1>(sB)); // BLK_K + CUTE_STATIC_ASSERT_V(size<1>(sA) == size<1>(sB)); // BLK_K + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sA)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sB)); // PIPE + + // Shift tensor so residue_k is at origin (Can't read any k_coord < residue_k) + // This aligns the tensor with BLK_K for all but the 0th k_tile + gA.data() = &gA(0, get<2>(residue_mnk), 0); + gB.data() = &gB(0, get<2>(residue_mnk), 0); + + // Partition the copying of A and B tiles across the threads + GmemTiledCopyA gmem_tiled_copy_A; + GmemTiledCopyB gmem_tiled_copy_B; + auto gmem_thr_copy_A = gmem_tiled_copy_A.get_slice(thread_idx); + auto gmem_thr_copy_B = gmem_tiled_copy_B.get_slice(thread_idx); + + Tensor tAgA = gmem_thr_copy_A.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k) + Tensor tAsA = gmem_thr_copy_A.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE) + Tensor tBgB = gmem_thr_copy_B.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k) + Tensor tBsB = gmem_thr_copy_B.partition_D(sB); // (BCPY,BCPY_N,BCPY_K,PIPE) + + // + // PREDICATES + // + + // Allocate predicate tensors for m and n + Tensor tApA = make_tensor(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{}); + Tensor tBpB = make_tensor(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{}); + + // Construct identity layout for sA and sB + Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + + // Repeat the partitioning with identity layouts + Tensor tAcA = gmem_thr_copy_A.partition_S(cA); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tBcB = gmem_thr_copy_B.partition_S(cB); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + + // Set predicates for m bounds + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < size<0>(tApA); ++m) { + tApA(m,0) = get<0>(tAcA(0,m,0)) < get<0>(residue_mnk); // blk_m coord < residue_m + } + // Set predicates for n bounds + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < size<0>(tBpB); ++n) { + tBpB(n,0) = get<0>(tBcB(0,n,0)) < get<1>(residue_mnk); // blk_n coord < residue_n + } + + // + // PREFETCH + // + + // Clear the smem tiles to account for predicated off loads + clear(tAsA); + clear(tBsB); + + // Start async loads for 0th k-tile, where we take care of the k residue + { + constexpr int k_pipe = 0; + + Tensor tAgAk = tAgA(_,_,_,*k_tile_iter); + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < size<2>(tAsA); ++k) { + if (get<1>(tAcA(0,0,k)) >= -get<2>(residue_mnk)) { // blk_k coord < residue_k (gA shifted) + copy_if(gmem_tiled_copy_A, tApA(_,k), tAgAk(_,_,k), tAsA(_,_,k,k_pipe)); + } + } + Tensor tBgBk = tBgB(_,_,_,*k_tile_iter); + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < size<2>(tBsB); ++k) { + if (get<1>(tBcB(0,0,k)) >= -get<2>(residue_mnk)) { // blk_k coord < residue_k (gB shifted) + copy_if(gmem_tiled_copy_B, tBpB(_,k), tBgBk(_,_,k), tBsB(_,_,k,k_pipe)); + } + } + cp_async_fence(); + ++k_tile_iter; + --k_tile_count; + } + + // Start async loads for 1st k-tile onwards, no k-residue handling needed + MCTLASS_PRAGMA_UNROLL + for (int k_pipe = 1; k_pipe < DispatchPolicy::Stages-1; ++k_pipe) { + if (k_tile_count <= 0) { + clear(tApA); + clear(tBpB); + } + copy_if(gmem_tiled_copy_A, tApA, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,k_pipe)); // CpAsync + copy_if(gmem_tiled_copy_B, tBpB, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,k_pipe)); // CpAsync + cp_async_fence(); + ++k_tile_iter; + --k_tile_count; + } + + // + // MMA Atom partitioning + // + + // Tile MMA compute thread partitions and allocate accumulators + TiledMma tiled_mma; + auto thr_mma = tiled_mma.get_thread_slice(thread_idx); + Tensor tCrA = thr_mma.partition_fragment_A(sA(_,_,0)); // (MMA,MMA_M,MMA_K) + Tensor tCrB = thr_mma.partition_fragment_B(sB(_,_,0)); // (MMA,MMA_N,MMA_K) + + CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(accum)); // MMA_M + CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(src_accum)); // MMA_M + CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(accum)); // MMA_N + CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(src_accum)); // MMA_N + CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K + + // + // Copy Atom retiling + // + + auto smem_tiled_copy_A = make_tiled_copy_A(SmemCopyAtomA{}, tiled_mma); + auto smem_thr_copy_A = smem_tiled_copy_A.get_thread_slice(thread_idx); + Tensor tCsA = smem_thr_copy_A.partition_S(sA); // (CPY,CPY_M,CPY_K,PIPE) + Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA); // (CPY,CPY_M,CPY_K) + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // CPY_M + CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrA_copy_view)); // CPY_K + + auto smem_tiled_copy_B = make_tiled_copy_B(SmemCopyAtomB{}, tiled_mma); + auto smem_thr_copy_B = smem_tiled_copy_B.get_thread_slice(thread_idx); + Tensor tCsB = smem_thr_copy_B.partition_S(sB); // (CPY,CPY_N,CPY_K,PIPE) + Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB); // (CPY,CPY_N,CPY_K) + CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // CPY_N + CUTE_STATIC_ASSERT_V(size<2>(tCsB) == size<2>(tCrB_copy_view)); // CPY_K + + // + // PIPELINED MAIN LOOP + // + + // Current pipe index in smem to read from + int smem_pipe_read = 0; + // Current pipe index in smem to write to + int smem_pipe_write = DispatchPolicy::Stages-1; + + Tensor tCsA_p = tCsA(_,_,_,smem_pipe_read); + Tensor tCsB_p = tCsB(_,_,_,smem_pipe_read); + + // Size of the register pipeline + auto K_BLOCK_MAX = size<2>(tCrA); + + // PREFETCH register pipeline + if (K_BLOCK_MAX > 1) { + // Wait until our first prefetched tile is loaded in + cp_async_wait(); + __syncthreads(); + + // Prefetch the first rmem from the first k-tile + copy(smem_tiled_copy_A, tCsA_p(_,_,Int<0>{}), tCrA_copy_view(_,_,Int<0>{})); + copy(smem_tiled_copy_B, tCsB_p(_,_,Int<0>{}), tCrB_copy_view(_,_,Int<0>{})); + } + + + MCTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > -(DispatchPolicy::Stages-1); --k_tile_count) + { + // Pipeline the outer products with a static for loop. + // + // Note, the for_each() function is required here to ensure `k_block` is of type Int. + for_each(make_int_sequence{}, [&] (auto k_block) + { + if (k_block == K_BLOCK_MAX - 1) + { + // Slice the smem_pipe_read smem + tCsA_p = tCsA(_,_,_,smem_pipe_read); + tCsB_p = tCsB(_,_,_,smem_pipe_read); + + // Commit the smem for smem_pipe_read + cp_async_wait(); + __syncthreads(); + } + + // Load A, B shmem->regs for k_block+1 + auto k_block_next = (k_block + Int<1>{}) % K_BLOCK_MAX; // static + copy(smem_tiled_copy_A, tCsA_p(_,_,k_block_next), tCrA_copy_view(_,_,k_block_next)); + copy(smem_tiled_copy_B, tCsB_p(_,_,k_block_next), tCrB_copy_view(_,_,k_block_next)); + // Copy gmem to smem before computing gemm on each k-pipe + if (k_block == 0) + { + // Set all predicates to false if we are going to overshoot bounds + if (k_tile_count <= 0) { + clear(tApA); + clear(tBpB); + } + copy_if(gmem_tiled_copy_A, tApA, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,smem_pipe_write)); + copy_if(gmem_tiled_copy_B, tBpB, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,smem_pipe_write)); + cp_async_fence(); + ++k_tile_iter; + + // Advance the pipe -- Doing it here accounts for K_BLOCK_MAX = 1 (no rmem pipe) + smem_pipe_write = smem_pipe_read; + ++smem_pipe_read; + smem_pipe_read = (smem_pipe_read == DispatchPolicy::Stages) ? 0 : smem_pipe_read; + } + + // Transform before compute + cute::transform(tCrA(_,_,k_block), TransformA{}); + cute::transform(tCrB(_,_,k_block), TransformB{}); + // Thread-level register gemm for k_block + cute::gemm(tiled_mma, accum, tCrA(_,_,k_block), tCrB(_,_,k_block), src_accum); + }); + + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_multistage_gmma_ss.hpp b/csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_multistage_gmma_ss.hpp new file mode 100644 index 0000000..fa39f64 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_multistage_gmma_ss.hpp @@ -0,0 +1,600 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/dispatch_policy.hpp" +#include "mctlass/pipeline/pipeline.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "mctlass/arch/reg_reconfig.h" + +#include "cute/arch/copy_sm90.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/algorithm/gemm.hpp" + +#include + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::collective { +using namespace cute; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int Stages, + class ClusterShape, + class TileShape_, + class ElementA_, + class StrideA_, + class ElementB_, + class StrideB_, + class TiledMma_, + class GmemTiledCopyA_, + class SmemLayoutAtomA_, + class SmemCopyAtomA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomB_, + class SmemCopyAtomB_, + class TransformB_> +struct CollectiveMma< + MainloopSm90CpAsyncGmmaUnpredicated, + TileShape_, + ElementA_, + StrideA_, + ElementB_, + StrideB_, + TiledMma_, + GmemTiledCopyA_, + SmemLayoutAtomA_, + SmemCopyAtomA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomB_, + SmemCopyAtomB_, + TransformB_> +{ + // + // Type Aliases + // + using DispatchPolicy = MainloopSm90CpAsyncGmmaUnpredicated; + using TileShape = TileShape_; + using ElementA = ElementA_; + using StrideA = StrideA_; + using ElementB = ElementB_; + using StrideB = StrideB_; + using TiledMma = TiledMma_; + using ElementAccumulator = typename TiledMma::ValTypeC; + using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomA = SmemLayoutAtomA_; + using SmemLayoutAtomB = SmemLayoutAtomB_; + using SmemCopyAtomA = SmemCopyAtomA_; + using SmemCopyAtomB = SmemCopyAtomB_; + using TransformA = TransformA_; + using TransformB = TransformB_; + using ArchTag = typename DispatchPolicy::ArchTag; + + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + using SmemLayoutA = decltype(tile_to_shape( + SmemLayoutAtomA{}, + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int{}))); + using SmemLayoutB = decltype(tile_to_shape( + SmemLayoutAtomB{}, + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}))); + + static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more."); + static_assert(cute::is_base_of::value && + cute::is_base_of::value, + "MMA atom must source both A and B operand from smem_desc for this mainloop."); + + struct SharedStorage + { + cute::array_aligned> smem_a; + cute::array_aligned> smem_b; + }; + + struct Arguments { + ElementA const* ptr_A; + StrideA dA; + ElementB const* ptr_B; + StrideB dB; + }; + + using Params = Arguments; + + // + // Methods + // + + CollectiveMma() = default; + + template + static constexpr Params + to_underlying_arguments(ProblemShape const& _, Arguments const& args, void* workspace) { + (void) workspace; + return args; + } + + /// Perform a collective-scoped matrix multiply-accumulate + template < + class TensorA, + class TensorB, + class FrgTensorC, + class KTileIterator, + class ResidueMNK + > + MCTLASS_DEVICE void + operator() ( + TensorA gA, + TensorB gB, + FrgTensorC& accum, + KTileIterator k_tile_iter, int k_tile_count, + ResidueMNK residue_mnk, + int thread_idx, + char *smem_buf, + Params const& mainloop_params) + { + using namespace cute; + + (void) residue_mnk; + + static_assert(is_gmem::value, "A tensor must be gmem resident."); + static_assert(is_gmem::value, "B tensor must be gmem resident."); + static_assert(is_rmem::value, "C tensor must be rmem resident."); + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2."); + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2."); + static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3."); + static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3."); + static_assert(cute::is_same::value, + "SM90 warpgroup MMA must specify transforms through MMA_Atom."); + static_assert(cute::is_same::value, + "SM90 warpgroup MMA must specify transforms through MMA_Atom."); + static_assert(cute::is_same::value, + "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions."); + static_assert(cute::is_same::value, + "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions."); + + SharedStorage& storage = *reinterpret_cast(smem_buf); + Tensor sA = make_tensor(make_smem_ptr(storage.smem_a.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(storage.smem_b.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + // Partition the copying of A and B tiles across the threads + GmemTiledCopyA gmem_tiled_copy_a; + GmemTiledCopyB gmem_tiled_copy_b; + auto gmem_thr_copy_a = gmem_tiled_copy_a.get_slice(thread_idx); + auto gmem_thr_copy_b = gmem_tiled_copy_b.get_slice(thread_idx); + + Tensor tAgA = gmem_thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k) + Tensor tAsA = gmem_thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE) + Tensor tBgB = gmem_thr_copy_b.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k) + Tensor tBsB = gmem_thr_copy_b.partition_D(sB); // (BCPY,BCPY_N,BCPY_K,PIPE) + + // Tile MMA atom and compute thread partitions across A, B and C + TiledMma tiled_mma; + auto thr_mma = tiled_mma.get_thread_slice(thread_idx); + + // Allocate registers for pipelining + Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE) + + Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_N,PIPE) + + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum)); // M + CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N + CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K + CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE + CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tAsA)); // PIPE + CUTE_STATIC_ASSERT_V(size<3>(tCsB) == size<3>(tBsB)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sA)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sB)); // PIPE + + // + // Prologue + // + + MCTLASS_PRAGMA_UNROLL + for (int k_pipe = 0; k_pipe < DispatchPolicy::Stages-1; ++k_pipe) { + copy(gmem_tiled_copy_a, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,k_pipe)); + copy(gmem_tiled_copy_b, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,k_pipe)); + cp_async_fence(); + ++k_tile_iter; + --k_tile_count; + } + + // Current pipe index in smem to read from + int smem_pipe_read = 0; + // Current pipe index in smem to write to + int smem_pipe_write = DispatchPolicy::Stages-1; + + // + // Pipelined Main Loop + // + MCTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > -(DispatchPolicy::Stages-1); --k_tile_count) + { + // Copy gmem to smem before computing gemm on each k-pipe + // pipe index in smem where the next gmem tile will be read into + copy(gmem_tiled_copy_a, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,smem_pipe_write)); + copy(gmem_tiled_copy_b, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,smem_pipe_write)); + cp_async_fence(); + if (k_tile_count > 0) { ++k_tile_iter; } + + // + // Compute on k_tile + // + warpgroup_fence_operand(accum); + warpgroup_arrive(); + + cp_async_wait(); + cute::gemm(tiled_mma, tCrA(_,_,_,smem_pipe_read), tCrB(_,_,_,smem_pipe_read), accum); + warpgroup_commit_batch(); + + // + // Advance the pipe + // + ++smem_pipe_read; + smem_pipe_read = (smem_pipe_read == DispatchPolicy::Stages) ? smem_pipe_read = 0 : smem_pipe_read; + + ++smem_pipe_write; + smem_pipe_write = (smem_pipe_write == DispatchPolicy::Stages) ? smem_pipe_write = 0 : smem_pipe_write; + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(accum); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int Stages, + class ClusterShape, + class TileShape_, + class ElementA_, + class StrideA_, + class ElementB_, + class StrideB_, + class TiledMma_, + class GmemTiledCopyA_, + class SmemLayoutAtomA_, + class SmemCopyAtomA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomB_, + class SmemCopyAtomB_, + class TransformB_> +struct CollectiveMma< + MainloopSm90CpAsyncGmma, + TileShape_, + ElementA_, + StrideA_, + ElementB_, + StrideB_, + TiledMma_, + GmemTiledCopyA_, + SmemLayoutAtomA_, + SmemCopyAtomA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomB_, + SmemCopyAtomB_, + TransformB_> +{ + // + // Type Aliases + // + using DispatchPolicy = MainloopSm90CpAsyncGmma; + using TileShape = TileShape_; + using ElementA = ElementA_; + using StrideA = StrideA_; + using ElementB = ElementB_; + using StrideB = StrideB_; + using TiledMma = TiledMma_; + using ElementAccumulator = typename TiledMma::ValTypeC; using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomA = SmemLayoutAtomA_; + using SmemLayoutAtomB = SmemLayoutAtomB_; + using SmemCopyAtomA = SmemCopyAtomA_; + using SmemCopyAtomB = SmemCopyAtomB_; + using TransformA = TransformA_; + using TransformB = TransformB_; + using ArchTag = typename DispatchPolicy::ArchTag; + + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + using SmemLayoutA = decltype(tile_to_shape( + SmemLayoutAtomA{}, + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int{}))); + using SmemLayoutB = decltype(tile_to_shape( + SmemLayoutAtomB{}, + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}))); + + static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more."); + static_assert(cute::is_base_of::value && + cute::is_base_of::value, + "MMA atom must source both A and B operand from smem_desc for this mainloop."); + + struct SharedStorage + { + cute::array_aligned> smem_a; + cute::array_aligned> smem_b; + }; + + struct Arguments { + ElementA const* ptr_A; + StrideA dA; + ElementB const* ptr_B; + StrideB dB; + }; + + using Params = Arguments; + + // + // Methods + // + + CollectiveMma() = default; + + template + static constexpr Params + to_underlying_arguments(ProblemShape const& _, Arguments const& args, void* workspace) { + (void) workspace; + return args; + } + + /// Perform a collective-scoped matrix multiply-accumulate + template < + class FrgTensorD, + class TensorA, + class TensorB, + class FrgTensorC, + class KTileIterator, + class ResidueMNK + > + MCTLASS_DEVICE void + operator() ( + FrgTensorD &accum, + TensorA gA, + TensorB gB, + FrgTensorC const &src_accum, + KTileIterator k_tile_iter, int k_tile_count, + ResidueMNK residue_mnk, + int thread_idx, + char *smem_buf) + { + using namespace cute; + + static_assert(is_rmem::value, "D tensor must be rmem resident."); + static_assert(is_gmem::value, "A tensor must be gmem resident."); + static_assert(is_gmem::value, "B tensor must be gmem resident."); + static_assert(is_rmem::value, "C tensor must be rmem resident."); + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2."); + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2."); + static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3."); + static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3."); + static_assert(cute::is_same::value, + "SM90 warpgroup MMA must specify transforms through MMA_Atom."); + static_assert(cute::is_same::value, + "SM90 warpgroup MMA must specify transforms through MMA_Atom."); + static_assert(cute::is_same::value, + "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions."); + static_assert(cute::is_same::value, + "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions."); + + SharedStorage& storage = *reinterpret_cast(smem_buf); + Tensor sA = make_tensor(make_smem_ptr(storage.smem_a.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(storage.smem_b.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + // Shift tensor so residue_k is at origin (Can't read any k_coord < residue_k) + // This aligns the tensor with BLK_K for all but the 0th k_tile + gA.data() = &gA(0, get<2>(residue_mnk), 0); + gB.data() = &gB(0, get<2>(residue_mnk), 0); + + // Partition the copying of A and B tiles across the threads + GmemTiledCopyA gmem_tiled_copy_a; + GmemTiledCopyB gmem_tiled_copy_b; + auto gmem_thr_copy_a = gmem_tiled_copy_a.get_slice(thread_idx); + auto gmem_thr_copy_b = gmem_tiled_copy_b.get_slice(thread_idx); + + Tensor tAgA = gmem_thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k) + Tensor tAsA = gmem_thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE) + Tensor tBgB = gmem_thr_copy_b.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k) + Tensor tBsB = gmem_thr_copy_b.partition_D(sB); // (BCPY,BCPY_N,BCPY_K,PIPE) + + // + // PREDICATES + // + + // Allocate predicate tensors for m and n + Tensor tApA = make_tensor(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{}); + Tensor tBpB = make_tensor(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{}); + + // Construct identity layout for sA and sB + Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + + // Repeat the partitioning with identity layouts + Tensor tAcA = gmem_thr_copy_a.partition_S(cA); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tBcB = gmem_thr_copy_b.partition_S(cB); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + + // Set predicates for m bounds + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < size<0>(tApA); ++m) { + tApA(m,0) = get<0>(tAcA(0,m,0)) < get<0>(residue_mnk); // blk_m coord < residue_m + } + // Set predicates for n bounds + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < size<0>(tBpB); ++n) { + tBpB(n,0) = get<0>(tBcB(0,n,0)) < get<1>(residue_mnk); // blk_n coord < residue_n + } + + // + // Prologue/PREFETCH + // + + // Clear the smem tiles to account for predicated off loads + clear(tAsA); + clear(tBsB); + + // Start async loads for 0th k-tile, where we take care of the k residue + { + constexpr int k_pipe = 0; + + Tensor tAgAk = tAgA(_,_,_,*k_tile_iter); + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < size<2>(tAsA); ++k) { + if (get<1>(tAcA(0,0,k)) >= -get<2>(residue_mnk)) { // blk_k coord < residue_k (gA shifted) + copy_if(gmem_tiled_copy_a, tApA(_,k), tAgAk(_,_,k), tAsA(_,_,k,k_pipe)); + } + } + Tensor tBgBk = tBgB(_,_,_,*k_tile_iter); + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < size<2>(tBsB); ++k) { + if (get<1>(tBcB(0,0,k)) >= -get<2>(residue_mnk)) { // blk_k coord < residue_k (gB shifted) + copy_if(gmem_tiled_copy_b, tBpB(_,k), tBgBk(_,_,k), tBsB(_,_,k,k_pipe)); + } + } + cp_async_fence(); + ++k_tile_iter; + --k_tile_count; + } + + // Start async loads for 1st k-tile onwards, no k-residue handling needed + MCTLASS_PRAGMA_UNROLL + for (int k_pipe = 1; k_pipe < DispatchPolicy::Stages-1; ++k_pipe) { + if (k_tile_count <= 0) { + clear(tApA); + clear(tBpB); + } + copy_if(gmem_tiled_copy_a, tApA, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,k_pipe)); // CpAsync + copy_if(gmem_tiled_copy_b, tBpB, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,k_pipe)); // CpAsync + cp_async_fence(); + ++k_tile_iter; + --k_tile_count; + } + + // + // MMA Atom partitioning + // + + // Tile MMA atom and compute thread partitions across A, B and C + TiledMma tiled_mma; + auto thr_mma = tiled_mma.get_thread_slice(thread_idx); + + // Allocate registers for pipelining + Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE) + + Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_N,PIPE) + + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum)); // M + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(src_accum)); // M + CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N + CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(src_accum)); // N + CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K + CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE + CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tAsA)); // PIPE + CUTE_STATIC_ASSERT_V(size<3>(tCsB) == size<3>(tBsB)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sA)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sB)); // PIPE + + // Current pipe index in smem to read from + int smem_pipe_read = 0; + // Current pipe index in smem to write to + int smem_pipe_write = DispatchPolicy::Stages-1; + + // + // Pipelined Main Loop + // + MCTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > -(DispatchPolicy::Stages-1); --k_tile_count) + { + // + // Copy gmem to smem for *k_tile_iter + // + if (k_tile_count <= 0) { + clear(tApA); + clear(tBpB); + } + copy_if(gmem_tiled_copy_a, tApA, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,smem_pipe_write)); // CpAsync + copy_if(gmem_tiled_copy_b, tBpB, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,smem_pipe_write)); // CpAsync + cp_async_fence(); + ++k_tile_iter; + + // + // Compute on k_tile + // + warpgroup_fence_operand(accum); + warpgroup_arrive(); + + cp_async_wait(); + cute::gemm(tiled_mma, accum, tCrA(_,_,_,smem_pipe_read), tCrB(_,_,_,smem_pipe_read), src_accum); + warpgroup_commit_batch(); + + // + // Advance the pipe + // + ++smem_pipe_read; + smem_pipe_read = (smem_pipe_read == DispatchPolicy::Stages) ? smem_pipe_read = 0 : smem_pipe_read; + + ++smem_pipe_write; + smem_pipe_write = (smem_pipe_write == DispatchPolicy::Stages) ? smem_pipe_write = 0 : smem_pipe_write; + + // Wait for the pipeline MMAs to drain + warpgroup_wait<0>(); + warpgroup_fence_operand(accum); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp b/csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp new file mode 100644 index 0000000..f457fff --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp @@ -0,0 +1,594 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/mctlass.h" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/arch/copy_sm90.hpp" +#include "mctlass/gemm/gemm.h" +#include "mctlass/detail/dependent_false.hpp" +#include "mctlass/gemm/dispatch_policy.hpp" + +#include "cute/algorithm/functional.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/algorithm/gemm.hpp" +#include "cute/tensor_predicate.hpp" +#include "cute/numeric/arithmetic_tuple.hpp" +#include "mctlass/pipeline/pipeline.hpp" +#include "mctlass/transform/collective/sm90_wgmma_transpose.hpp" +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::collective { +using namespace cute; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// WarpSpecialized Mainloop that source A operand from registers +template < + int Stages, + class ClusterShape, + class KernelSchedule, + int PipelineAsyncMmaStages, + class TileShape_, + class ElementA_, + class StrideA_, + class ElementB_, + class StrideB_, + class TiledMma_, + class GmemTiledCopyA_, + class SmemLayoutAtomA_, + class SmemCopyAtomA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomB_, + class SmemCopyAtomB_, + class TransformB_> +struct CollectiveMma< + MainloopSm90TmaGmmaRmemAWarpSpecialized, + TileShape_, + ElementA_, + StrideA_, + ElementB_, + StrideB_, + TiledMma_, + GmemTiledCopyA_, + SmemLayoutAtomA_, + SmemCopyAtomA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomB_, + SmemCopyAtomB_, + TransformB_> +{ + // + // Type Aliases + // + using DispatchPolicy = MainloopSm90TmaGmmaRmemAWarpSpecialized; + using TileShape = TileShape_; + using ElementA = ElementA_; + using StrideA = StrideA_; + using ElementB = ElementB_; + using StrideB = StrideB_; + using TiledMma = TiledMma_; + using ElementAccumulator = typename TiledMma::ValTypeC; + using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomA = SmemLayoutAtomA_; + using SmemLayoutAtomB = SmemLayoutAtomB_; + using SmemCopyAtomA = SmemCopyAtomA_; + using SmemCopyAtomB = SmemCopyAtomB_; + + // Swap and transpose A/B for A k-major layout and B mn-major layout since WGMMA is k-major only (e.g. tf32, Fp32, Int8 WGMMA) + static constexpr bool IsLayoutAkBmn = + cute::is_same_v, layout::RowMajor> && + cute::is_same_v, layout::RowMajor>; + + static constexpr bool IsInputSizeTwoBytes = sizeof(ElementA) == 2 && sizeof(ElementB) == 2; + static constexpr bool SwapAB = !IsInputSizeTwoBytes && IsLayoutAkBmn; + using InternalSmemLayoutAtomA = cute::conditional_t; + using InternalSmemLayoutAtomB = cute::conditional_t; + using InternalSmemCopyAtomA = cute::conditional_t; + using InternalSmemCopyAtomB = cute::conditional_t; + + using TransformA = TransformA_; + using TransformB = TransformB_; + using ArchTag = typename DispatchPolicy::ArchTag; + + using MainloopPipeline = mctlass::PipelineTmaAsync< + DispatchPolicy::Stages, + typename DispatchPolicy::ClusterShape>; + using PipelineState = mctlass::PipelineState; + + using PipelineParams = typename MainloopPipeline::Params; + + static_assert(rank(InternalSmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<0>(TileShape{}) % size<0>(InternalSmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(InternalSmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + static_assert(rank(InternalSmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<1>(TileShape{}) % size<0>(InternalSmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(InternalSmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + // Tile along K mode first before tiling over MN. PIPE mode last as usual. + // This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs. + using SmemLayoutA = decltype(tile_to_shape( + InternalSmemLayoutAtomA{}, + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int{}), + Step<_2,_1,_3>{})); + using SmemLayoutB = decltype(tile_to_shape( + InternalSmemLayoutAtomB{}, + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}), + Step<_2,_1,_3>{})); + + // If A mn-layout and B mn-layout, transposing B matrix since WGMMA is k-major only (e.g. tf32, fp32, fp8, int8). + static constexpr bool IsLayoutAmnBmn = + cute::is_same_v, layout::ColumnMajor> && + cute::is_same_v, layout::RowMajor>; + static constexpr bool TransposeB = !IsInputSizeTwoBytes && IsLayoutAmnBmn; + + static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more."); + static_assert(not cute::is_base_of::value && + cute::is_base_of::value, + "MMA atom must source A from rmem and B operand from smem_desc for this mainloop."); + static_assert(cute::is_same_v || cute::is_same_v, + "GmemTiledCopy - invalid SM90 TMA copy atom specified."); + static_assert(cute::is_same_v || cute::is_same_v, + "GmemTiledCopy - invalid SM90 TMA copy atom specified."); + + // TMA converts f32 input to tf32 when copying from GMEM to SMEM + // For all other types, cast to size equivalent uint type to avoid any rounding by TMA. + static constexpr bool ConvertF32toTF32A = cute::is_same_v; + static constexpr bool ConvertF32toTF32B = cute::is_same_v; + using ConvertedElementA = cute::conditional_t>>; + using ConvertedElementB = cute::conditional_t>>; + using InternalElementA = cute::conditional_t; + using InternalElementB = cute::conditional_t; + using InternalStrideA = cute::conditional_t; + using InternalStrideB = cute::conditional_t; + + using GmmaSmemLayoutAtomB = decltype(transform::collective::detail::gmma_smem_transpose_or_passthrough< + TransposeB, InternalSmemLayoutAtomB, InternalElementB>()); + + // SmemLayoutB for GMMA is different from SmemLayoutB for TMA if TransposeB + using GmmaSmemLayoutB = decltype(tile_to_shape( + GmmaSmemLayoutAtomB{}, + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}), + Step<_2,_1,_3>{})); + + static_assert(!SwapAB || !TransposeB, "Cannot SwapAB and TransposeB at the same time."); + static_assert(TransposeB || (cute::is_same_v), + "Should be same layout if not TransposeB."); + static_assert(!TransposeB || size<1>(SmemLayoutB{}) * sizeof(InternalElementB) == 128, + "SmemLayoutB K must be 128bytes to be transposed."); + static_assert(!transform::collective::detail::use_universal_transposition(), + "Warp specialized ARF kernels have not supported universal B transposition yet."); + static_assert(!TransposeB || !cute::is_same_v, + "Transpose RS kernel requires kernel schedule schmem is not KernelTmaWarpSpecializedCooperative."); + + struct SharedStorage + { + struct TensorStorage : cute::aligned_struct<128> { + cute::array_aligned> smem_A; + cute::array_aligned> smem_B; + } tensors; + + using PipelineStorage = typename MainloopPipeline::SharedStorage; + PipelineStorage pipeline; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + using PipelineStorage = typename SharedStorage::PipelineStorage; + + // Host side kernel arguments + struct Arguments { + ElementA const* ptr_A; + StrideA dA; + ElementB const* ptr_B; + StrideB dB; + }; + + // Device side kernel params + struct Params { + // Assumption: StrideA is congruent with Problem_MK + using TMA_A = decltype(make_tma_copy( + GmemTiledCopyA{}, + make_tensor(static_cast(nullptr), repeat_like(InternalStrideA{}, int32_t(0)), InternalStrideA{}), + SmemLayoutA{}(_,_,0), + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})), + size<1>(ClusterShape{}))); // mcast along N mode for this M load, if any + // Assumption: StrideB is congruent with Problem_NK + using TMA_B = decltype(make_tma_copy( + GmemTiledCopyB{}, + make_tensor(static_cast(nullptr), repeat_like(InternalStrideB{}, int32_t(0)), InternalStrideB{}), + SmemLayoutB{}(_,_,0), + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), + size<0>(ClusterShape{}))); // mcast along M mode for this N load, if any + TMA_A tma_load_a; + TMA_B tma_load_b; + }; + + // + // Methods + // + + template + static constexpr Params + to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) { + (void) workspace; + + // Optionally append _1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK) + auto problem_shape_MNKL = append<4>(problem_shape, Int<1>{}); + auto M = get<0>(problem_shape_MNKL); + auto N = get<1>(problem_shape_MNKL); + auto K = get<2>(problem_shape_MNKL); + auto L = get<3>(problem_shape_MNKL); + + if constexpr (SwapAB) { + M = get<1>(problem_shape_MNKL); + N = get<0>(problem_shape_MNKL); + } + + InternalElementA const* ptr_A; + InternalStrideA dA; + InternalElementB const* ptr_B; + InternalStrideB dB; + + if constexpr (not SwapAB) { + ptr_A = reinterpret_cast(args.ptr_A); + ptr_B = reinterpret_cast(args.ptr_B); + dA = args.dA; + dB = args.dB; + } + else { + ptr_A = reinterpret_cast(args.ptr_B); + ptr_B = reinterpret_cast(args.ptr_A); + dA = args.dB; + dB = args.dA; + } + + Tensor tensor_a = make_tensor(ptr_A, make_layout(make_shape(M,K,L), dA)); + Tensor tensor_b = make_tensor(ptr_B, make_layout(make_shape(N,K,L), dB)); + typename Params::TMA_A tma_load_a = make_tma_copy( + GmemTiledCopyA{}, + tensor_a, + SmemLayoutA{}(_,_,cute::Int<0>{}), + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})), + size<1>(ClusterShape{})); // mcast along N mode for this M load, if any + typename Params::TMA_B tma_load_b = make_tma_copy( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,cute::Int<0>{}), + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), + size<0>(ClusterShape{})); // mcast along M mode for this N load, if any + return { + tma_load_a, + tma_load_b + }; + } + + static constexpr int K_PIPE_MAX = DispatchPolicy::Stages; + static constexpr int K_PIPE_MMAS = DispatchPolicy::PipelineAsyncMmaStages; + static_assert(K_PIPE_MMAS == 0, "no MMA stage should be asynchronous for this mainloop for now."); + static constexpr uint32_t TmaTransactionBytes = + (size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast(sizeof(InternalElementA)))+ + (size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(sizeof(InternalElementB))); + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + MCTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& mainloop_params) + { + cute::prefetch_tma_descriptor(mainloop_params.tma_load_a.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_b.get_tma_descriptor()); + } + + /// Perform a collective-scoped matrix multiply-accumulate + /// Producer Perspective + template < + class TensorA, class TMA_LOAD_A, + class TensorB, class TMA_LOAD_B, + class KTileIterator + > + MCTLASS_DEVICE void + load( + MainloopPipeline pipeline, + PipelineState smem_pipe_write, + TensorA const& gA, TMA_LOAD_A& tma_load_a, + TensorB const& gB, TMA_LOAD_B& tma_load_b, + KTileIterator k_tile_iter, int k_tile_count, + int thread_idx, + TensorStorage& shared_tensors) + { + + using namespace cute; + int warp_idx = canonical_warp_idx(); + int warp_idx_in_warp_group = warp_idx % 4; + int lane_predicate = cute::elect_one_sync(); + + if (warp_idx_in_warp_group == 0 and lane_predicate) { + Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + // + // Prepare the TMA loads for A and B + // + + dim3 cluster_local_block_id = cute::block_id_in_cluster(); + auto block_tma_a = tma_load_a.get_slice(cluster_local_block_id.y); + auto block_tma_b = tma_load_b.get_slice(cluster_local_block_id.x); + + // Applies the mapping from block_tma_a + Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k) + Tensor tAsA = block_tma_a.partition_D(sA); // (TMA,TMA_M,TMA_K,PIPE) + + Tensor tBgB = block_tma_b.partition_S(gB); // (TMA,TMA_N,TMA_K,k) + Tensor tBsB = block_tma_b.partition_D(sB); // (TMA,TMA_N,TMA_K,PIPE) + + uint16_t mcast_mask_a = 0; + uint16_t mcast_mask_b = 0; + + // Issue TmaLoads + // Maps the tile -> block, value + if constexpr (cute::is_same_v) { + auto block_layout = Layout{}; // (m,n) -> block_id + for (int n = 0; n < size<1>(block_layout); ++n) { + mcast_mask_a |= (uint16_t(1) << block_layout(cluster_local_block_id.x,n,Int<0>{})); + } + } + + if constexpr (cute::is_same_v) { + auto block_layout = Layout{}; // (m,n) -> block_id + for (int m = 0; m < size<0>(block_layout); ++m) { + mcast_mask_b |= (uint16_t(1) << block_layout(m,cluster_local_block_id.y,Int<0>{})); + } + } + + // Mainloop + MCTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + // LOCK smem_pipe_write for _writing_ + pipeline.producer_acquire(smem_pipe_write); + + // + // Copy gmem to smem for *k_tile_iter + // + + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write); + + int write_stage = smem_pipe_write.index(); + copy(tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage)); + copy(tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage)); + ++k_tile_iter; + + // Advance smem_pipe_write + ++smem_pipe_write; + } + } + } + + /// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster + MCTLASS_DEVICE void + load_tail(MainloopPipeline pipeline, PipelineState smem_pipe_write) + { + int warp_idx = canonical_warp_idx(); + int warp_idx_in_warp_group = warp_idx % 4; + int lane_predicate = cute::elect_one_sync(); + + // Issue the epilogue waits + if (warp_idx_in_warp_group == 0 and lane_predicate) { + /* This helps avoid early exit of blocks in Cluster + * Waits for all stages to either be released (all + * Consumer UNLOCKs), or if the stage was never used + * then would just be acquired since the phase was + * still inverted from make_producer_start_state + */ + pipeline.producer_tail(smem_pipe_write); + } + } + + /// Perform a collective-scoped matrix multiply-accumulate + /// Consumer Perspective + template < + class FrgTensorC + > + MCTLASS_DEVICE void + mma(MainloopPipeline pipeline, + PipelineState smem_pipe_read, + FrgTensorC& accum, + int k_tile_count, + int thread_idx, + TensorStorage& shared_tensors, + Params const& mainloop_params) + { + using namespace cute; + static_assert(is_rmem::value, "C tensor must be rmem resident."); + static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3."); + static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3."); + static_assert(rank(InternalSmemLayoutAtomA{}) == 2, "InternalSmemLayoutAtomA must be rank 2."); + static_assert(rank(InternalSmemLayoutAtomB{}) == 2, "InternalSmemLayoutAtomB must be rank 2."); + static_assert(!cute::is_void_v, + "SM90 GMMA mainloops must specify a non-void copy atom for smem sourced instructions."); + static_assert(cute::is_void_v, + "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions."); + + // Obtain warp index + int warp_idx = canonical_warp_idx(); + int warp_idx_in_warp_group = warp_idx % 4; + int warp_group_thread_idx = thread_idx % 128; + + Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + // If TransposeB, GMMA will read from transposed B layout SMEM + Tensor gmma_sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), GmmaSmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + // + // Define C accumulators and A/B partitioning + // + + TiledMma tiled_mma; + auto thread_mma = tiled_mma.get_thread_slice(thread_idx); + + // Allocate fragments and descriptors + Tensor tCsA = thread_mma.partition_A(sA); + Tensor tCrA = thread_mma.partition_fragment_A(sA(_,_,Int<0>{})); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCsB = thread_mma.partition_B(gmma_sB); // (MMA,MMA_N,MMA_K,PIPE) + Tensor tCrB = thread_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE) + + // + // Copy Atom A retiling + // + + auto smem_tiled_copy_A = make_tiled_copy_A(InternalSmemCopyAtomA{}, tiled_mma); + auto smem_thr_copy_A = smem_tiled_copy_A.get_thread_slice(thread_idx); + Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA); // (CPY,CPY_M,CPY_K) + + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // CPY_M + CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrA_copy_view)); // CPY_K + CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(accum)); // MMA_M + CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N + CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K + CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sA)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sB)); // PIPE + + // + // PIPELINED MAIN LOOP + // + static_assert((0 <= K_PIPE_MMAS) && (K_PIPE_MMAS < K_PIPE_MAX), + "ERROR : Incorrect number of MMAs in flight"); + + // We release buffers to producer warps(dma load) with some mmas in flight + PipelineState smem_pipe_release = smem_pipe_read; + + // Prologue GMMAs + int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count); + + tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; + + warpgroup_fence_operand(accum); + MCTLASS_PRAGMA_UNROLL + for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue) { + // WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value) + pipeline.consumer_wait(smem_pipe_read); + + // copy smem->rmem for A operand + copy(smem_tiled_copy_A, tCsA(_,_,_,smem_pipe_read.index()), tCrA_copy_view); + // transpose B operand in SMEM + if constexpr (TransposeB) { + transform::collective::detail::transpose_b_operand( + sB, gmma_sB, smem_pipe_read, warp_idx_in_warp_group, warp_group_thread_idx, + tiled_mma, SmemLayoutB{}, InternalSmemLayoutAtomB{}, InternalElementB{}); + } // if TransposeB + + int read_stage = smem_pipe_read.index(); + warpgroup_arrive(); + // Unroll the K mode manually to set scale D to 1 + MCTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + // (V,M,K) x (V,N,K) => (V,M,N) + cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block,read_stage), accum); + tiled_mma.accumulate_ = GMMA::ScaleOut::One; + } + + warpgroup_commit_batch(); + + ++smem_pipe_read; + } + warpgroup_fence_operand(accum); + // Mainloop GMMAs + k_tile_count -= prologue_mma_count; + MCTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) { + // WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value) + pipeline.consumer_wait(smem_pipe_read); + + // + // Compute on k_tile + // + // copy smem->rmem for A operand + copy(smem_tiled_copy_A, tCsA(_,_,_,smem_pipe_read.index()), tCrA_copy_view); + // transpose B operand in SMEM + if constexpr (TransposeB) { + transform::collective::detail::transpose_b_operand( + sB, gmma_sB, smem_pipe_read, warp_idx_in_warp_group, warp_group_thread_idx, + tiled_mma, SmemLayoutB{}, InternalSmemLayoutAtomB{}, InternalElementB{}); + } // if TransposeB + + int read_stage = smem_pipe_read.index(); + warpgroup_fence_operand(accum); + warpgroup_arrive(); + // Unroll the K mode manually to set scale D to 1 + MCTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + // (V,M,K) x (V,N,K) => (V,M,N) + cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block,read_stage), accum); + tiled_mma.accumulate_ = GMMA::ScaleOut::One; + } + warpgroup_commit_batch(); + + /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed + warpgroup_wait(); + warpgroup_fence_operand(accum); + + pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it + + // Advance smem_pipe_read and smem_pipe_release + ++smem_pipe_read; + ++smem_pipe_release; + } + + warpgroup_fence_operand(accum); + } + + /// Perform a Consumer Epilogue to release all buffers + MCTLASS_DEVICE void + mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) { + // Prologue GMMAs + int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count); + k_tile_count -= prologue_mma_count; + + smem_pipe_release.advance(k_tile_count); + + // Wait on all GMMAs to complete + warpgroup_wait<0>(); + + for (int count = 0; count < prologue_mma_count; ++count) { + pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it + ++smem_pipe_release; + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp b/csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp new file mode 100644 index 0000000..d76fb31 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp @@ -0,0 +1,497 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/mctlass.h" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/arch/copy_sm90.hpp" +#include "mctlass/gemm/dispatch_policy.hpp" + +#include "cute/algorithm/functional.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/algorithm/gemm.hpp" +#include "cute/tensor_predicate.hpp" +#include "cute/numeric/arithmetic_tuple.hpp" +#include "mctlass/pipeline/pipeline.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::collective { +using namespace cute; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + int Stages, + class ClusterShape, + int PipelineAsyncMmaStages, + class TileShape_, + class ElementA_, + class StrideA_, + class ElementB_, + class StrideB_, + class TiledMma_, + class GmemTiledCopyA_, + class SmemLayoutAtomA_, + class SmemCopyAtomA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomB_, + class SmemCopyAtomB_, + class TransformB_> +struct CollectiveMma< + MainloopSm90TmaGmma, + TileShape_, + ElementA_, + StrideA_, + ElementB_, + StrideB_, + TiledMma_, + GmemTiledCopyA_, + SmemLayoutAtomA_, + SmemCopyAtomA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomB_, + SmemCopyAtomB_, + TransformB_> +{ + // + // Type Aliases + // + using DispatchPolicy = MainloopSm90TmaGmma; + using TileShape = TileShape_; + using ElementA = ElementA_; + using StrideA = StrideA_; + using ElementB = ElementB_; + using StrideB = StrideB_; + using TiledMma = TiledMma_; + using ElementAccumulator = typename TiledMma::ValTypeC; + using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomA = SmemLayoutAtomA_; + using SmemLayoutAtomB = SmemLayoutAtomB_; + using SmemCopyAtomA = SmemCopyAtomA_; + using SmemCopyAtomB = SmemCopyAtomB_; + using TransformA = TransformA_; + using TransformB = TransformB_; + using ArchTag = typename DispatchPolicy::ArchTag; + + using MainloopPipeline = mctlass::PipelineTmaAsync< + DispatchPolicy::Stages, + typename DispatchPolicy::ClusterShape>; + + using PipelineParams = typename MainloopPipeline::Params; + using PipelineState = typename mctlass::PipelineState; + + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + // Tile along K mode first before tiling over MN. PIPE mode last as usual. + // This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs. + using SmemLayoutA = decltype(tile_to_shape( + SmemLayoutAtomA{}, + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int{}), + Step<_2,_1,_3>{})); + using SmemLayoutB = decltype(tile_to_shape( + SmemLayoutAtomB{}, + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}), + Step<_2,_1,_3>{})); + + static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more."); + static_assert(cute::is_base_of::value && + cute::is_base_of::value, + "MMA atom must source both A and B operand from smem_desc for this mainloop."); + static_assert(cute::is_same_v || cute::is_same_v, + "GmemTiledCopy - invalid SM90 TMA copy atom specified."); + static_assert(cute::is_same_v || cute::is_same_v, + "GmemTiledCopy - invalid SM90 TMA copy atom specified."); + + // TMA converts f32 input to tf32 when copying from GMEM to SMEM + // For all other types, cast to size equivalent uint type to avoid any rounding by TMA. + static constexpr bool ConvertF32toTF32A = cute::is_same_v; + static constexpr bool ConvertF32toTF32B = cute::is_same_v; + using InternalElementA = cute::conditional_t>>; + using InternalElementB = cute::conditional_t>>; + + struct SharedStorage + { + cute::array_aligned> smem_A; + cute::array_aligned> smem_B; + + using PipelineStorage = typename MainloopPipeline::SharedStorage; + alignas(16) PipelineStorage pipeline_storage; + }; + + // Host side kernel arguments + struct Arguments { + ElementA const* ptr_A; + StrideA dA; + ElementB const* ptr_B; + StrideB dB; + }; + + // Device side kernel params + struct Params { + // Assumption: StrideA is congruent with Problem_MK + using TMA_A = decltype(make_tma_copy( + GmemTiledCopyA{}, + make_tensor(static_cast(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}), + SmemLayoutA{}(_,_,0), + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})), + size<1>(ClusterShape{}))); // mcast along N mode for this M load, if any + // Assumption: StrideB is congruent with Problem_NK + using TMA_B = decltype(make_tma_copy( + GmemTiledCopyB{}, + make_tensor(static_cast(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}), + SmemLayoutB{}(_,_,0), + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), + size<0>(ClusterShape{}))); // mcast along M mode for this N load, if any + TMA_A tma_load_a; + TMA_B tma_load_b; + }; + + // + // Methods + // + + template + static constexpr Params + to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) { + (void) workspace; + + // Optionally append _1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK) + auto problem_shape_MNKL = append<4>(problem_shape, Int<1>{}); + auto M = get<0>(problem_shape_MNKL); + auto N = get<1>(problem_shape_MNKL); + auto K = get<2>(problem_shape_MNKL); + auto L = get<3>(problem_shape_MNKL); + + auto ptr_A = reinterpret_cast(args.ptr_A); + auto ptr_B = reinterpret_cast(args.ptr_B); + + Tensor tensor_a = make_tensor(ptr_A, make_layout(make_shape(M,K,L), args.dA)); + Tensor tensor_b = make_tensor(ptr_B, make_layout(make_shape(N,K,L), args.dB)); + typename Params::TMA_A tma_load_a = make_tma_copy( + GmemTiledCopyA{}, + tensor_a, + SmemLayoutA{}(_,_,cute::Int<0>{}), + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})), + size<1>(ClusterShape{})); // mcast along N mode for this M load, if any + typename Params::TMA_B tma_load_b = make_tma_copy( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,cute::Int<0>{}), + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), + size<0>(ClusterShape{})); // mcast along M mode for this N load, if any + return { + tma_load_a, + tma_load_b + }; + } + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + MCTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& mainloop_params) + { + cute::prefetch_tma_descriptor(mainloop_params.tma_load_a.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_b.get_tma_descriptor()); + } + + /// Perform a collective-scoped matrix multiply-accumulate + template < + class TensorA, class TMA_LOAD_A, + class TensorB, class TMA_LOAD_B, + class FrgTensorC, + class KTileIterator + > + MCTLASS_DEVICE void + operator() ( + TensorA const& gA, TMA_LOAD_A& tma_load_a, + TensorB const& gB, TMA_LOAD_B& tma_load_b, + FrgTensorC& accum, + KTileIterator k_tile_iter, int k_tile_count, + int thread_idx, + char* shared_memory, + Params const& mainloop_params) + { + using namespace cute; + + static_assert(is_rmem::value, "C tensor must be rmem resident."); + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2."); + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2."); + static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3."); + static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3."); + static_assert(cute::is_void_v, + "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions."); + static_assert(cute::is_void_v, + "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions."); + + SharedStorage& storage = *reinterpret_cast(shared_memory); + Tensor sA = make_tensor(make_smem_ptr(storage.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(storage.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + // + // Prepare the TMA loads for A and B + // + dim3 cluster_local_block_id = cute::block_id_in_cluster(); + auto block_tma_a = tma_load_a.get_slice(cluster_local_block_id.y); + auto block_tma_b = tma_load_b.get_slice(cluster_local_block_id.x); + + // Applies the mapping from block_tma_a + Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k) + Tensor tAsA = block_tma_a.partition_D(sA); // (TMA,TMA_M,TMA_K,PIPE) + + Tensor tBgB = block_tma_b.partition_S(gB); // (TMA,TMA_N,TMA_K,k) + Tensor tBsB = block_tma_b.partition_D(sB); // (TMA,TMA_N,TMA_K,PIPE) + + // + // Prepare TMA membars and PREFETCH + // + + // Number of pipelined k-tiles in smem + constexpr int K_PIPE_MAX = DispatchPolicy::Stages; + + // NOTE: Another parameter: Partition the pipeline between active MMAs and active TMAs + // Tunable via the dispatch policy to tollerate latencies evenly across the math and compute stages + // K_PIPE_MMAS: The max number of active MMA pipes at beginning of every loop + // K_PIPE_TMAS: The max number of active TMA pipes at beginning of every loop (geq 1) + constexpr int K_PIPE_MMAS = DispatchPolicy::PipelineAsyncMmaStages; + constexpr int K_PIPE_TMAS = K_PIPE_MAX - K_PIPE_MMAS; + static_assert(0 <= K_PIPE_MMAS && K_PIPE_MMAS < K_PIPE_MAX); + static_assert(0 < K_PIPE_TMAS && K_PIPE_TMAS <= K_PIPE_MAX); + + static_assert(K_PIPE_MMAS < K_PIPE_MAX - 1); + + // Set the bytes transferred in this TMA transaction (may involve multiple issues) + constexpr uint32_t TmaTransactionBytes = static_cast( + (size<0>(sA) * size<1>(sA) * sizeof(InternalElementA)) + + (size<0>(sB) * size<1>(sB) * sizeof(InternalElementB))); + + + // Obtain warp index + int warp_idx = canonical_warp_idx(); + int warp_group_thread_idx = thread_idx % NumThreadsPerWarpGroup; + + PipelineParams params; + params.transaction_bytes = TmaTransactionBytes; + params.role = MainloopPipeline::ThreadCategory::ProducerConsumer; + params.is_leader = warp_group_thread_idx == 0; + params.num_consumers = NumThreadsPerWarpGroup; + + MainloopPipeline pipeline( + storage.pipeline_storage, + params); + + // State variables used for iterating the circular buffer + // smem_pipe_read / release is used by the consumer of SMEM data - i.e MMA + // smem_pipe_write is used by the producer of SMEM data - i.e TMA + PipelineState smem_pipe_read; + PipelineState smem_pipe_release; + PipelineState smem_pipe_write = mctlass::make_producer_start_state(); + + // We need this to guarantee that the Pipeline init is visible + // To all producers and consumer blocks in the Cluster + if constexpr (size(ClusterShape{}) > 1) { + cute::cluster_arrive_relaxed(); + cute::cluster_wait(); + } + else { + __syncthreads(); + } + + // Set predicate for the lowest lane_id in the warp + int lane_predicate = cute::elect_one_sync(); + + uint16_t mcast_mask_a = 0; + uint16_t mcast_mask_b = 0; + // Keep a copy to know when to stop issuing loads + int k_tile_count_tma = k_tile_count; + + // Issue TmaLoads (Prologue fetches) + if (warp_idx == 0 && lane_predicate == 1) { + // Maps the tile -> block, value + if constexpr (cute::is_same_v) { + auto block_layout = Layout{}; // (m,n) -> block_id + for (int n = 0; n < size<1>(block_layout); ++n) { + mcast_mask_a |= (uint16_t(1) << block_layout(cluster_local_block_id.x,n,Int<0>{})); + } + } + + if constexpr (cute::is_same_v) { + auto block_layout = Layout{}; // (m,n) -> block_id + for (int m = 0; m < size<0>(block_layout); ++m) { + mcast_mask_b |= (uint16_t(1) << block_layout(m,cluster_local_block_id.y,Int<0>{})); + } + } + + // Issue the prologue loads + int prologue_tma_count = min(K_PIPE_MAX, k_tile_count); + MCTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < prologue_tma_count; ++stage) { + pipeline.producer_acquire(smem_pipe_write); + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write); + + copy(tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,stage)); + copy(tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,stage)); + ++k_tile_iter; + ++smem_pipe_write; + } + k_tile_count_tma -= prologue_tma_count; + } + + // + // Define C accumulators and A/B partitioning + // + + TiledMma tiled_mma; + auto thread_mma = tiled_mma.get_thread_slice(thread_idx); + + Tensor tCsA = thread_mma.partition_A(sA); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCsB = thread_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE) + + // Allocate "fragments/descriptors" + Tensor tCrA = thread_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCrB = thread_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE) + + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum)); // M + CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N + CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K + CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE + CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tAsA)); // PIPE + CUTE_STATIC_ASSERT_V(size<3>(tCsB) == size<3>(tBsB)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sA)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sB)); // PIPE + + __syncthreads(); + + tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; + + warpgroup_fence_operand(accum); + // Prologue MMAs + MCTLASS_PRAGMA_UNROLL + for (int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count); + prologue_mma_count > 0; --prologue_mma_count) + { + // WAIT on smem_pipe_read until it's data is available + pipeline.consumer_wait(smem_pipe_read); + warpgroup_arrive(); + // Unroll the K mode manually to set scale D to 1 + MCTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + // (V,M,K) x (V,N,K) => (V,M,N) + cute::gemm(tiled_mma, tCrA(_,_,k_block,smem_pipe_read.index()), tCrB(_,_,k_block,smem_pipe_read.index()), accum); + tiled_mma.accumulate_ = GMMA::ScaleOut::One; + } + + warpgroup_commit_batch(); + ++smem_pipe_read; + --k_tile_count; + } + warpgroup_fence_operand(accum); + + // + // PIPELINED MAIN LOOP + // + + MCTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) + { + // WAIT on smem_pipe_read until data is available + pipeline.consumer_wait(smem_pipe_read); + + // + // Compute on k_tile + // + + warpgroup_fence_operand(accum); + warpgroup_arrive(); + // Unroll the K mode manually to set scale D to 1 + MCTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + // (V,M,K) x (V,N,K) => (V,M,N) + cute::gemm(tiled_mma, tCrA(_,_,k_block,smem_pipe_read.index()), tCrB(_,_,k_block,smem_pipe_read.index()), accum); + tiled_mma.accumulate_ = GMMA::ScaleOut::One; + } + warpgroup_commit_batch(); + + /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed + warpgroup_wait(); + warpgroup_fence_operand(accum); + + pipeline.consumer_release(smem_pipe_release); // UNLOCK wr stage, done _computing_ on it + + // + // Copy gmem to smem for *k_tile_iter + // + + // Do Acquire & Load only if needed - helps with both performance and also corner case illegal barrier-ops + if (warp_idx == 0 && lane_predicate == 1 && (k_tile_count_tma > 0) ) { + pipeline.producer_acquire(smem_pipe_write); // LOCK wr stage, for _writing_ + + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write); + + copy(tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,smem_pipe_write.index())); + copy(tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,smem_pipe_write.index())); + ++smem_pipe_write; + ++k_tile_iter; + --k_tile_count_tma; + } + + // Advance consumer pipeline + ++smem_pipe_read; + ++smem_pipe_release; + } + + // Wait on all GMMAs + warpgroup_wait<0>(); + warpgroup_fence_operand(accum); + + // Workaround for ensuring Smem destruction doesn't happen accidentally + if constexpr (size(typename DispatchPolicy::ClusterShape{}) > 1) { + cute::cluster_arrive(); + cute::cluster_wait(); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp b/csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp new file mode 100644 index 0000000..27335fd --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp @@ -0,0 +1,501 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/mctlass.h" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/arch/copy_sm90.hpp" +#include "mctlass/gemm/dispatch_policy.hpp" + +#include "cute/algorithm/functional.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/algorithm/gemm.hpp" +#include "cute/tensor_predicate.hpp" +#include "cute/numeric/arithmetic_tuple.hpp" +#include "mctlass/pipeline/pipeline.hpp" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::collective { +using namespace cute; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// WarpSpecialized Mainloop +template < + int Stages, + class ClusterShape, + class KernelSchedule, + class TileShape_, + class ElementA_, + class StrideA_, + class ElementB_, + class StrideB_, + class TiledMma_, + class GmemTiledCopyA_, + class SmemLayoutAtomA_, + class SmemCopyAtomA_, + class TransformA_, + class GmemTiledCopyB_, + class SmemLayoutAtomB_, + class SmemCopyAtomB_, + class TransformB_> +struct CollectiveMma< + MainloopSm90TmaGmmaWarpSpecialized, + TileShape_, + ElementA_, + StrideA_, + ElementB_, + StrideB_, + TiledMma_, + GmemTiledCopyA_, + SmemLayoutAtomA_, + SmemCopyAtomA_, + TransformA_, + GmemTiledCopyB_, + SmemLayoutAtomB_, + SmemCopyAtomB_, + TransformB_> +{ + // + // Type Aliases + // + using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecialized; + using TileShape = TileShape_; + using ElementA = ElementA_; + using StrideA = StrideA_; + using ElementB = ElementB_; + using StrideB = StrideB_; + using TiledMma = TiledMma_; + using ElementAccumulator = typename TiledMma::ValTypeC; + using GmemTiledCopyA = GmemTiledCopyA_; + using GmemTiledCopyB = GmemTiledCopyB_; + using SmemLayoutAtomA = SmemLayoutAtomA_; + using SmemLayoutAtomB = SmemLayoutAtomB_; + using SmemCopyAtomA = SmemCopyAtomA_; + using SmemCopyAtomB = SmemCopyAtomB_; + using TransformA = TransformA_; + using TransformB = TransformB_; + using ArchTag = typename DispatchPolicy::ArchTag; + + using MainloopPipeline = mctlass::PipelineTmaAsync< + DispatchPolicy::Stages, + typename DispatchPolicy::ClusterShape>; + using PipelineState = mctlass::PipelineState; + + using PipelineParams = typename MainloopPipeline::Params; + + static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)"); + static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape."); + + // Tile along K mode first before tiling over MN. PIPE mode last as usual. + // This maximizes TMA boxes due to better smem-K vectorization, reducing total issued TMAs. + using SmemLayoutA = decltype(tile_to_shape( + SmemLayoutAtomA{}, + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int{}), + Step<_2,_1,_3>{})); + using SmemLayoutB = decltype(tile_to_shape( + SmemLayoutAtomB{}, + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int{}), + Step<_2,_1,_3>{})); + + static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 2 or more."); + static_assert(cute::is_base_of::value && + cute::is_base_of::value, + "MMA atom must source both A and B operand from smem_desc for this mainloop."); + static_assert(cute::is_same_v || cute::is_same_v, + "GmemTiledCopy - invalid SM90 TMA copy atom specified."); + static_assert(cute::is_same_v || cute::is_same_v, + "GmemTiledCopy - invalid SM90 TMA copy atom specified."); + + // TMA converts f32 input to tf32 when copying from GMEM to SMEM + // For all other types, cast to size equivalent uint type to avoid any rounding by TMA. + static constexpr bool ConvertF32toTF32A = cute::is_same_v; + static constexpr bool ConvertF32toTF32B = cute::is_same_v; + using InternalElementA = cute::conditional_t>>; + using InternalElementB = cute::conditional_t>>; + + struct SharedStorage + { + struct TensorStorage : cute::aligned_struct<128> { + cute::array_aligned> smem_A; + cute::array_aligned> smem_B; + } tensors; + + using PipelineStorage = typename MainloopPipeline::SharedStorage; + PipelineStorage pipeline; + }; + using TensorStorage = typename SharedStorage::TensorStorage; + using PipelineStorage = typename SharedStorage::PipelineStorage; + + // Host side kernel arguments + struct Arguments { + ElementA const* ptr_A; + StrideA dA; + ElementB const* ptr_B; + StrideB dB; + }; + + // Device side kernel params + struct Params { + // Assumption: StrideA is congruent with Problem_MK + using TMA_A = decltype(make_tma_copy( + GmemTiledCopyA{}, + make_tensor(static_cast(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}), + SmemLayoutA{}(_,_,0), + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})), + size<1>(ClusterShape{}))); // mcast along N mode for this M load, if any + // Assumption: StrideB is congruent with Problem_NK + using TMA_B = decltype(make_tma_copy( + GmemTiledCopyB{}, + make_tensor(static_cast(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}), + SmemLayoutB{}(_,_,0), + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), + size<0>(ClusterShape{}))); // mcast along M mode for this N load, if any + TMA_A tma_load_a; + TMA_B tma_load_b; + }; + + // + // Methods + // + + template + static constexpr Params + to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) { + (void) workspace; + + // Optionally append _1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK) + auto problem_shape_MNKL = append<4>(problem_shape, Int<1>{}); + auto M = get<0>(problem_shape_MNKL); + auto N = get<1>(problem_shape_MNKL); + auto K = get<2>(problem_shape_MNKL); + auto L = get<3>(problem_shape_MNKL); + + auto ptr_A = reinterpret_cast(args.ptr_A); + auto ptr_B = reinterpret_cast(args.ptr_B); + + Tensor tensor_a = make_tensor(ptr_A, make_layout(make_shape(M,K,L), args.dA)); + Tensor tensor_b = make_tensor(ptr_B, make_layout(make_shape(N,K,L), args.dB)); + typename Params::TMA_A tma_load_a = make_tma_copy( + GmemTiledCopyA{}, + tensor_a, + SmemLayoutA{}(_,_,cute::Int<0>{}), + make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})), + size<1>(ClusterShape{})); // mcast along N mode for this M load, if any + typename Params::TMA_B tma_load_b = make_tma_copy( + GmemTiledCopyB{}, + tensor_b, + SmemLayoutB{}(_,_,cute::Int<0>{}), + make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), + size<0>(ClusterShape{})); // mcast along M mode for this N load, if any + return { + tma_load_a, + tma_load_b + }; + } + + static constexpr int K_PIPE_MAX = DispatchPolicy::Stages; + static constexpr int K_PIPE_MMAS = 1; + static constexpr uint32_t TmaTransactionBytes = + (size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast(sizeof(ElementA)))+ + (size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast(sizeof(ElementB))); + + /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance + MCTLASS_DEVICE + static void prefetch_tma_descriptors(Params const& mainloop_params) + { + cute::prefetch_tma_descriptor(mainloop_params.tma_load_a.get_tma_descriptor()); + cute::prefetch_tma_descriptor(mainloop_params.tma_load_b.get_tma_descriptor()); + } + + /// Perform a collective-scoped matrix multiply-accumulate + /// Producer Perspective + template < + class TensorA, class TMA_LOAD_A, + class TensorB, class TMA_LOAD_B, + class KTileIterator + > + MCTLASS_DEVICE void + load( + MainloopPipeline pipeline, + PipelineState smem_pipe_write, + TensorA const& gA, TMA_LOAD_A& tma_load_a, + TensorB const& gB, TMA_LOAD_B& tma_load_b, + KTileIterator k_tile_iter, int k_tile_count, + int thread_idx, + TensorStorage& shared_tensors) + { + + using namespace cute; + int warp_idx = canonical_warp_idx(); + int warp_idx_in_warp_group = warp_idx % 4; + int lane_predicate = cute::elect_one_sync(); + + if (warp_idx_in_warp_group == 0 and lane_predicate) { + Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + // + // Prepare the TMA loads for A and B + // + + dim3 cluster_local_block_id = cute::block_id_in_cluster(); + auto block_tma_a = tma_load_a.get_slice(cluster_local_block_id.y); + auto block_tma_b = tma_load_b.get_slice(cluster_local_block_id.x); + + // Applies the mapping from block_tma_a + Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k) + Tensor tAsA = block_tma_a.partition_D(sA); // (TMA,TMA_M,TMA_K,PIPE) + + Tensor tBgB = block_tma_b.partition_S(gB); // (TMA,TMA_N,TMA_K,k) + Tensor tBsB = block_tma_b.partition_D(sB); // (TMA,TMA_N,TMA_K,PIPE) + + uint16_t mcast_mask_a = 0; + uint16_t mcast_mask_b = 0; + + // Issue TmaLoads + // Maps the tile -> block, value + if constexpr (cute::is_same_v) { + auto block_layout = Layout{}; // (m,n) -> block_id + for (int n = 0; n < size<1>(block_layout); ++n) { + mcast_mask_a |= (uint16_t(1) << block_layout(cluster_local_block_id.x,n,Int<0>{})); + } + } + + if constexpr (cute::is_same_v) { + auto block_layout = Layout{}; // (m,n) -> block_id + for (int m = 0; m < size<0>(block_layout); ++m) { + mcast_mask_b |= (uint16_t(1) << block_layout(m,cluster_local_block_id.y,Int<0>{})); + } + } + + // Mainloop + MCTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) + { + // LOCK smem_pipe_write for _writing_ + pipeline.producer_acquire(smem_pipe_write); + + // + // Copy gmem to smem for *k_tile_iter + // + + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write); + + int write_stage = smem_pipe_write.index(); + copy(tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage)); + copy(tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage)); + ++k_tile_iter; + + // Advance smem_pipe_write + ++smem_pipe_write; + } + } + } + + /// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster + MCTLASS_DEVICE void + load_tail( + MainloopPipeline pipeline, + PipelineState smem_pipe_write) + { + int warp_idx = canonical_warp_idx(); + int warp_idx_in_warp_group = warp_idx % 4; + int lane_predicate = cute::elect_one_sync(); + + // Issue the epilogue waits + if (warp_idx_in_warp_group == 0 and lane_predicate) { + /* This helps avoid early exit of blocks in Cluster + * Waits for all stages to either be released (all + * Consumer UNLOCKs), or if the stage was never used + * then would just be acquired since the phase was + * still inverted from make_producer_start_state + */ + pipeline.producer_tail(smem_pipe_write); + } + } + + /// Perform a collective-scoped matrix multiply-accumulate + /// Consumer Perspective + template < + class FrgTensorC + > + MCTLASS_DEVICE void + mma(MainloopPipeline pipeline, + PipelineState smem_pipe_read, + FrgTensorC& accum, + int k_tile_count, + int thread_idx, + TensorStorage& shared_tensors, + Params const& mainloop_params) + { + using namespace cute; + + static_assert(is_rmem::value, "C tensor must be rmem resident."); + static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3."); + static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3."); + static_assert(cute::is_void_v, + "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions."); + static_assert(cute::is_void_v, + "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions."); + + Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE) + Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE) + + // + // Define C accumulators and A/B partitioning + // + + TiledMma tiled_mma; + auto thread_mma = tiled_mma.get_thread_slice(thread_idx); + + Tensor tCsA = thread_mma.partition_A(sA); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCsB = thread_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE) + + // Allocate "fragments/descriptors" + Tensor tCrA = thread_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCrB = thread_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE) + + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum)); // M + CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N + CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K + CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sA)); // PIPE + CUTE_STATIC_ASSERT_V(Int{} == size<2>(sB)); // PIPE + + // + // PIPELINED MAIN LOOP + // + static_assert((0 <= K_PIPE_MMAS) && (K_PIPE_MMAS < K_PIPE_MAX), + "ERROR : Incorrect number of MMAs in flight"); + + // We release buffers to producer warps(dma load) with some mmas in flight + PipelineState smem_pipe_release = smem_pipe_read; + + // Prologue GMMAs + int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count); + + tiled_mma.accumulate_ = GMMA::ScaleOut::Zero; + + warpgroup_fence_operand(accum); + MCTLASS_PRAGMA_UNROLL + for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue) + { + // WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value) + pipeline.consumer_wait(smem_pipe_read); + + int read_stage = smem_pipe_read.index(); + warpgroup_arrive(); + // Unroll the K mode manually to set scale D to 1 + MCTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + // (V,M,K) x (V,N,K) => (V,M,N) + cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accum); + tiled_mma.accumulate_ = GMMA::ScaleOut::One; + } + + warpgroup_commit_batch(); + + ++smem_pipe_read; + } + + warpgroup_fence_operand(accum); + // Mainloop GMMAs + k_tile_count -= prologue_mma_count; + + MCTLASS_PRAGMA_NO_UNROLL + for ( ; k_tile_count > 0; --k_tile_count) + { + // WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value) + pipeline.consumer_wait(smem_pipe_read); + + // + // Compute on k_tile + // + + int read_stage = smem_pipe_read.index(); + warpgroup_fence_operand(accum); + warpgroup_arrive(); + // Unroll the K mode manually to set scale D to 1 + MCTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) { + // (V,M,K) x (V,N,K) => (V,M,N) + cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accum); + tiled_mma.accumulate_ = GMMA::ScaleOut::One; + } + warpgroup_commit_batch(); + + /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed + warpgroup_wait(); + warpgroup_fence_operand(accum); + + // UNLOCK smem_pipe_release, done _computing_ on it + pipeline.consumer_release(smem_pipe_release); + + // Advance smem_pipe_read and smem_pipe_release + ++smem_pipe_read; + ++smem_pipe_release; + } + + warpgroup_fence_operand(accum); + } + + /// Perform a Consumer Epilogue to release all buffers + MCTLASS_DEVICE void + mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) { + // Prologue GMMAs + int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count); + k_tile_count -= prologue_mma_count; + + smem_pipe_release.advance(k_tile_count); + + // Wait on all GMMAs to complete + warpgroup_wait<0>(); + + for (int count = 0; count < prologue_mma_count; ++count) { + pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it + ++smem_pipe_release; + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/base_grouped.h b/csrc/mctlass/include/mctlass/gemm/device/base_grouped.h new file mode 100644 index 0000000..801acf7 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/base_grouped.h @@ -0,0 +1,478 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Base device-level grouped kernel. +*/ + +#pragma once + +#include +#include +#include + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/gemm_universal.h" + +#include "mctlass/gemm/kernel/default_gemm_universal.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +#include "mctlass/trace.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// GEMM Grouped +template +class BaseGrouped { +public: + + using BaseKernel = BaseKernel_; + + using ElementA = typename BaseKernel::ElementA; + using LayoutA = typename BaseKernel::LayoutA; + using TensorRefA = TensorRef; + static ComplexTransform const kTransformA = BaseKernel::kTransformA; + static int const kAlignmentA = BaseKernel::kAlignmentA; + + using ElementB = typename BaseKernel::ElementB; + using LayoutB = typename BaseKernel::LayoutB; + using TensorRefB = TensorRef; + static ComplexTransform const kTransformB = BaseKernel::kTransformB; + static int const kAlignmentB = BaseKernel::kAlignmentB; + + using ElementC = typename BaseKernel::ElementC; + using LayoutC = typename BaseKernel::LayoutC; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + static int const kAlignmentC = BaseKernel::kAlignmentC; + + using ElementAccumulator = typename BaseKernel::Mma::Policy::Operator::ElementC; + + using EpilogueOutputOp = typename BaseKernel::EpilogueOutputOp; + using ThreadblockSwizzle = typename BaseKernel::ThreadblockSwizzle; + + using Operator = typename BaseKernel::Operator; + using WarpMmaOperator = typename BaseKernel::Mma::Policy::Operator; + + using ArchMmaOperator = typename WarpMmaOperator::ArchMmaOperator; + using MathOperator = typename WarpMmaOperator::MathOperator; + using OperatorClass = typename WarpMmaOperator::OperatorClass; + using ArchTag = typename WarpMmaOperator::ArchTag; + using ThreadblockShape = typename BaseKernel::Mma::Shape; + using WarpShape = typename BaseKernel::WarpShape; + using InstructionShape = typename BaseKernel::InstructionShape; + static int const kStages = BaseKernel::Mma::kStages; + + /// Argument structure + using Arguments = typename BaseKernel::Arguments; + + using ProblemInfo = typename BaseKernel::ProblemVisitor::ProblemInfo; + +protected: + + /// Kernel parameters object + typename BaseKernel::Params params_; + +private: + + /// Get the number of tiles across all problems in a group + static int32_t group_tile_count(const mctlass::gemm::GemmCoord* problem_sizes_ptr, int problem_count) { + int32_t tiles = 0; + for (int32_t i = 0; i < problem_count; ++i) { + mctlass::gemm::GemmCoord problem = problem_sizes_ptr[i]; + BaseKernel::ProblemVisitor::possibly_transpose_problem(problem); + tiles += problem_tile_count(problem); + } + return tiles; + } + + /// Copy from `data` to `workspace` + Status copy_to_workspace(void* workspace, void* data, size_t bytes) { + mcError_t mc_error = mcMemcpy(workspace, data, bytes, mcMemcpyHostToDevice); + if (mc_error != mcSuccess) { + // Call mcGetLastError() to clear the error bit + mc_error = mcGetLastError(); + MCTLASS_TRACE_HOST( + " mcMemcpy() returned error " + << mcGetErrorString(mc_error)); + return Status::kErrorInternal; + } + + return Status::kSuccess; + } + + /// Precomputes scheduling information for the grouped GEMM + Status precompute(Arguments const &args, int32_t tile_count, void* workspace) { + size_t workspace_bytes = get_workspace_size(args); + std::vector host_workspace(workspace_bytes); + BaseKernel::ProblemVisitor::host_precompute(args.host_problem_sizes, + args.problem_count, + args.threadblock_count, + (void*)host_workspace.data()); + return copy_to_workspace(workspace, host_workspace.data(), workspace_bytes); + } + + /// Reorder `data` according to `indices` + template + static void reorder_array(T* data, const std::vector& indices) { + // For now, simply create a copy of the data and then copy over to the original. + std::vector copy(indices.size()); + for (size_t i = 0; i < indices.size(); ++i) { + copy.at(i) = data[indices[i]]; + } + + memcpy(data, copy.data(), indices.size() * sizeof(T)); + } + +public: + + /// Constructs the GEMM. + BaseGrouped() { } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + return BaseKernel::can_implement(args); + } + + /// Get the number of tiles in a problem + static int32_t problem_tile_count(mctlass::gemm::GemmCoord const &problem) { + auto grid = BaseKernel::ProblemVisitor::grid_shape(problem); + return BaseKernel::ProblemVisitor::tile_count(grid); + } + + /// Get the number of tiles across all problems in a group + static int32_t group_tile_count(Arguments const &args) { + if (args.host_problem_sizes == nullptr) { + MCTLASS_TRACE_HOST("Received nullptr for `args.host_problem_sizes"); + return -1; + } + + return group_tile_count(args.host_problem_sizes, args.problem_count); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + if (BaseKernel::ProblemVisitor::kRequiresPrecomputation) { + return BaseKernel::ProblemVisitor::get_workspace_size(args.host_problem_sizes, + args.problem_count, + args.threadblock_count); + } else { + return 0; + } + } + + /// Computes the grid shape + static dim3 get_grid_shape(Arguments const &args) { + + return dim3(args.threadblock_count, 1, 1); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int smem_capacity = -1) { + + MCTLASS_TRACE_HOST("BaseGrouped::maximum_active_blocks()"); + + int smem_size = int(sizeof(typename BaseKernel::SharedStorage)); + + MCTLASS_TRACE_HOST(" smem_size: " << smem_size << " bytes"); + + mcError_t result; + if (smem_size > (64 << 10)) { + result = mcFuncSetAttribute(Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + // Call mcGetLastError() to clear the error bit + result = mcGetLastError(); + MCTLASS_TRACE_HOST( + " mcFuncSetAttribute() returned error " + << mcGetErrorString(result)); + return -1; + } + } + + int max_active_blocks = -1; + result = mcOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks, + Kernel, + BaseKernel::kThreadCount, + smem_size); + + if (result != mcSuccess) { + // Call mcGetLastError() to clear the error bit + result = mcGetLastError(); + MCTLASS_TRACE_HOST( + " mcOccupancyMaxActiveBlocksPerMultiprocessor() returned error " + << mcGetErrorString(result)); + return -1; + } + + MCTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks); + return max_active_blocks; + } + + /// Sorts each pointer passed in according to the indices that sort + /// `problem_sizes_ptr` in descending order of problem-K dimension. + static void sort_problems(int problem_count, + mctlass::gemm::GemmCoord* problem_sizes_ptr, + int64_t* lda_host_ptr, + int64_t* ldb_host_ptr, + int64_t* ldc_host_ptr, + int64_t* ldd_host_ptr, + int64_t* offset_A_ptr, + int64_t* offset_B_ptr, + int64_t* offset_C_ptr, + int64_t* offset_D_ptr) + { + std::vector indices(problem_count); + std::iota(indices.begin(), indices.end(), 0); + std::stable_sort(indices.begin(), indices.end(), + [&problem_sizes_ptr](size_t i, size_t j) { + return problem_sizes_ptr[i].k() > problem_sizes_ptr[j].k(); + }); + + reorder_array(problem_sizes_ptr, indices); + reorder_array(lda_host_ptr, indices); + reorder_array(ldb_host_ptr, indices); + reorder_array(ldc_host_ptr, indices); + reorder_array(ldd_host_ptr, indices); + reorder_array(offset_A_ptr, indices); + reorder_array(offset_B_ptr, indices); + reorder_array(offset_C_ptr, indices); + reorder_array(offset_D_ptr, indices); + } + + /// Computes the number of threadblocks to launch for the grouped kernel + static int sufficient(const mctlass::gemm::GemmCoord* problem_sizes_ptr=nullptr, + int problem_count=0, + int available_sm_count=-1) { + // Determine the number of blocks that would be launched to fill up a single + // wave on the GPU with each SM having maximum occupancy. + mcDeviceProp_t properties; + int device_idx; + mcError_t result = mcGetDevice(&device_idx); + if (result != mcSuccess) { + // Call mcGetLastError() to clear the error bit + result = mcGetLastError(); + MCTLASS_TRACE_HOST(" mcGetDevice() returned error " + << mcGetErrorString(result)); + return 0; + } + + int multiprocessor_count; + result = mcDeviceGetAttribute(&multiprocessor_count, + mcDeviceAttributeMultiProcessorCount, device_idx); + if (result != mcSuccess) { + MCTLASS_TRACE_HOST( + " mcDeviceGetAttribute() returned error " + << mcGetErrorString(result)); + return 0; + } + + bool override_sm_count = (available_sm_count < 0 || available_sm_count > multiprocessor_count); + if (override_sm_count) { + available_sm_count = multiprocessor_count; + } + + int max_active_blocks = maximum_active_blocks(); + if (max_active_blocks <= 0) { + return 0; + } + + int occupancy_based_block_count = available_sm_count * max_active_blocks; + + if (problem_sizes_ptr == nullptr || problem_count == 0) { + return occupancy_based_block_count; + } + + int total_tiles = group_tile_count(problem_sizes_ptr, problem_count); + + // If the group contains a single problem, launching the exact number of + // threadblocks needed to cover the problem minimizes the work performed + // per threadblock in finding the next tile to compute. We return total_tiles + // unless the user has provided the SM count. + if (problem_count == 1 && override_sm_count) { + return total_tiles; + } + + // Choose between the full wave of threadblocks and the tile count. If there + // are fewer tiles in the group than threadblocks in the full wave, only + // some threadblocks will be assigned tiles. Those threadblocks + // which are not assigned tiles still need to perform the work of iterating through + // problem sizes to determine that they have no work to do. This competes for cycles + // with those threadblocks that are assigned tiles to compute. + return std::min(total_tiles, occupancy_based_block_count); + } + + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + MCTLASS_TRACE_HOST("BaseGrouped::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + // Workspace + size_t workspace_bytes = get_workspace_size(args); + + if (workspace_bytes && !workspace) { + return Status::kErrorWorkspaceNull; + } + + if (BaseKernel::ProblemVisitor::kRequiresPrecomputation) { + int32_t tile_count = group_tile_count(args); + Status status = precompute(args, tile_count, workspace); + if (status != Status::kSuccess) { + return status; + } + + params_ = typename BaseKernel::Params(args, workspace, tile_count); + } else { + params_ = typename BaseKernel::Params(args, workspace); + } + + // Specify shared memory capacity for kernel. + int smem_size = int(sizeof(typename BaseKernel::SharedStorage)); + + if (smem_size >= (64 << 10)) { + mcError_t result = mcFuncSetAttribute(Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + size_t workspace_bytes = get_workspace_size(args); + + if (workspace_bytes && !workspace) { + return Status::kErrorWorkspaceNull; + } + + if (BaseKernel::ProblemVisitor::kRequiresPrecomputation) { + int32_t tile_count = group_tile_count(args); + Status status = precompute(args, tile_count, workspace); + if (status != Status::kSuccess) { + return status; + } + + params_.update(args, workspace, tile_count); + } else { + params_.update(args, workspace); + } + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + // + // Configure grid and block dimensions + // + + if (!params_.problem_visitor.problem_count) { + return Status::kSuccess; + } + + dim3 grid(params_.threadblock_count, 1, 1); + dim3 block(BaseKernel::kThreadCount, 1, 1); + + int smem_size = int(sizeof(typename BaseKernel::SharedStorage)); + + // + // Launch kernel + // + + // Launch + mctlass::Kernel<<>>(params_); + + // + // Query for errors + // + mcError_t result = mcGetLastError(); + + if (result != mcSuccess) { + MCTLASS_TRACE_HOST(" grid launch failed with error " << mcGetErrorString(result)); + return Status::kErrorInternal; + } + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Initializes and runs the kernel. + Status operator()( + Arguments const &args, + void *workspace, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/default_gemm_configuration.h b/csrc/mctlass/include/mctlass/gemm/device/default_gemm_configuration.h new file mode 100644 index 0000000..87d6860 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/default_gemm_configuration.h @@ -0,0 +1,815 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Definitions for GEMM structures +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/arch/mma.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/epilogue/thread/linear_combination.h" +#include "mctlass/epilogue/thread/linear_combination_clamp.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename OperatorClass, + typename ArchTag, + typename ElementA, + typename ElementB, + typename ElementC, + typename ElementAccumulator +> +struct DefaultGemmConfiguration; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ArchTag, + typename ElementA, + typename ElementB, + typename ElementC, + typename ElementAccumulator> +struct DefaultGemmConfiguration< + arch::OpClassSimt, + ArchTag, + ElementA, + ElementB, + ElementC, + ElementAccumulator> { + + static int const kAlignmentA = 1; + static int const kAlignmentB = 1; + using ThreadblockShape = GemmShape<128, 128, 8>; + using WarpShape = GemmShape<32, 64, 8>; + using InstructionShape = GemmShape<1, 1, 1>; + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombination< + ElementC, + 1, + ElementAccumulator, + ElementAccumulator + >; + + using Operator = arch::OpMultiplyAdd; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ArchTag, + typename ElementC> +struct DefaultGemmConfiguration { + + static int const kAlignmentA = 4; + static int const kAlignmentB = 4; + using ThreadblockShape = GemmShape<128, 128, 32>; + using WarpShape = GemmShape<32, 64, 32>; + using InstructionShape = GemmShape<1, 1, 4>; + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, + 1, + int32_t, + float + >; + + using Operator = arch::OpMultiplyAdd; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ArchTag, + typename ElementA, + typename ElementB, + typename ElementC, + typename ElementAccumulator> +struct DefaultGemmConfiguration< + arch::OpClassWmmaTensorOp, + ArchTag, + ElementA, + ElementB, + ElementC, + ElementAccumulator> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombination< + ElementC, + 128 / sizeof_bits::value, + ElementAccumulator, + ElementAccumulator + >; + + using Operator = arch::OpMultiplyAdd; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementA, + typename ElementB, + typename ElementC, + typename ElementAccumulator> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm70, + ElementA, + ElementB, + ElementC, + ElementAccumulator> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 32>; + using WarpShape = GemmShape<64, 64, 32>; + using InstructionShape = GemmShape<8, 8, 4>; + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombination< + ElementC, + 128 / sizeof_bits::value, + ElementAccumulator, + ElementAccumulator + >; + + using Operator = arch::OpMultiplyAdd; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementA, + typename ElementB, + typename ElementC, + typename ElementAccumulator> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm75, + ElementA, + ElementB, + ElementC, + ElementAccumulator> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + using ThreadblockShape = GemmShape<128, 256, 32>; + using WarpShape = GemmShape<64, 64, 32>; + using InstructionShape = GemmShape<16, 8, 8>; + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombination< + ElementC, + 128 / sizeof_bits::value, + ElementAccumulator, + ElementAccumulator + >; + + using Operator = typename platform::conditional< + (platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value), + arch::OpMultiplyAddSaturate, arch::OpMultiplyAdd>::type; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm75, + int8_t, + int8_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 64>; + using WarpShape = GemmShape<64, 64, 64>; + using InstructionShape = GemmShape<8, 8, 16>; + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm75, + int8_t, + uint8_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 64>; + using WarpShape = GemmShape<64, 64, 64>; + using InstructionShape = GemmShape<8, 8, 16>; + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm75, + uint8_t, + int8_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 64>; + using WarpShape = GemmShape<64, 64, 64>; + using InstructionShape = GemmShape<8, 8, 16>; + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm75, + uint8_t, + uint8_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 64>; + using WarpShape = GemmShape<64, 64, 64>; + using InstructionShape = GemmShape<8, 8, 16>; + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm75, + int4b_t, + int4b_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 128>; + using WarpShape = GemmShape<64, 64, 128>; + using InstructionShape = GemmShape<8, 8, 32>; + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm75, + int4b_t, + uint4b_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 128>; + using WarpShape = GemmShape<64, 64, 128>; + using InstructionShape = GemmShape<8, 8, 32>; + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm75, + uint4b_t, + int4b_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 128>; + using WarpShape = GemmShape<64, 64, 128>; + using InstructionShape = GemmShape<8, 8, 32>; + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm75, + uint4b_t, + uint4b_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 128>; + using WarpShape = GemmShape<64, 64, 128>; + using InstructionShape = GemmShape<8, 8, 32>; + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm75, + uint1b_t, + uint1b_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 512>; + using WarpShape = GemmShape<64, 64, 512>; + using InstructionShape = GemmShape<8, 8, 128>; + static int const kStages = 2; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpXorPopc; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template +struct DefaultGemmConfiguration { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 64>; + using WarpShape = GemmShape<64, 64, 64>; + using InstructionShape = GemmShape<16, 8, 16>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombination< + ElementC, 128 / sizeof_bits::value, ElementAccumulator, + ElementAccumulator>; + + using Operator = typename platform::conditional< + (platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value), + arch::OpMultiplyAddSaturate, arch::OpMultiplyAdd>::type; +}; + +//////////////////////////////////////////////////////////////////////////////// +template +struct DefaultGemmConfiguration { + + static int const kAlignmentA = 1; + static int const kAlignmentB = 1; + + using ThreadblockShape = GemmShape<128, 128, 16>; + using WarpShape = GemmShape<32, 64, 16>; + using InstructionShape = GemmShape<8, 8, 4>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombination< + ElementC, 128 / sizeof_bits::value, ElementAccumulator, + ElementAccumulator>; + + using Operator = arch::OpMultiplyAdd; +}; + + +template <> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm80, + complex, + complex, + complex, + complex + > { + + static int const kAlignmentA = 1; + static int const kAlignmentB = 1; + + using ThreadblockShape = GemmShape<64, 64, 16>; + using WarpShape = GemmShape<32, 32, 16>; + using InstructionShape = GemmShape<8, 8, 4>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombination< + complex, 1, complex, + complex>; + + using Operator = arch::OpMultiplyAddComplex; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm80, + int8_t, + int8_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 64>; + using WarpShape = GemmShape<64, 64, 64>; + using InstructionShape = GemmShape<16, 8, 32>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm80, + int8_t, + uint8_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 64>; + using WarpShape = GemmShape<64, 64, 64>; + using InstructionShape = GemmShape<16, 8, 32>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm80, + uint8_t, + int8_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 64>; + using WarpShape = GemmShape<64, 64, 64>; + using InstructionShape = GemmShape<16, 8, 32>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm80, + uint8_t, + uint8_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 64>; + using WarpShape = GemmShape<64, 64, 64>; + using InstructionShape = GemmShape<16, 8, 32>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm80, + int4b_t, + int4b_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 128>; + using WarpShape = GemmShape<64, 64, 128>; + using InstructionShape = GemmShape<16, 8, 64>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm80, + int4b_t, + uint4b_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 128>; + using WarpShape = GemmShape<64, 64, 128>; + using InstructionShape = GemmShape<16, 8, 64>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm80, + uint4b_t, + int4b_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 128>; + using WarpShape = GemmShape<64, 64, 128>; + using InstructionShape = GemmShape<16, 8, 64>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm80, + uint4b_t, + uint4b_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 128>; + using WarpShape = GemmShape<64, 64, 128>; + using InstructionShape = GemmShape<16, 8, 64>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAddSaturate; +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementC> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm80, + uint1b_t, + uint1b_t, + ElementC, + int32_t> { + + static int const kAlignmentA = 128 / sizeof_bits::value; + static int const kAlignmentB = 128 / sizeof_bits::value; + + using ThreadblockShape = GemmShape<128, 256, 512>; + using WarpShape = GemmShape<64, 64, 512>; + using InstructionShape = GemmShape<16, 8, 256>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombinationClamp< + ElementC, 128 / sizeof_bits::value, int32_t, float>; + + using Operator = arch::OpMultiplyAdd; +}; + +//////////////////////////////////////////////////////////////////////////////// +template +struct DefaultGemmConfiguration { + + static int const kAlignmentA = 1; + static int const kAlignmentB = 1; + + using ThreadblockShape = GemmShape<128, 256, 64>; + using WarpShape = GemmShape<64, 64, 64>; + using InstructionShape = GemmShape<16, 8, 4>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombination< + ElementC, 128 / sizeof_bits::value, ElementAccumulator, + ElementAccumulator>; + + using Operator = arch::OpMultiplyAdd; +}; + +template <> +struct DefaultGemmConfiguration< + arch::OpClassTensorOp, + arch::Sm90, + complex, + complex, + complex, + complex + > { + + static int const kAlignmentA = 1; + static int const kAlignmentB = 1; + + using ThreadblockShape = GemmShape<64, 64, 16>; + using WarpShape = GemmShape<32, 32, 16>; + using InstructionShape = GemmShape<16, 8, 4>; + static int const kStages = 3; + + using EpilogueOutputOp = epilogue::thread::LinearCombination< + complex, 1, complex, + complex>; + + using Operator = arch::OpMultiplyAddComplex; +}; + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/ell_gemm.h b/csrc/mctlass/include/mctlass/gemm/device/ell_gemm.h new file mode 100644 index 0000000..b62ced9 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/ell_gemm.h @@ -0,0 +1,848 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a Block-Ell sparse gemm kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/ell_gemm.h" + +#include "mctlass/gemm/kernel/default_ell_gemm.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/*! Blocked-Ell sparse gemm device-level operator. This is an interface to efficient MCTLASS + Blocked-Ell kernels that may be invoked from host code. + + The contributions of this class are: + + 1. At compile time, it maps data types and high-level structural parameters onto + specific MCTLASS components. + + 2. At runtime, it maps logical arguments to Blocked-Ell problems to kernel parameters. + + 3. At runtime, it launches kernels on the device. + + Example of a MCTLASS EllGemm operator is as follows: + + // + // Instantiate the MCTLASS EllGemm operator. + // + + mctlass::gemm::device::EllGemm< + mctlass::half_t, + mctlass::layout::RowMajor, + mctlass::half_t, + mctlass::layout::ColumnMajor, + mctlass::half_t, + mctlass::layout::ColumnMajor, + float, + mctlass::arch::OpClassTensorOp, + mctlass::arch::Sm80, + mctlass::gemm::GemmShape<128, 128, 32>, + mctlass::gemm::GemmShape<64, 64, 32>, + mctlass::gemm::GemmShape<16, 8, 16>, + mctlass::epilogue::thread::LinearCombination< + mctlass::half_t, 128 / mctlass::sizeof_bits::value, + float, float>, + mctlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<8>, + 4, // Stages + 128 / mctlass::sizeof_bits::value, // Alignment A + 128 / mctlass::sizeof_bits::value // Alignment B + > ellgemm_op; + + // + // Launch the EllGemm operation on the device + // + + Description of parameters and tensors used to represent the Blocked-Ellpack (ELL) format: + a_rows - Rows in the sparse matrix. + a_cols - Colums in the sparse matrix. + BlockedEllA - Packed matrix (ellValue matrix) that stores non-zero values in + consecutive blocks, whose size is (a_rows * a_ell_num_columns) + ell_idx - Blocked-ELL Column indices (ellColInd) matrix, whose size is + (a_rows / a_ell_blocksize) * (a_ell_num_columns / a_ell_blocksize) + a_ell_blocksize - Size of the ELL-Blocks. + a_ell_num_columns - Number of columns in the Blocked-Ellpack format (ellValue columns) + B - Input dense matrix whose size is (a_cols * n) + C/D - Output dense matrix whose size is (a_rows * n) + + mctlass::Status status = ellgemm_op({ + {a_rows, n, a_cols}, // GemmCoord problem_size + {BlockedEllA, lda}, // TensorRef ref_BlockedEllA + {B, ldb}, // TensorRef ref_B, + {C, ldc}, // TensorRef ref_C, + {D, ldd}, // TensorRef ref_D, + ell_idx, // Blocked-ELL Column indices or ellColInd matrix (const int*) + a_ell_num_columns, // Columns in the Blocked-Ellpack (ellValue) matrix (int) + a_ell_blocksize, // Size of the ELL-Blocks (int) + a_ell_base, // Base index of ellColInd (int) - Zero or One + {alpha, beta} // EpilogueOutputOp::Params epilogue_op_params + }); + + A simplified view of the template is listed below. + + template < + /// Element type for A matrix operand + typename ElementA, + + /// Layout type for A matrix operand + typename LayoutA, + + /// Element type for B matrix operand + typename ElementB, + + /// Layout type for B matrix operand + typename LayoutB, + + /// Element type for C and D matrix operands + typename ElementC, + + /// Layout type for C and D matrix operands + typename LayoutC, + + /// Element type for internal accumulation + typename ElementAccumulator, + + /// Operator class tag + typename OperatorClass, + + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag, + + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + + /// Epilogue output operator + typename EpilogueOutputOp, + + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + + /// Number of stages used in the pipelined mainloop + int Stages + + /// Access granularity of A matrix in units of elements + int AlignmentA, + + /// Access granularity of B matrix in units of elements + int AlignmentB, + + /// Supports split-K with serial reduction + bool SplitKSerial, + + /// Operation performed by GEMM + typename Operator, + + /// Sparse matrix is A or not + bool IsASparse + > + class EllGemm; +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassTensorOp, + /// Tag indicating architecture to tune for + typename ArchTag_ = arch::Sm80, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = + typename threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial = false, + /// Operation performed by GEMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator, + /// Sparse matrix is A or not + bool IsASparse = true + > +class EllGemm { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = SplitKSerial; + static ComplexTransform const kTransformA = ComplexTransform::kNone; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + static bool const kIsASparse = IsASparse; + + /// Define the kernel + using GemmKernel = typename kernel::DefaultEllGemm< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + kSplitKSerial, + Operator, + kIsASparse + >::GemmKernel; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + TensorRef ref_B; + TensorRef ref_C; + TensorRef ref_D; + const int* ell_idx; + int ell_ncol; + int ell_blocksize; + int ell_base_idx; + typename EpilogueOutputOp::Params epilogue; + int split_k_slices; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments(): problem_size(0, 0, 0), split_k_slices(1) { + + } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + TensorRef ref_A_, + TensorRef ref_B_, + TensorRef ref_C_, + TensorRef ref_D_, + const int* ell_idx_, + int ell_ncol_, + int ell_blocksize_, + int ell_base_idx_, + typename EpilogueOutputOp::Params epilogue_ = + typename EpilogueOutputOp::Params(), + int split_k_slices = 1 + ): + problem_size(problem_size_), + ref_A(ref_A_), + ref_B(ref_B_), + ref_C(ref_C_), + ref_D(ref_D_), + ell_idx(ell_idx_), + ell_ncol(ell_ncol_), + ell_blocksize(ell_blocksize_), + ell_base_idx(ell_base_idx_), + epilogue(epilogue_), + split_k_slices(split_k_slices) { + + } + }; + +private: + + /// Kernel parameters object + typename GemmKernel::Params params_; + +public: + + /// Constructs the GEMM. + EllGemm() { } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + if (!kSplitKSerial && args.split_k_slices > 1) { + return Status::kErrorInvalidProblem; + } + + Status status = GemmKernel::can_implement( + args.problem_size, + args.ref_A.non_const_ref(), + args.ref_B.non_const_ref(), + args.ref_C.non_const_ref(), + args.ref_D + ); + + if (status != Status::kSuccess) { + return status; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + size_t bytes = 0; + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {args.ell_blocksize, + ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + tiled_shape.m() *= (args.ell_blocksize + ThreadblockShape::kM - 1 ) / ThreadblockShape::kM; + + if (kSplitKSerial && args.split_k_slices > 1) { + + bytes += sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n()); + } + + return bytes; + } + + Status set(Arguments const &args, mctlass::gemm::GemmCoord const &grid_shape, void *workspace){ + // Initialize the Params structure + params_ = typename GemmKernel::Params{ + args.problem_size, + grid_shape, + args.ref_A.non_const_ref(), + args.ref_B.non_const_ref(), + args.ref_C.non_const_ref(), + args.ref_D, + args.ell_idx, + args.ell_ncol, + args.ell_blocksize, + args.ell_base_idx, + args.epilogue, + static_cast(workspace) + }; + return Status::kSuccess; + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {args.ell_blocksize, ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + grid_shape.m() *= (args.ell_blocksize + ThreadblockShape::kM - 1 ) / ThreadblockShape::kM; + + if (kSplitKSerial) { + if (args.split_k_slices > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + size_t bytes = get_workspace_size(args); + + mcError_t result = mcMemsetAsync(workspace, 0, bytes, stream); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + } + else { + + if (args.split_k_slices > 1) { + return Status::kErrorInvalidProblem; + } + } + + return set(args, grid_shape, workspace); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + if (kSplitKSerial && args.split_k_slices > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + } + + params_.ref_A.reset(args.ref_A.non_const_ref().data()); + params_.ref_B.reset(args.ref_B.non_const_ref().data()); + params_.ref_C.reset(args.ref_C.non_const_ref().data()); + params_.ref_D.reset(args.ref_D.data()); + params_.output_op = args.epilogue; + params_.semaphore = static_cast(workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(GemmKernel::kThreadCount, 1, 1); + + mcError_t result; + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + + if (smem_size >= (48 << 10)) { + result = mcFuncSetAttribute(Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + mctlass::Kernel<<>>(params_); + + result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output exchanges problem size and operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB, + /// If true, kernel supports split-K as a serial reduction + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator_, + /// Sparse matrix is A or not + bool IsASparse> +class EllGemm { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static ComplexTransform const kTransformA = ComplexTransform::kNone; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + static bool const kSplitKSerial = SplitKSerial; + static bool const kIsASparse = false; + + using UnderlyingOperator = EllGemm< + ElementB, + typename layout::LayoutTranspose::type, + ElementA, + typename layout::LayoutTranspose::type, + ElementC, + layout::RowMajor, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + kAlignmentB, + kAlignmentA, + SplitKSerial, + Operator, + kIsASparse + >; + + using UnderlyingArguments = typename UnderlyingOperator::Arguments; + using GemmKernel = typename UnderlyingOperator::GemmKernel; + static int const kAlignmentC = UnderlyingOperator::kAlignmentC; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + TensorRef ref_B; + TensorRef ref_C; + TensorRef ref_D; + const int* ell_idx; + int ell_ncol; + int ell_blocksize; + int ell_base_idx; + typename EpilogueOutputOp::Params epilogue; + int split_k_slices; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + TensorRef ref_A_, + TensorRef ref_B_, + TensorRef ref_C_, + TensorRef ref_D_, + const int* ell_idx_, + int ell_ncol_, + int ell_blocksize_, + int ell_base_idx_, + typename EpilogueOutputOp::Params epilogue_ = + typename EpilogueOutputOp::Params(), + int split_k_slices = 1 + ): + problem_size(problem_size_), + ref_A(ref_A_), + ref_B(ref_B_), + ref_C(ref_C_), + ref_D(ref_D_), + ell_idx(ell_idx_), + ell_ncol(ell_ncol_), + ell_blocksize(ell_blocksize_), + ell_base_idx(ell_base_idx_), + epilogue(epilogue_), + split_k_slices(split_k_slices) { } + }; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the GEMM. + EllGemm() { } + + /// Helper to construct a transposed equivalent for the underying GEMM operator + static UnderlyingArguments to_underlying_arguments(Arguments const &args) { + return UnderlyingArguments( + {args.problem_size.n(), args.problem_size.m(), args.problem_size.k()}, + {args.ref_B.data(), args.ref_B.stride(0)}, + {args.ref_A.data(), args.ref_A.stride(0)}, + {args.ref_C.data(), args.ref_C.stride(0)}, + {args.ref_D.data(), args.ref_D.stride(0)}, + args.ell_idx, + args.ell_ncol, + args.ell_blocksize, + args.ell_base_idx, + args.epilogue, + args.split_k_slices + ); + } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + size_t bytes = 0; + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, args.ell_blocksize, ThreadblockShape::kK}, + args.split_k_slices); + + tiled_shape.n() *= (args.ell_blocksize + ThreadblockShape::kN - 1 ) / ThreadblockShape::kN; + + if (kSplitKSerial && args.split_k_slices > 1) { + + bytes += sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n()); + } + + return bytes; + } + + Status set(Arguments const &args, mctlass::gemm::GemmCoord const &grid_shape, void *workspace){ + // Initialize the Params structure + return underlying_operator_.set(to_underlying_arguments(args), grid_shape, workspace); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape( + {args.problem_size.n(), args.problem_size.m(), args.problem_size.k()}, + {ThreadblockShape::kM, args.ell_blocksize, ThreadblockShape::kK}, + args.split_k_slices); + + grid_shape.n() *= (args.ell_blocksize + ThreadblockShape::kN - 1 ) / ThreadblockShape::kN; + + if (kSplitKSerial) { + if (args.split_k_slices > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + size_t bytes = get_workspace_size(args); + + mcError_t result = mcMemsetAsync(workspace, 0, bytes, stream); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + } + else { + + if (args.split_k_slices > 1) { + return Status::kErrorInvalidProblem; + } + } + + // Initialize the Params structure + set(args, grid_shape, workspace); + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm.h b/csrc/mctlass/include/mctlass/gemm/device/gemm.h new file mode 100644 index 0000000..e305dc3 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm.h @@ -0,0 +1,771 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the domcmentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/gemm.h" + +#include "mctlass/gemm/kernel/default_gemm.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +#include "mctlass/layout/permute.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/*! Gemm device-level operator. This is an interface to efficient MCTLASS GEMM kernels that may + be invoked from host code. + + The contributions of this class are: + + 1. At compile time, it maps data types and high-level structural parameters onto + specific MCTLASS components. + + 2. At runtime, it maps logical arguments to GEMM problems to kernel parameters. + + 3. At runtime, it launches kernels on the device. + + The intent is to provide a convenient mechanism for interacting with most plausible GEMM + configurations for each supported architecture. Consequently, not all parameters are exposed + to the top-level interface. Rather, sensible defaults at each level of the MCTLASS hierarchy + are selected to tradeoff simplicity of the interface with flexibility. We expect + most configurations to be specified at this level. Applications with more exotic requirements + may construct their kernels of interest using MCTLASS components at the threadblock, warp, + and thread levels of abstraction. + + MCTLASS exposes computations using the functor design pattern in which objects compose some + internal state with an overloaded function call operator. This enables decoupling of + initialization from exemction, possibly reducing overhead during steady state phases of + application exemction. + + MCTLASS device-level operators expose an Arguments structure encompassing each logical + input to the computation. This is distinct from the kernel-level Params structure pattern + which contains application-specific precomputed state needed by the device code. + + Example of a MCTLASS GEMM operator implementing the functionality of mcBLAS's SGEMM NN + is as follows: + + // + // Instantiate the MCTLASS GEMM operator. + // + + mctlass::gemm::device::Gemm< + float, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::ColumnMajor + > gemm_op; + + // + // Launch the GEMM operation on the device + // + + mctlass::Status status = gemm_op({ + {m, n, k}, // GemmCoord problem_size, + {A, lda}, // TensorRef ref_A, + {B, ldb}, // TensorRef ref_B, + {C, ldc}, // TensorRef ref_C, + {D, ldd}, // TensorRef ref_D, + {alpha, beta} // EpilogueOutputOp::Params epilogue_op_params + }); + + + A simplified view of the template is listed below. + + template < + /// Element type for A matrix operand + typename ElementA, + + /// Layout type for A matrix operand + typename LayoutA, + + /// Element type for B matrix operand + typename ElementB, + + /// Layout type for B matrix operand + typename LayoutB, + + /// Element type for C and D matrix operands + typename ElementC, + + /// Layout type for C and D matrix operands + typename LayoutC, + + /// Element type for internal acmcmulation + typename ElementAccumulator, + + /// Operator class tag + typename OperatorClass, + + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag, + + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + + /// Epilogue output operator + typename EpilogueOutputOp, + + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + + /// Number of stages used in the pipelined mainloop + int Stages + > + class Gemm; +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal acmcmulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassSimt, + /// Tag indicating architecture to tune for + typename ArchTag_ = arch::Sm70, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = + typename threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial = false, + /// Operation performed by GEMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator, + /// Gather operand A by using an index array + bool GatherA = false, + /// Gather operand B by using an index array + bool GatherB = false, + /// Scatter result D by using an index array + bool ScatterD = false, + /// Permute result D + typename PermuteDLayout = layout::NoPermute> +class Gemm { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = SplitKSerial; + static ComplexTransform const kTransformA = ComplexTransform::kNone; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + + /// Define the kernel + using GemmKernel = typename kernel::DefaultGemm< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + kSplitKSerial, + Operator, + SharedMemoryClearOption::kNone, + GatherA, + GatherB, + ScatterD, + PermuteDLayout + >::GemmKernel; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + TensorRef ref_B; + TensorRef ref_C; + TensorRef ref_D; + typename EpilogueOutputOp::Params epilogue; + int split_k_slices; + // For gather+scatter operations + int const *gather_A_indices; + int const *gather_B_indices; + int const *scatter_D_indices; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments(): problem_size(0, 0, 0), split_k_slices(1) { + + } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + TensorRef ref_A_, + TensorRef ref_B_, + TensorRef ref_C_, + TensorRef ref_D_, + typename EpilogueOutputOp::Params epilogue_ = + typename EpilogueOutputOp::Params(), + int split_k_slices = 1, + int const *gather_A_indices_ = nullptr, + int const *gather_B_indices_ = nullptr, + int const *scatter_D_indices_ = nullptr + ): + problem_size(problem_size_), + ref_A(ref_A_), + ref_B(ref_B_), + ref_C(ref_C_), + ref_D(ref_D_), + epilogue(epilogue_), + split_k_slices(split_k_slices), + gather_A_indices(gather_A_indices_), + gather_B_indices(gather_B_indices_), + scatter_D_indices(scatter_D_indices_) { + + } + }; + +private: + + /// Kernel parameters object + typename GemmKernel::Params params_; + +public: + + /// Constructs the GEMM. + Gemm() { } + + /// Determines whether the GEMM can exemcte the given problem. + static Status can_implement(Arguments const &args) { + + if (!kSplitKSerial && args.split_k_slices > 1) { + return Status::kErrorInvalidProblem; + } + + Status status = GemmKernel::can_implement( + args.problem_size, + args.ref_A.non_const_ref(), + args.ref_B.non_const_ref(), + args.ref_C.non_const_ref(), + args.ref_D + ); + + if (status != Status::kSuccess) { + return status; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + size_t bytes = 0; + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + if (kSplitKSerial && args.split_k_slices > 1) { + + bytes += sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n()); + } + + return bytes; + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + if (kSplitKSerial) { + if (args.split_k_slices > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + size_t bytes = get_workspace_size(args); + + mcError_t result = mcMemsetAsync(workspace, 0, bytes, stream); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + } + else { + + if (args.split_k_slices > 1) { + return Status::kErrorInvalidProblem; + } + } + + // Initialize the Params structure + params_ = typename GemmKernel::Params{ + args.problem_size, + grid_shape, + args.ref_A.non_const_ref(), + args.ref_B.non_const_ref(), + args.ref_C.non_const_ref(), + args.ref_D, + args.epilogue, + static_cast(workspace), + args.gather_A_indices, + args.gather_B_indices, + args.scatter_D_indices + }; + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + if (kSplitKSerial && args.split_k_slices > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + } + + params_.ref_A.reset(args.ref_A.non_const_ref().data()); + params_.ref_B.reset(args.ref_B.non_const_ref().data()); + params_.ref_C.reset(args.ref_C.non_const_ref().data()); + params_.ref_D.reset(args.ref_D.data()); + params_.output_op = args.epilogue; + params_.semaphore = static_cast(workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(GemmKernel::kThreadCount, 1, 1); + + mcError_t result; + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + + if (smem_size >= (64 << 10)) { + result = mcFuncSetAttribute(Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + mctlass::Kernel<<>>(params_); + + result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output exchanges problem size and operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal acmcmulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB, + /// If true, kernel supports split-K as a serial reduction + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator_, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Scatter result D by using an index array + bool ScatterD, + /// Permute result D + typename PermuteDLayout +> +class Gemm { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static ComplexTransform const kTransformA = ComplexTransform::kNone; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + static bool const kSplitKSerial = SplitKSerial; + + using UnderlyingOperator = Gemm< + ElementB, + typename layout::LayoutTranspose::type, + ElementA, + typename layout::LayoutTranspose::type, + ElementC, + layout::RowMajor, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + kAlignmentB, + kAlignmentA, + SplitKSerial, + Operator, + GatherB, + GatherA, + ScatterD, + PermuteDLayout + >; + + using UnderlyingArguments = typename UnderlyingOperator::Arguments; + using GemmKernel = typename UnderlyingOperator::GemmKernel; + static int const kAlignmentC = UnderlyingOperator::kAlignmentC; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + TensorRef ref_B; + TensorRef ref_C; + TensorRef ref_D; + typename EpilogueOutputOp::Params epilogue; + int split_k_slices; + // For gather+scatter operations + int *gather_A_indices; + int *gather_B_indices; + int *scatter_D_indices; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + TensorRef ref_A_, + TensorRef ref_B_, + TensorRef ref_C_, + TensorRef ref_D_, + typename EpilogueOutputOp::Params epilogue_ = + typename EpilogueOutputOp::Params(), + int split_k_slices = 1, + int *gather_A_indices_ = nullptr, + int *gather_B_indices_ = nullptr, + int *scatter_D_indices_ = nullptr + ): + problem_size(problem_size_), + ref_A(ref_A_), + ref_B(ref_B_), + ref_C(ref_C_), + ref_D(ref_D_), + epilogue(epilogue_), + split_k_slices(split_k_slices), + gather_A_indices(gather_A_indices_), + gather_B_indices(gather_B_indices_), + scatter_D_indices(scatter_D_indices_) { } + }; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the GEMM. + Gemm() { } + + /// Helper to construct a transposed equivalent for the underying GEMM operator + static UnderlyingArguments to_underlying_arguments(Arguments const &args) { + return UnderlyingArguments( + {args.problem_size.n(), args.problem_size.m(), args.problem_size.k()}, + {args.ref_B.data(), args.ref_B.stride(0)}, + {args.ref_A.data(), args.ref_A.stride(0)}, + {args.ref_C.data(), args.ref_C.stride(0)}, + {args.ref_D.data(), args.ref_D.stride(0)}, + args.epilogue, + args.split_k_slices, + args.gather_B_indices, + args.gather_A_indices, + args.scatter_D_indices + ); + } + + /// Determines whether the GEMM can exemcte the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_array.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_array.h new file mode 100644 index 0000000..a341dfd --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_array.h @@ -0,0 +1,737 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the domcmentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/gemm_array.h" + +#include "mctlass/gemm/kernel/default_gemm.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +//////////////////////////////////////////////////////////////////////////////// + +/*! Gemm device-level operator. This is an interface to efficient MCTLASS GEMM kernels that may + be invoked from host code. + + The contributions of this class are: + + 1. At compile time, it maps data types and high-level structural parameters onto + specific MCTLASS components. + + 2. At runtime, it maps logical arguments to GEMM problems to kernel parameters. + + 3. At runtime, it launches kernels on the device. + + The intent is to provide a convenient mechanism for interacting with most plausible GEMM + configurations for each supported architecture. Consequently, not all parameters are exposed + to the top-level interface. Rather, sensible defaults at each level of the MCTLASS hierarchy + are selected to tradeoff simplicity of the interface with flexibility. We expect + most configurations to be specified at this level. Applications with more exotic requirements + may construct their kernels of interest using MCTLASS components at the threadblock, warp, + and thread levels of abstraction. + + MCTLASS exposes computations using the functor design pattern in which objects compose some + internal state with an overloaded function call operator. This enables decoupling of + initialization from exemction, possibly reducing overhead during steady state phases of + application exemction. + + MCTLASS device-level operators expose an Arguments structure encompassing each logical + input to the computation. This is distinct from the kernel-level Params structure pattern + which contains application-specific precomputed state needed by the device code. + + Example of a MCTLASS GEMM operator implementing the functionality of mcBLAS's SGEMM NN + is as follows: + + // + // Instantiate the MCTLASS GEMM operator. + // + + mctlass::gemm::device::Gemm< + float, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::ColumnMajor + > gemm_op; + + // + // Launch the GEMM operation on the device + // + + mctlass::Status status = gemm_op({ + {m, n, k}, // GemmCoord problem_size, + {A, lda}, // TensorRef ref_A, + {B, ldb}, // TensorRef ref_B, + {C, ldc}, // TensorRef ref_C, + {D, ldd}, // TensorRef ref_D, + {alpha, beta} // EpilogueOutputOp::Params epilogue_op_params + }); + + + A simplified view of the template is listed below. + + template < + /// Element type for A matrix operand + typename ElementA, + + /// Layout type for A matrix operand + typename LayoutA, + + /// Element type for B matrix operand + typename ElementB, + + /// Layout type for B matrix operand + typename LayoutB, + + /// Element type for C and D matrix operands + typename ElementC, + + /// Layout type for C and D matrix operands + typename LayoutC, + + /// Element type for internal acmcmulation + typename ElementAccumulator, + + /// Operator class tag + typename OperatorClass, + + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag, + + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + + /// Epilogue output operator + typename EpilogueOutputOp, + + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + + /// Number of stages used in the pipelined mainloop + int Stages + > + class Gemm; +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal acmcmulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassSimt, + /// Tag indicating architecture to tune for + typename ArchTag_ = arch::Sm70, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = threadblock::GemmBatchedIdentityThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// Operation performed by GEMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator +> +class GemmArray { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + using Operator = Operator_; + + /// Define the kernel + using DefaultGemmKernel = typename kernel::DefaultGemm< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + false, + Operator + >::GemmKernel; + + using GemmKernel = kernel::GemmArray; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + + ElementA const * const *ptr_A; + LayoutA layout_A; + + ElementB const * const *ptr_B; + LayoutB layout_B; + + ElementC const * const *ptr_C; + LayoutC layout_C; + + ElementC * const * ptr_D; + LayoutC layout_D; + + typename EpilogueOutputOp::Params epilogue; + int batch_count; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + ElementA const * const *ptr_A_, + LayoutA layout_A_, + ElementB const * const *ptr_B_, + LayoutB layout_B_, + ElementC const * const *ptr_C_, + LayoutC layout_C_, + ElementC * const * ptr_D_, + LayoutC layout_D_, + typename EpilogueOutputOp::Params epilogue_, + int batch_count_ + ): + problem_size(problem_size_), + ptr_A(ptr_A_), + layout_A(layout_A_), + ptr_B(ptr_B_), + layout_B(layout_B_), + ptr_C(ptr_C_), + layout_C(layout_C_), + ptr_D(ptr_D_), + layout_D(layout_D_), + epilogue(epilogue_), + batch_count(batch_count_) { } + }; + +private: + + /// Kernel parameters object + typename GemmKernel::Params params_; + +public: + + /// Constructs the GEMM. + GemmArray() { } + + /// Determines whether the GEMM can exemcte the given problem. + static Status can_implement(Arguments const &args) { + + if (args.layout_A.stride(0) % kAlignmentA) { + return Status::kErrorMisalignedOperand; + } + + if (args.layout_B.stride(0) % kAlignmentB) { + return Status::kErrorMisalignedOperand; + } + + if (args.layout_C.stride(0) % kAlignmentC) { + return Status::kErrorMisalignedOperand; + } + + if (args.layout_D.stride(0) % kAlignmentC) { + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + return 0; + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count); + + // Initialize the Params structure + params_ = typename GemmKernel::Params{ + args.problem_size, + grid_shape, + args.ptr_A, + args.layout_A, + args.ptr_B, + args.layout_B, + args.ptr_C, + args.layout_C, + args.ptr_D, + args.layout_D, + args.epilogue, + args.batch_count + }; + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + args.batch_count, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}); + + params_ = typename GemmKernel::Params{ + args.problem_size, + grid_shape, + args.ptr_A, + args.layout_A, + args.ptr_B, + args.layout_B, + args.ptr_C, + args.layout_C, + args.ptr_D, + args.layout_D, + args.epilogue, + args.batch_count + }; + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(GemmKernel::kThreadCount, 1, 1); + + mcError_t result; + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + if (smem_size >= (64 << 10)) { + result = mcFuncSetAttribute(Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + mctlass::Kernel<<>>(params_); + + result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output exchanges problem size and operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal acmcmulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB, + typename Operator_ +> +class GemmArray< + ElementA_, + LayoutA_, + ElementB_, + LayoutB_, + ElementC_, + layout::ColumnMajor, + ElementAccumulator_, + OperatorClass_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + AlignmentA, + AlignmentB, + Operator_ +> { +public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static int const kStages = Stages; + + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = false; + + // + using UnderlyingOperator = GemmArray< + ElementB, + typename layout::LayoutTranspose::type, + ElementA, + typename layout::LayoutTranspose::type, + ElementC, + layout::RowMajor, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + kAlignmentB, + kAlignmentA + >; + + using UnderlyingArguments = typename UnderlyingOperator::Arguments; + using GemmKernel = typename UnderlyingOperator::GemmKernel; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + + ElementA const * const *ptr_A; + LayoutA layout_A; + + ElementB const * const *ptr_B; + LayoutB layout_B; + + ElementC const * const *ptr_C; + LayoutC layout_C; + + ElementC * const * ptr_D; + LayoutC layout_D; + + typename EpilogueOutputOp::Params epilogue; + int batch_count; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + ElementA const * const *ptr_A_, + LayoutA layout_A_, + ElementB const * const *ptr_B_, + LayoutB layout_B_, + ElementC const * const *ptr_C_, + LayoutC layout_C_, + ElementC * const * ptr_D_, + LayoutC layout_D_, + typename EpilogueOutputOp::Params epilogue_, + int batch_count_ + ): + problem_size(problem_size_), + ptr_A(ptr_A_), + layout_A(layout_A_), + ptr_B(ptr_B_), + layout_B(layout_B_), + ptr_C(ptr_C_), + layout_C(layout_C_), + ptr_D(ptr_D_), + layout_D(layout_D_), + epilogue(epilogue_), + batch_count(batch_count_) { } + }; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the GEMM. + GemmArray() { } + + /// Helper to construct a transposed equivalent for the underying GEMM operator + static UnderlyingArguments to_underlying_arguments(Arguments const &args) { + + GemmCoord problem_size{ + args.problem_size.n(), + args.problem_size.m(), + args.problem_size.k() + }; + + return UnderlyingArguments( + problem_size, + args.ptr_B, + args.layout_B.stride(), + args.ptr_A, + args.layout_A.stride(), + args.ptr_C, + args.layout_C.stride(), + args.ptr_D, + args.layout_D.stride(), + args.epilogue, + args.batch_count + ); + } + + /// Determines whether the GEMM can exemcte the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } + +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_batched.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_batched.h new file mode 100644 index 0000000..5992bbb --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_batched.h @@ -0,0 +1,703 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the domcmentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined batch GEMM kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/gemm_batched.h" + +#include "mctlass/gemm/kernel/default_gemm.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +//////////////////////////////////////////////////////////////////////////////// + +/*! Gemm device-level operator. This is an interface to efficient MCTLASS GEMM kernels that may + be invoked from host code. + + The contributions of this class are: + + 1. At compile time, it maps data types and high-level structural parameters onto + specific MCTLASS components. + + 2. At runtime, it maps logical arguments to GEMM problems to kernel parameters. + + 3. At runtime, it launches kernels on the device. + + The intent is to provide a convenient mechanism for interacting with most plausible GEMM + configurations for each supported architecture. Consequently, not all parameters are exposed + to the top-level interface. Rather, sensible defaults at each level of the MCTLASS hierarchy + are selected to tradeoff simplicity of the interface with flexibility. We expect + most configurations to be specified at this level. Applications with more exotic requirements + may construct their kernels of interest using MCTLASS components at the threadblock, warp, + and thread levels of abstraction. + + MCTLASS exposes computations using the functor design pattern in which objects compose some + internal state with an overloaded function call operator. This enables decoupling of + initialization from exemction, possibly reducing overhead during steady state phases of + application exemction. + + MCTLASS device-level operators expose an Arguments structure encompassing each logical + input to the computation. This is distinct from the kernel-level Params structure pattern + which contains application-specific precomputed state needed by the device code. + + Example of a MCTLASS GEMM operator implementing the functionality of mcBLAS's SGEMM NN + is as follows: + + // + // Instantiate the MCTLASS GEMM operator. + // + + mctlass::gemm::device::Gemm< + float, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::ColumnMajor + > gemm_op; + + // + // Launch the GEMM operation on the device + // + + mctlass::Status status = gemm_op({ + {m, n, k}, // GemmCoord problem_size, + {A, lda}, // TensorRef ref_A, + {B, ldb}, // TensorRef ref_B, + {C, ldc}, // TensorRef ref_C, + {D, ldd}, // TensorRef ref_D, + {alpha, beta} // EpilogueOutputOp::Params epilogue_op_params + }); + + + A simplified view of the template is listed below. + + template < + /// Element type for A matrix operand + typename ElementA, + + /// Layout type for A matrix operand + typename LayoutA, + + /// Element type for B matrix operand + typename ElementB, + + /// Layout type for B matrix operand + typename LayoutB, + + /// Element type for C and D matrix operands + typename ElementC, + + /// Layout type for C and D matrix operands + typename LayoutC, + + /// Element type for internal acmcmulation + typename ElementAccumulator, + + /// Operator class tag + typename OperatorClass, + + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag, + + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + + /// Epilogue output operator + typename EpilogueOutputOp, + + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + + /// Number of stages used in the pipelined mainloop + int Stages + > + class Gemm; +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal acmcmulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassSimt, + /// Tag indicating architecture to tune for + typename ArchTag_ = arch::Sm70, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = threadblock::GemmBatchedIdentityThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// Operation performed by GEMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator +> +class GemmBatched { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + using Operator = Operator_; + + /// Define the kernel + using DefaultGemmKernel = typename kernel::DefaultGemm< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + false, + Operator + >::GemmKernel; + + using GemmKernel = kernel::GemmBatched; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + int64_t stride_A; + TensorRef ref_B; + int64_t stride_B; + TensorRef ref_C; + int64_t stride_C; + TensorRef ref_D; + int64_t stride_D; + typename EpilogueOutputOp::Params epilogue; + int batch_count; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + TensorRef ref_A_, + int64_t stride_A_, + TensorRef ref_B_, + int64_t stride_B_, + TensorRef ref_C_, + int64_t stride_C_, + TensorRef ref_D_, + int64_t stride_D_, + typename EpilogueOutputOp::Params epilogue_, + int batch_count_ + ): + problem_size(problem_size_), + ref_A(ref_A_), + stride_A(stride_A_), + ref_B(ref_B_), + stride_B(stride_B_), + ref_C(ref_C_), + stride_C(stride_C_), + ref_D(ref_D_), + stride_D(stride_D_), + epilogue(epilogue_), + batch_count(batch_count_) { } + }; + +private: + + /// Kernel parameters object + typename GemmKernel::Params params_; + +public: + + /// Constructs the GEMM. + GemmBatched() { } + + /// Determines whether the GEMM can exemcte the given problem. + static Status can_implement(Arguments const &args) { + + if (!TensorRef_aligned(args.ref_A, kAlignmentA) || (args.stride_A % kAlignmentA)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(args.ref_B, kAlignmentB) || (args.stride_B % kAlignmentB)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(args.ref_C, kAlignmentC) || (args.stride_C % kAlignmentC)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(args.ref_D, kAlignmentC) || (args.stride_D % kAlignmentC)) { + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + return 0; + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count); + + // Initialize the Params structure + params_ = typename GemmKernel::Params{ + args.problem_size, + grid_shape, + args.ref_A.non_const_ref(), + args.stride_A, + args.ref_B.non_const_ref(), + args.stride_B, + args.ref_C.non_const_ref(), + args.stride_C, + args.ref_D, + args.stride_D, + args.epilogue, + args.batch_count + }; + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + params_.ref_A.reset(args.ref_A.non_const_ref().data()); + params_.ref_B.reset(args.ref_B.non_const_ref().data()); + params_.ref_C.reset(args.ref_C.non_const_ref().data()); + params_.ref_D.reset(args.ref_D.data()); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(GemmKernel::kThreadCount, 1, 1); + + mcError_t result; + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + if (smem_size >= (64 << 10)) { + result = mcFuncSetAttribute(Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + mctlass::Kernel<<>>(params_); + + result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output exchanges problem size and operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal acmcmulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB, + typename Operator_ +> +class GemmBatched< + ElementA_, + LayoutA_, + ElementB_, + LayoutB_, + ElementC_, + layout::ColumnMajor, + ElementAccumulator_, + OperatorClass_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + AlignmentA, + AlignmentB, + Operator_ +> { +public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static int const kStages = Stages; + + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = false; + + // + using UnderlyingOperator = GemmBatched< + ElementB, + typename layout::LayoutTranspose::type, + ElementA, + typename layout::LayoutTranspose::type, + ElementC, + layout::RowMajor, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + kAlignmentB, + kAlignmentA + >; + + using UnderlyingArguments = typename UnderlyingOperator::Arguments; + using GemmKernel = typename UnderlyingOperator::GemmKernel; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + int64_t stride_A; + TensorRef ref_B; + int64_t stride_B; + TensorRef ref_C; + int64_t stride_C; + TensorRef ref_D; + int64_t stride_D; + typename EpilogueOutputOp::Params epilogue; + int batch_count; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + TensorRef ref_A_, + int64_t stride_A_, + TensorRef ref_B_, + int64_t stride_B_, + TensorRef ref_C_, + int64_t stride_C_, + TensorRef ref_D_, + int64_t stride_D_, + typename EpilogueOutputOp::Params epilogue_, + int batch_count_ + ): + problem_size(problem_size_), + ref_A(ref_A_), + stride_A(stride_A_), + ref_B(ref_B_), + stride_B(stride_B_), + ref_C(ref_C_), + stride_C(stride_C_), + ref_D(ref_D_), + stride_D(stride_D_), + epilogue(epilogue_), + batch_count(batch_count_) { } + }; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the GEMM. + GemmBatched() { } + + /// Helper to construct a transposed equivalent for the underying GEMM operator + static UnderlyingArguments to_underlying_arguments(Arguments const &args) { + return UnderlyingArguments( + {args.problem_size.n(), args.problem_size.m(), args.problem_size.k()}, + {args.ref_B.data(), args.ref_B.stride(0)}, + args.stride_B, + {args.ref_A.data(), args.ref_A.stride(0)}, + args.stride_A, + {args.ref_C.data(), args.ref_C.stride(0)}, + args.stride_C, + {args.ref_D.data(), args.ref_D.stride(0)}, + args.stride_D, + args.epilogue, + args.batch_count + ); + } + + /// Determines whether the GEMM can exemcte the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } + +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_complex.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_complex.h new file mode 100644 index 0000000..342c8c3 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_complex.h @@ -0,0 +1,717 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the domcmentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/gemm.h" + +#include "mctlass/gemm/kernel/default_gemm_complex.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/*! Gemm device-level operator. This is an interface to efficient MCTLASS GEMM + kernels that may be invoked from host code. + + The contributions of this class are: + + 1. At compile time, it maps data types and high-level structural parameters + onto specific MCTLASS components. + + 2. At runtime, it maps logical arguments to GEMM problems to kernel + parameters. + + 3. At runtime, it launches kernels on the device. + + The intent is to provide a convenient mechanism for interacting with most + plausible GEMM configurations for each supported architecture. Consequently, + not all parameters are exposed to the top-level interface. Rather, sensible + defaults at each level of the MCTLASS hierarchy are selected to tradeoff + simplicity of the interface with flexibility. We expect most configurations to + be specified at this level. Applications with more exotic requirements may + construct their kernels of interest using MCTLASS components at the + threadblock, warp, and thread levels of abstraction. + + MCTLASS exposes computations using the functor design pattern in which objects + compose some internal state with an overloaded function call operator. This + enables decoupling of initialization from exemction, possibly reducing + overhead during steady state phases of application exemction. + + MCTLASS device-level operators expose an Arguments structure encompassing each + logical input to the computation. This is distinct from the kernel-level + Params structure pattern which contains application-specific precomputed state + needed by the device code. + + Example of a MCTLASS GEMM operator implementing the functionality of mcBLAS's + SGEMM NN is as follows: + + // + // Instantiate the MCTLASS GEMM operator. + // + + mctlass::gemm::device::Gemm< + float, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::ColumnMajor + > gemm_op; + + // + // Launch the GEMM operation on the device + // + + mctlass::Status status = gemm_op({ + {m, n, k}, // GemmCoord problem_size, + {A, lda}, // TensorRef ref_A, + {B, ldb}, // TensorRef ref_B, + {C, ldc}, // TensorRef ref_C, + {D, ldd}, // TensorRef ref_D, + {alpha, beta} // EpilogueOutputOp::Params epilogue_op_params + }); + + + A simplified view of the template is listed below. + + template < + /// Element type for A matrix operand + typename ElementA, + + /// Layout type for A matrix operand + typename LayoutA, + + /// Element type for B matrix operand + typename ElementB, + + /// Layout type for B matrix operand + typename LayoutB, + + /// Element type for C and D matrix operands + typename ElementC, + + /// Layout type for C and D matrix operands + typename LayoutC, + + /// Element type for internal acmcmulation + typename ElementAccumulator, + + /// Operator class tag + typename OperatorClass, + + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag, + + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + + /// Epilogue output operator + typename EpilogueOutputOp, + + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + + /// Number of stages used in the pipelined mainloop + int Stages + > + class Gemm; +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal acmcmulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassSimt, + /// Tag indicating architecture to tune for. + typename ArchTag_ = arch::Sm70, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = + threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB = ComplexTransform::kNone, + /// Multiply-add operator + // (selects complex or gaussian complex) + typename Operator_ = arch::OpMultiplyAddComplex, + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial = false> +class GemmComplex { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static int const kStages = Stages; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + using Operator = Operator_; + static bool const kSplitKSerial = SplitKSerial; + static int const kAlignmentA = 1; + static int const kAlignmentB = 1; + static int const kAlignmentC = EpilogueOutputOp::kCount; + + /// Define the kernel + using GemmKernel = typename kernel::DefaultGemmComplex< + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + kTransformA, + kTransformB, + Operator, + kSplitKSerial + >::GemmKernel; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + TensorRef ref_B; + TensorRef ref_C; + TensorRef ref_D; + typename EpilogueOutputOp::Params epilogue; + int split_k_slices; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments(): problem_size(0, 0, 0), split_k_slices(1) { + + } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + TensorRef ref_A_, + TensorRef ref_B_, + TensorRef ref_C_, + TensorRef ref_D_, + typename EpilogueOutputOp::Params epilogue_ = + typename EpilogueOutputOp::Params(), + int split_k_slices = 1 + ): + problem_size(problem_size_), + ref_A(ref_A_), + ref_B(ref_B_), + ref_C(ref_C_), + ref_D(ref_D_), + epilogue(epilogue_), + split_k_slices(split_k_slices) { + + } + }; + +private: + + /// Kernel parameters object + typename GemmKernel::Params params_; + +public: + + /// Constructs the GEMM. + GemmComplex() { } + + /// Determines whether the GEMM can exemcte the given problem. + static Status can_implement(Arguments const &args) { + + if (!kSplitKSerial && args.split_k_slices > 1) { + return Status::kErrorInvalidProblem; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + if (kSplitKSerial && args.split_k_slices > 1) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + return sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n()); + } + + return 0; + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + if (kSplitKSerial) { + if (args.split_k_slices > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + size_t bytes = get_workspace_size(args); + + mcError_t result = mcMemsetAsync(workspace, 0, bytes, stream); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + } + else { + + if (args.split_k_slices > 1) { + return Status::kErrorInvalidProblem; + } + } + + // Initialize the Params structure + params_ = typename GemmKernel::Params{ + args.problem_size, + grid_shape, + args.ref_A.non_const_ref(), + args.ref_B.non_const_ref(), + args.ref_C.non_const_ref(), + args.ref_D, + args.epilogue, + static_cast(workspace) + }; + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + if (kSplitKSerial && args.split_k_slices > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + } + + params_.ref_A.reset(args.ref_A.non_const_ref().data()); + params_.ref_B.reset(args.ref_B.non_const_ref().data()); + params_.ref_C.reset(args.ref_C.non_const_ref().data()); + params_.ref_D.reset(args.ref_D.data()); + params_.semaphore = static_cast(workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(GemmKernel::kThreadCount, 1, 1); + + mcError_t result; + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + if (smem_size >= (64 << 10)) { + result = mcFuncSetAttribute(Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + mctlass::Kernel<<>>(params_); + + result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output exchanges problem size and operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal acmcmulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Multiply-add operator + // (selects complex or gaussian complex) + typename Operator_, + /// If true, kernel supports split-K as a serial reduction + bool SplitKSerial +> +class GemmComplex< + ElementA_, + LayoutA_, + ElementB_, + LayoutB_, + ElementC_, + layout::ColumnMajor, // partially specialized on LayoutC + ElementAccumulator_, + OperatorClass_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + TransformA, + TransformB, + Operator_, + SplitKSerial +> { +public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static int const kStages = Stages; + using Operator = Operator_; + static bool const kSplitKSerial = SplitKSerial; + + using UnderlyingOperator = GemmComplex< + ElementB, + typename layout::LayoutTranspose::type, + ElementA, + typename layout::LayoutTranspose::type, + ElementC, + layout::RowMajor, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + TransformB, + TransformA, + Operator, + SplitKSerial + >; + + static int const kAlignmentA = UnderlyingOperator::kAlignmentB; + static int const kAlignmentB = UnderlyingOperator::kAlignmentA; + static int const kAlignmentC = UnderlyingOperator::kAlignmentC; + static ComplexTransform const kTransformA = UnderlyingOperator::kTransformB; + static ComplexTransform const kTransformB = UnderlyingOperator::kTransformA; + + using UnderlyingArguments = typename UnderlyingOperator::Arguments; + using GemmKernel = typename UnderlyingOperator::GemmKernel; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + TensorRef ref_B; + TensorRef ref_C; + TensorRef ref_D; + typename EpilogueOutputOp::Params epilogue; + int split_k_slices; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + TensorRef ref_A_, + TensorRef ref_B_, + TensorRef ref_C_, + TensorRef ref_D_, + typename EpilogueOutputOp::Params epilogue_ = + typename EpilogueOutputOp::Params(), + int split_k_slices = 1 + ): + problem_size(problem_size_), + ref_A(ref_A_), + ref_B(ref_B_), + ref_C(ref_C_), + ref_D(ref_D_), + epilogue(epilogue_), + split_k_slices(split_k_slices) { } + }; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the GEMM. + GemmComplex() { } + + /// Helper to construct a transposed equivalent for the underying GEMM operator + static UnderlyingArguments to_underlying_arguments(Arguments const &args) { + return UnderlyingArguments( + {args.problem_size.n(), args.problem_size.m(), args.problem_size.k()}, + {args.ref_B.data(), args.ref_B.stride(0)}, + {args.ref_A.data(), args.ref_A.stride(0)}, + {args.ref_C.data(), args.ref_C.stride(0)}, + {args.ref_D.data(), args.ref_D.stride(0)}, + args.epilogue, + args.split_k_slices + ); + } + + /// Determines whether the GEMM can exemcte the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_grouped.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_grouped.h new file mode 100644 index 0000000..5669047 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_grouped.h @@ -0,0 +1,61 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Device-level grouped GEMM. +*/ + +#pragma once + +#include "mctlass/gemm/device/base_grouped.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// GEMM Grouped +template +class GemmGrouped : public BaseGrouped { +public: + using GemmKernel = GemmKernel_; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_layernorm_mainloop_fusion.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_layernorm_mainloop_fusion.h new file mode 100644 index 0000000..566f671 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_layernorm_mainloop_fusion.h @@ -0,0 +1,385 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Device-level GEMM with layernorm elementwise operations fused in mainloop +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/gemm_universal.h" + +#include "mctlass/gemm/kernel/default_gemm_layernorm_mainloop_fusion.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" +#include "mctlass/gemm/device/gemm_universal_base.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/*! + The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and + batched array variants. +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for Scale/Bias vectors + typename ElementScaleBias_, + /// Layout type for Scale/Bias vectors + typename LayoutScaleBias_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassSimt, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_ = arch::Sm70, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// Operation performed by GEMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator +> +class GemmLayernormMainloopFusion : + public GemmUniversalBase< + typename kernel::DefaultGemmLayernormMainloopFusion< + ElementA_, + LayoutA_, + AlignmentA, + ElementB_, + LayoutB_, + AlignmentB, + ElementScaleBias_, + LayoutScaleBias_, + ElementC_, + LayoutC_, + ElementAccumulator_, + OperatorClass_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + Operator_, + SharedMemoryClearOption::kNone + >::GemmKernel + > { + + public: + + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + + using Base = GemmUniversalBase< + typename kernel::DefaultGemmLayernormMainloopFusion< + ElementA_, + LayoutA_, + AlignmentA, + ElementB_, + LayoutB_, + AlignmentB, + ElementScaleBias_, + LayoutScaleBias_, + ElementC_, + LayoutC_, + ElementAccumulator_, + OperatorClass_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + Operator_, + SharedMemoryClearOption::kNone + >::GemmKernel + >; + + using Arguments = typename Base::Arguments; + using GemmKernel = typename Base::GemmKernel; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output exchanges problem size and operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for Scale/Bias vectors + typename ElementScaleBias_, + /// Layout type for Scale/Bias vectors + typename LayoutScaleBias_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB, + /// Operation performed by GEMM + typename Operator_ +> +class GemmLayernormMainloopFusion { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementScaleBias = ElementScaleBias_; + using LayoutScaleBias = LayoutScaleBias_; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + + using UnderlyingOperator = typename GemmLayernormMainloopFusion< + ElementB, + typename layout::LayoutTranspose::type, + ElementA, + typename layout::LayoutTranspose::type, + ElementScaleBias, + LayoutScaleBias, + ElementC, + layout::RowMajor, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + kAlignmentB, + kAlignmentA, + Operator + >::Base; + + using GemmKernel = typename UnderlyingOperator::GemmKernel; + static int const kAlignmentC = EpilogueOutputOp::kCount; + + /// Argument structure + using Arguments = typename UnderlyingOperator::Arguments; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the GEMM. + GemmLayernormMainloopFusion() { } + + /// Helper to construct a transposed equivalent for the underlying GEMM operator + static Arguments to_underlying_arguments(Arguments const &args) { + return args.transposed_problem(); + } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Computes the grid shape + static dim3 get_grid_shape(Arguments const &args) { + return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args)); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int smem_capacity = -1) { + return UnderlyingOperator::maximum_active_blocks(smem_capacity); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_sparse.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_sparse.h new file mode 100644 index 0000000..28bcb12 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_sparse.h @@ -0,0 +1,514 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the domcmentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/sparse_gemm.h" + +#include "mctlass/gemm/kernel/default_gemm_sparse.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/*! Gemm device-level operator. This is an interface to efficient MCTLASS GEMM kernels that may + be invoked from host code. + + The contributions of this class are: + + 1. At compile time, it maps data types and high-level structural parameters onto + specific MCTLASS components. + + 2. At runtime, it maps logical arguments to GEMM problems to kernel parameters. + + 3. At runtime, it launches kernels on the device. + + The intent is to provide a convenient mechanism for interacting with most plausible GEMM + configurations for each supported architecture. Consequently, not all parameters are exposed + to the top-level interface. Rather, sensible defaults at each level of the MCTLASS hierarchy + are selected to tradeoff simplicity of the interface with flexibility. We expect + most configurations to be specified at this level. Applications with more exotic requirements + may construct their kernels of interest using MCTLASS components at the threadblock, warp, + and thread levels of abstraction. + + MCTLASS exposes computations using the functor design pattern in which objects compose some + internal state with an overloaded function call operator. This enables decoupling of + initialization from exemction, possibly reducing overhead during steady state phases of + application exemction. + + MCTLASS device-level operators expose an Arguments structure encompassing each logical + input to the computation. This is distinct from the kernel-level Params structure pattern + which contains application-specific precomputed state needed by the device code. + + Example of a MCTLASS GEMM operator implementing the functionality of mcBLAS's SGEMM NN + is as follows: + + // + // Instantiate the MCTLASS GEMM operator. + // + + mctlass::gemm::device::Gemm< + float, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::ColumnMajor + > gemm_op; + + // + // Launch the GEMM operation on the device + // + + mctlass::Status status = gemm_op({ + {m, n, k}, // GemmCoord problem_size, + {A, lda}, // TensorRef ref_A, + {B, ldb}, // TensorRef ref_B, + {C, ldc}, // TensorRef ref_C, + {D, ldd}, // TensorRef ref_D, + {alpha, beta} // EpilogueOutputOp::Params epilogue_op_params + }); + + + A simplified view of the template is listed below. + + template < + /// Element type for A matrix operand + typename ElementA, + + /// Layout type for A matrix operand + typename LayoutA, + + /// Element type for B matrix operand + typename ElementB, + + /// Layout type for B matrix operand + typename LayoutB, + + /// Element type for C and D matrix operands + typename ElementC, + + /// Layout type for C and D matrix operands + typename LayoutC, + + /// Element type for internal acmcmulation + typename ElementAccumulator, + + /// Operator class tag + typename OperatorClass, + + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag, + + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + + /// Epilogue output operator + typename EpilogueOutputOp, + + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + + /// Number of stages used in the pipelined mainloop + int Stages + > + class Gemm; +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal acmcmulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassSimt, + /// Tag indicating architecture to tune for + typename ArchTag_ = arch::Sm70, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = + typename threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial = false, + /// Operation performed by GEMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator> +class SparseGemm { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + using MathOperator = Operator; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = SplitKSerial; + static ComplexTransform const kTransformA = ComplexTransform::kNone; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + + /// Define the kernel + using GemmKernel = typename kernel::DefaultSparseGemm< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + kSplitKSerial, + Operator + >::GemmKernel; + + using ElementE = typename GemmKernel::ElementE; + + using LayoutE = typename GemmKernel::LayoutE; + + static int const kAlignmentE = 128 / sizeof_bits::value; + + static int const kSparse = GemmKernel::kSparse; + static int const kMetaSizeInBits = GemmKernel::kMetaSizeInBits; + static int const kElementsPerElementE = GemmKernel::kElementsPerElementE; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + TensorRef ref_B; + TensorRef ref_C; + TensorRef ref_D; + TensorRef ref_E; + typename EpilogueOutputOp::Params epilogue; + int split_k_slices; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments(): problem_size(0, 0, 0), split_k_slices(1) { + + } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + TensorRef ref_A_, + TensorRef ref_B_, + TensorRef ref_C_, + TensorRef ref_D_, + TensorRef ref_E_, + typename EpilogueOutputOp::Params epilogue_ = + typename EpilogueOutputOp::Params(), + int split_k_slices = 1 + ): + problem_size(problem_size_), + ref_A(ref_A_), + ref_B(ref_B_), + ref_C(ref_C_), + ref_D(ref_D_), + ref_E(ref_E_), + epilogue(epilogue_), + split_k_slices(split_k_slices) { + + } + }; + +private: + + /// Kernel parameters object + typename GemmKernel::Params params_; + +public: + + /// Constructs the GEMM. + SparseGemm() { } + + /// Determines whether the GEMM can exemcte the given problem. + static Status can_implement(Arguments const &args) { + + if (!kSplitKSerial && args.split_k_slices > 1) { + return Status::kErrorInvalidProblem; + } + + Status status = GemmKernel::can_implement( + args.problem_size, + args.ref_A.non_const_ref(), + args.ref_B.non_const_ref(), + args.ref_C.non_const_ref(), + args.ref_D, + args.ref_E.non_const_ref() + ); + + if (status != Status::kSuccess) { + return status; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + size_t bytes = 0; + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + if (kSplitKSerial && args.split_k_slices > 1) { + + bytes += sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n()); + } + + return bytes; + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + if (kSplitKSerial) { + if (args.split_k_slices > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + size_t bytes = get_workspace_size(args); + + mcError_t result = mcMemsetAsync(workspace, 0, bytes, stream); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + } + else { + + if (args.split_k_slices > 1) { + return Status::kErrorInvalidProblem; + } + } + + // Initialize the Params structure + params_ = typename GemmKernel::Params{ + args.problem_size, + grid_shape, + args.ref_A.non_const_ref(), + args.ref_B.non_const_ref(), + args.ref_C.non_const_ref(), + args.ref_D, + args.ref_E.non_const_ref(), + args.epilogue, + static_cast(workspace) + }; + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + if (smem_size >= (64 << 10)) { + mcError_t result = mcFuncSetAttribute(Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + if (kSplitKSerial && args.split_k_slices > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + } + + params_.ref_A.reset(args.ref_A.non_const_ref().data()); + params_.ref_B.reset(args.ref_B.non_const_ref().data()); + params_.ref_C.reset(args.ref_C.non_const_ref().data()); + params_.ref_D.reset(args.ref_D.data()); + params_.ref_E.reset(args.ref_E.non_const_ref().data()); + params_.output_op = args.epilogue; + params_.semaphore = static_cast(workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(GemmKernel::kThreadCount, 1, 1); + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + + mctlass::Kernel<<>>(params_); + + mcError_t result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_sparse_row_broadcast.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_sparse_row_broadcast.h new file mode 100644 index 0000000..29bb52a --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_sparse_row_broadcast.h @@ -0,0 +1,514 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/sparse_gemm_row_broadcast.h" + +#include "mctlass/gemm/kernel/default_gemm_sparse_row_broadcast.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/*! Gemm device-level operator. This is an interface to efficient MCTLASS GEMM kernels that may + be invoked from host code. + + The contributions of this class are: + + 1. At compile time, it maps data types and high-level structural parameters onto + specific MCTLASS components. + + 2. At runtime, it maps logical arguments to GEMM problems to kernel parameters. + + 3. At runtime, it launches kernels on the device. + + The intent is to provide a convenient mechanism for interacting with most plausible GEMM + configurations for each supported architecture. Consequently, not all parameters are exposed + to the top-level interface. Rather, sensible defaults at each level of the MCTLASS hierarchy + are selected to tradeoff simplicity of the interface with flexibility. We expect + most configurations to be specified at this level. Applications with more exotic requirements + may construct their kernels of interest using MCTLASS components at the threadblock, warp, + and thread levels of abstraction. + + MCTLASS exposes computations using the functor design pattern in which objects compose some + internal state with an overloaded function call operator. This enables decoupling of + initialization from execution, possibly reducing overhead during steady state phases of + application execution. + + MCTLASS device-level operators expose an Arguments structure encompassing each logical + input to the computation. This is distinct from the kernel-level Params structure pattern + which contains application-specific precomputed state needed by the device code. + + Example of a MCTLASS GEMM operator implementing the functionality of cuBLAS's SGEMM NN + is as follows: + + // + // Instantiate the MCTLASS GEMM operator. + // + + mctlass::gemm::device::Gemm< + float, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::ColumnMajor + > gemm_op; + + // + // Launch the GEMM operation on the device + // + + mctlass::Status status = gemm_op({ + {m, n, k}, // GemmCoord problem_size, + {A, lda}, // TensorRef ref_A, + {B, ldb}, // TensorRef ref_B, + {C, ldc}, // TensorRef ref_C, + {D, ldd}, // TensorRef ref_D, + {alpha, beta} // EpilogueOutputOp::Params epilogue_op_params + }); + + + A simplified view of the template is listed below. + + template < + /// Element type for A matrix operand + typename ElementA, + + /// Layout type for A matrix operand + typename LayoutA, + + /// Element type for B matrix operand + typename ElementB, + + /// Layout type for B matrix operand + typename LayoutB, + + /// Element type for C and D matrix operands + typename ElementC, + + /// Layout type for C and D matrix operands + typename LayoutC, + + /// Element type for internal accumulation + typename ElementAccumulator, + + /// Operator class tag + typename OperatorClass, + + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag, + + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + + /// Epilogue output operator + typename EpilogueOutputOp, + + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + + /// Number of stages used in the pipelined mainloop + int Stages + > + class Gemm; +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassSimt, + /// Tag indicating architecture to tune for + typename ArchTag_ = arch::Sm70, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = + typename threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial = false, + /// Operation performed by GEMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator> +class SparseGemmRowBroadcast { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + using MathOperator = Operator; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = SplitKSerial; + static ComplexTransform const kTransformA = ComplexTransform::kNone; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + + /// Define the kernel + using GemmKernel = typename kernel::DefaultSparseGemmRowBroadcast< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + kSplitKSerial, + Operator + >::GemmKernel; + + using ElementE = typename GemmKernel::ElementE; + + using LayoutE = typename GemmKernel::LayoutE; + + static int const kAlignmentE = 128 / sizeof_bits::value; + + static int const kSparse = GemmKernel::kSparse; + static int const kMetaSizeInBits = GemmKernel::kMetaSizeInBits; + static int const kElementsPerElementE = GemmKernel::kElementsPerElementE; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + TensorRef ref_B; + TensorRef ref_C; + TensorRef ref_D; + TensorRef ref_E; + typename EpilogueOutputOp::Params epilogue; + int split_k_slices; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments(): problem_size(0, 0, 0), split_k_slices(1) { + + } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + TensorRef ref_A_, + TensorRef ref_B_, + TensorRef ref_C_, + TensorRef ref_D_, + TensorRef ref_E_, + typename EpilogueOutputOp::Params epilogue_ = + typename EpilogueOutputOp::Params(), + int split_k_slices = 1 + ): + problem_size(problem_size_), + ref_A(ref_A_), + ref_B(ref_B_), + ref_C(ref_C_), + ref_D(ref_D_), + ref_E(ref_E_), + epilogue(epilogue_), + split_k_slices(split_k_slices) { + + } + }; + +private: + + /// Kernel parameters object + typename GemmKernel::Params params_; + +public: + + /// Constructs the GEMM. + SparseGemmRowBroadcast() { } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + if (!kSplitKSerial && args.split_k_slices > 1) { + return Status::kErrorInvalidProblem; + } + + Status status = GemmKernel::can_implement( + args.problem_size, + args.ref_A.non_const_ref(), + args.ref_B.non_const_ref(), + args.ref_C.non_const_ref(), + args.ref_D, + args.ref_E.non_const_ref() + ); + + if (status != Status::kSuccess) { + return status; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + size_t bytes = 0; + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + if (kSplitKSerial && args.split_k_slices > 1) { + + bytes += sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n()); + } + + return bytes; + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + if (kSplitKSerial) { + if (args.split_k_slices > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + size_t bytes = get_workspace_size(args); + + mcError_t result = mcMemsetAsync(workspace, 0, bytes, stream); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + } + else { + + if (args.split_k_slices > 1) { + return Status::kErrorInvalidProblem; + } + } + + // Initialize the Params structure + params_ = typename GemmKernel::Params{ + args.problem_size, + grid_shape, + args.ref_A.non_const_ref(), + args.ref_B.non_const_ref(), + args.ref_C.non_const_ref(), + args.ref_D, + args.ref_E.non_const_ref(), + args.epilogue, + static_cast(workspace) + }; + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + if (smem_size >= (48 << 10)) { + mcError_t result = mcFuncSetAttribute(Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + if (kSplitKSerial && args.split_k_slices > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + } + + params_.ref_A.reset(args.ref_A.non_const_ref().data()); + params_.ref_B.reset(args.ref_B.non_const_ref().data()); + params_.ref_C.reset(args.ref_C.non_const_ref().data()); + params_.ref_D.reset(args.ref_D.data()); + params_.ref_E.reset(args.ref_E.non_const_ref().data()); + params_.output_op = args.epilogue; + params_.semaphore = static_cast(workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(GemmKernel::kThreadCount, 1, 1); + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + + mctlass::Kernel<<>>(params_); + + mcError_t result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_splitk_parallel.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_splitk_parallel.h new file mode 100644 index 0000000..eb8a91e --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_splitk_parallel.h @@ -0,0 +1,637 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for GEMM performing a reduction over K partitions in parallel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/gemm.h" + +#include "mctlass/gemm/kernel/default_gemm_splitk_parallel.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +#include "mctlass/epilogue/thread/conversion_op.h" +#include "mctlass/reduction/kernel/reduce_split_k.h" +#include "mctlass/reduction/thread/reduction_operators.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +//////////////////////////////////////////////////////////////////////////////// + +/*! + Gemm device-level operator performing parallel reduction over the K partition. + +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassSimt, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_ = arch::Sm70, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Epilogue output operator + typename ConvertScaledOp_ = mctlass::epilogue::thread::Convert< + ElementAccumulator_, + DefaultGemmConfiguration::EpilogueOutputOp::kCount, + ElementAccumulator_>, + /// Reduction operator + typename ReductionOp_ = mctlass::reduction::thread::ReduceAdd< + ElementAccumulator_, typename EpilogueOutputOp_::ElementAccumulator, + EpilogueOutputOp_::kCount>, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = + threadblock::GemmSplitKHorizontalThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int kAlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int kAlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// Operation performed by GEMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator> +class GemmSplitKParallel { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ConvertScaledOp = ConvertScaledOp_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ReductionOp = ReductionOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static int const kStages = Stages; + + /// GEMM kernel + using GemmKernel = typename kernel::DefaultGemmSplitKParallel< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + ConvertScaledOp, + ThreadblockSwizzle, + kStages, + Operator + >::GemmKernel; + + /// Reduction kernel + using ReductionKernel = mctlass::reduction::kernel::ReduceSplitK< + mctlass::MatrixShape<4, 32 * EpilogueOutputOp::kCount>, + EpilogueOutputOp, + ReductionOp + >; + + // + // + // + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + TensorRef ref_B; + TensorRef ref_C; + TensorRef ref_D; + typename EpilogueOutputOp::Params epilogue; + int split_k_slices; + typename ConvertScaledOp::Params convert; + typename ReductionOp::Params reduction; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + TensorRef ref_A_, + TensorRef ref_B_, + TensorRef ref_C_, + TensorRef ref_D_, + typename EpilogueOutputOp::Params epilogue_ = + typename EpilogueOutputOp::Params(), + int split_k_slices = 1, + typename ConvertScaledOp::Params convert_ = + typename ConvertScaledOp::Params(), + typename ReductionOp::Params reduction_ = + typename ReductionOp::Params() + ): + problem_size(problem_size_), + ref_A(ref_A_), + ref_B(ref_B_), + ref_C(ref_C_), + ref_D(ref_D_), + epilogue(epilogue_), + split_k_slices(split_k_slices), + convert(convert_), + reduction(reduction_) { } + }; + +private: + + /// Kernel parameters object + typename GemmKernel::Params gemm_params_; + + /// Reduction kernel parameters object + typename ReductionKernel::Params reduction_params_; + +public: + + /// Constructs the GEMM. + GemmSplitKParallel() { } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + return sizeof(ElementAccumulator_) * size_t(args.problem_size.m()) * size_t(args.problem_size.n()) * grid_shape.k(); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.split_k_slices); + + // Define a reference to the workspace - this is an aligned region in device memory. + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + TensorRef ref_workspace( + static_cast(workspace), + args.problem_size.n()); + + int64_t partition_stride = int64_t(args.problem_size.m()) * int64_t(args.problem_size.n()); + + // Initialize the Params structure + gemm_params_ = typename GemmKernel::Params{ + args.problem_size, + grid_shape, + args.ref_A.non_const_ref(), + args.ref_B.non_const_ref(), + ref_workspace, + args.convert, + partition_stride + }; + + reduction_params_ = typename ReductionKernel::Params( + args.problem_size.mn(), + grid_shape.k(), + partition_stride, + ref_workspace, + args.ref_D, + args.ref_C.non_const_ref(), + args.epilogue + ); + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + gemm_params_.ref_A.reset(args.ref_A.data()); + gemm_params_.ref_B.reset(args.ref_B.data()); + gemm_params_.ref_D.reset(workspace); + + reduction_params_.ref_D.reset(args.ref_D.data()); + reduction_params_.ref_C.reset(args.ref_C.data()); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + // + // Launch GEMM kernel + // + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(gemm_params_.grid_tiled_shape); + dim3 block(GemmKernel::kThreadCount, 1, 1); + + mcError_t result; + + int smem_size = int(sizeof(typename GemmKernel::SharedStorage)); + if (smem_size >= (64 << 10)) { + + result = mcFuncSetAttribute( + Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + Kernel<<>>(gemm_params_); + + result = mcGetLastError(); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + + // + // Launch reduction kernel + // + + block = ReductionKernel::block_shape(); + grid = ReductionKernel::grid_shape(gemm_params_.problem_size.mn()); + + Kernel<<< grid, block, 0, stream >>>(reduction_params_); + + result = mcGetLastError(); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Epilogue output operator + typename ConvertScaledOp_, + /// Reduction operator + typename ReductionOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, int kAlignmentA, int kAlignmentB, + /// Operation performed by GEMM + typename Operator_> +class GemmSplitKParallel { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ConvertScaledOp = ConvertScaledOp_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ReductionOp = ReductionOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static int const kStages = Stages; + + using UnderlyingOperator = GemmSplitKParallel< + ElementB, + typename layout::LayoutTranspose::type, + ElementA, + typename layout::LayoutTranspose::type, + ElementC, + layout::RowMajor, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ConvertScaledOp, + ReductionOp, + ThreadblockSwizzle, + Stages, + kAlignmentA, + kAlignmentB, + Operator + >; + + using UnderlyingArguments = typename UnderlyingOperator::Arguments; + using GemmKernel = typename UnderlyingOperator::GemmKernel; + using ReductionKernel = typename UnderlyingOperator::ReductionKernel; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord problem_size; + TensorRef ref_A; + TensorRef ref_B; + TensorRef ref_C; + TensorRef ref_D; + typename EpilogueOutputOp::Params epilogue; + int split_k_slices; + typename ConvertScaledOp::Params convert; + typename ReductionOp::Params reduction; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() { } + + /// Constructs an Arguments structure + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord problem_size_, + TensorRef ref_A_, + TensorRef ref_B_, + TensorRef ref_C_, + TensorRef ref_D_, + typename EpilogueOutputOp::Params epilogue_ = + typename EpilogueOutputOp::Params(), + int split_k_slices = 1, + typename ConvertScaledOp::Params convert_ = + typename ConvertScaledOp::Params(), + typename ReductionOp::Params reduction_ = + typename ReductionOp::Params() + ): + problem_size(problem_size_), + ref_A(ref_A_), + ref_B(ref_B_), + ref_C(ref_C_), + ref_D(ref_D_), + epilogue(epilogue_), + split_k_slices(split_k_slices), + convert(convert_), + reduction(reduction_) { } + }; + +private: + + /// Kernel parameters object + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the GEMM. + GemmSplitKParallel() { } + + /// Helper to construct a transposed equivalent for the underying GEMM operator + static UnderlyingArguments to_underlying_arguments(Arguments const &args) { + return UnderlyingArguments( + {args.problem_size.n(), args.problem_size.m(), args.problem_size.k()}, + {args.ref_B.data(), args.ref_B.stride(0)}, + {args.ref_A.data(), args.ref_A.stride(0)}, + {args.ref_C.data(), args.ref_C.stride(0)}, + {args.ref_D.data(), args.ref_D.stride(0)}, + args.epilogue, + args.split_k_slices, + args.convert, + args.reduction + ); + } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_universal.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_universal.h new file mode 100644 index 0000000..c65cb8a --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_universal.h @@ -0,0 +1,441 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/gemm_universal.h" + +#include "mctlass/gemm/kernel/default_gemm_universal.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" +#include "mctlass/gemm/device/gemm_universal_base.h" + +#include "mctlass/layout/permute.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/*! + GemmUniversal is a stateful, reusable GEMM handle. Once initialized for a given GEMM computation + (problem geometry and data references), it can be reused across different GEMM problems having the + geometry. (Once initialized, details regarding problem geometry and references to workspace memory + cannot be updated.) + + The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and + batched array variants. +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassSimt, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_ = arch::Sm70, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// Operation performed by GEMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB = ComplexTransform::kNone, + /// Gather operand A by using an index array + bool GatherA = false, + /// Gather operand B by using an index array + bool GatherB = false, + /// Scatter result D by using an index array + bool ScatterD = false, + /// Permute result D + typename PermuteDLayout_ = layout::NoPermute, + /// Permute operand A + typename PermuteALayout_ = layout::NoPermute, + /// Permute operand B + typename PermuteBLayout_ = layout::NoPermute +> +class GemmUniversal : + public GemmUniversalBase< + typename kernel::DefaultGemmUniversal< + ElementA_, + LayoutA_, + TransformA, + AlignmentA, + ElementB_, + LayoutB_, + TransformB, + AlignmentB, + ElementC_, + LayoutC_, + ElementAccumulator_, + OperatorClass_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + Operator_, + SharedMemoryClearOption::kNone, + GatherA, + GatherB, + ScatterD, + PermuteDLayout_, + PermuteALayout_, + PermuteBLayout_ + >::GemmKernel + > { + + public: + + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + using PermuteDLayout = PermuteDLayout_; + using PermuteALayout = PermuteALayout_; + using PermuteBLayout = PermuteBLayout_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + + using Base = GemmUniversalBase< + typename kernel::DefaultGemmUniversal< + ElementA_, + LayoutA_, + TransformA, + AlignmentA, + ElementB_, + LayoutB_, + TransformB, + AlignmentB, + ElementC_, + LayoutC_, + ElementAccumulator_, + OperatorClass_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + Operator_, + SharedMemoryClearOption::kNone, + GatherA, + GatherB, + ScatterD, + PermuteDLayout_, + PermuteALayout_, + PermuteBLayout_ + >::GemmKernel + >; + + using Arguments = typename Base::Arguments; + using GemmKernel = typename Base::GemmKernel; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output exchanges problem size and operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB, + /// Operation performed by GEMM + typename Operator_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Scatter result D by using an index array + bool ScatterD, + /// Permute result D + typename PermuteDLayout_, + /// Permute operand A + typename PermuteALayout_, + /// Permute operand B + typename PermuteBLayout_ +> +class GemmUniversal { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + using PermuteDLayout = PermuteDLayout_; + using PermuteALayout = PermuteALayout_; + using PermuteBLayout = PermuteBLayout_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + + using UnderlyingOperator = typename GemmUniversal< + ElementB, + typename layout::LayoutTranspose::type, + ElementA, + typename layout::LayoutTranspose::type, + ElementC, + layout::RowMajor, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + kAlignmentB, + kAlignmentA, + Operator, + kTransformB, + kTransformA, + GatherB, + GatherA, + ScatterD, + PermuteDLayout, + PermuteBLayout, + PermuteALayout + >::Base; + + using GemmKernel = typename UnderlyingOperator::GemmKernel; + static int const kAlignmentC = EpilogueOutputOp::kCount; + + /// Argument structure + using Arguments = typename UnderlyingOperator::Arguments; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the GEMM. + GemmUniversal() { } + + /// Helper to construct a transposed equivalent for the underying GEMM operator + static Arguments to_underlying_arguments(Arguments const &args) { + return args.transposed_problem(); + } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Computes the grid shape + static dim3 get_grid_shape(Arguments const &args) { + return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args)); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int smem_capacity = -1) { + return UnderlyingOperator::maximum_active_blocks(smem_capacity); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_universal_adapter.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_universal_adapter.h new file mode 100644 index 0000000..bebcba0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_universal_adapter.h @@ -0,0 +1,554 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and + batched array variants. +*/ + +#pragma once + +// common +#include "mctlass/mctlass.h" +#include "mctlass/trace.h" +#include "mctlass/cluster_launch.hpp" +#include "mctlass/device_kernel.h" +#include "mctlass/gemm/gemm.h" + +// 2.x +#include "mctlass/gemm/device/gemm_universal_base.h" +#include "mctlass/gemm/kernel/gemm_transpose_operands.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +// 3.x +#include "mctlass/gemm/kernel/gemm_universal.hpp" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::device { + +//////////////////////////////////////////////////////////////////////////////// + +/*! + GemmUniversalAdapter is a stateful, reusable GEMM handle built around a kernel + of type mctlass::gemm::kernel::Gemm or mctlass::gemm::kernel::GemmUniversal. + + It manages the lifetime of the underlying `kernel::Params` struct, and exposes APIs + to create it from the host facing arguments. For power users, new static methods + are exposed in 3.x APIs that bypass the stateful methods or args->params lowering. + + It supports kernel types that implement both the 2.x and 3.0 APIs, + however, this is done by specializing the implementation of GemmUniversalAdapter + on the two kernel API types, and thus, GemmUniversalAdapter's behaviour might + differ between the two specializations. +*/ +template +class GemmUniversalAdapter; + +//////////////////////////////////////////////////////////////////////////////// +////////////////////////////// MCTLASS 3.x API ///////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template +class GemmUniversalAdapter< + GemmKernel_, + cute::enable_if_t::value>> +{ +public: + using GemmKernel = GemmKernel_; + using TileShape = typename GemmKernel::TileShape; + using ElementA = typename GemmKernel::ElementA; + using ElementB = typename GemmKernel::ElementB; + using ElementC = typename GemmKernel::ElementC; + using ElementD = typename GemmKernel::ElementD; + using ElementAccumulator = typename GemmKernel::TiledMma::ValTypeC; + using DispatchPolicy = typename GemmKernel::DispatchPolicy; + using CollectiveMainloop = typename GemmKernel::CollectiveMainloop; + using CollectiveEpilogue = typename GemmKernel::CollectiveEpilogue; + + // Map back to 2.x type as best as possible + using LayoutA = gemm::detail::StrideToLayoutTagA_t; + using LayoutB = gemm::detail::StrideToLayoutTagB_t; + using LayoutC = gemm::detail::StrideToLayoutTagC_t; + using LayoutD = gemm::detail::StrideToLayoutTagC_t; + + // NOTE: 3.0 kernels do not support complex transforms for now ... + static ComplexTransform const kTransformA = ComplexTransform::kNone; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + + // Legacy: Assume MultiplyAdd only since we do not use this tag type in 3.0 + using MathOperator = mctlass::arch::OpMultiplyAdd; + + // If our TiledMMA's instruction thread layout size is larger than 1, we know its a tensorop! + using OperatorClass = cute::conditional_t< + (cute::size(typename GemmKernel::TiledMma::AtomThrID{}) > 1), + mctlass::arch::OpClassTensorOp, mctlass::arch::OpClassSimt>; + + using ArchTag = typename GemmKernel::ArchTag; + + // NOTE: Assume identity swizzle for now + static_assert(cute::is_void_v, + "MCTLASS 3.x kernel types do not support grid swizzle functors yet."); + using ThreadblockSwizzle = mctlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>; + + // Assume TiledMma's ShapeMNK is the same as 2.x's ThreadblockShape + using ThreadblockShape = mctlass::gemm::GemmShape< + cute::size<0>(TileShape{}), + cute::size<1>(TileShape{}), + cute::size<2>(TileShape{})>; + + using ClusterShape = mctlass::gemm::GemmShape< + cute::size<0>(typename GemmKernel::DispatchPolicy::ClusterShape{}), + cute::size<1>(typename GemmKernel::DispatchPolicy::ClusterShape{}), + cute::size<2>(typename GemmKernel::DispatchPolicy::ClusterShape{})>; + + // Instruction shape is easy too, since we get that directly from our TiledMma's atom shape + using InstructionShape = mctlass::gemm::GemmShape< + cute::size<0>(typename CollectiveMainloop::TiledMma::AtomShape_MNK{}), + cute::size<1>(typename CollectiveMainloop::TiledMma::AtomShape_MNK{}), + cute::size<2>(typename CollectiveMainloop::TiledMma::AtomShape_MNK{})>; + + // Legacy: provide a correct warp count, but no reliable warp shape + static int const kThreadCount = GemmKernel::MaxThreadsPerBlock; + + // Warp shape is not a primary API type in 3.x + // But we can best approximate it by inspecting the TiledMma::TiledShape_MNK + // For this, we make the assumption that we always have 4 warps along M, and rest along N, none along K + // We also always round up the warp count to 4 if the tiled mma is smaller than 128 threads + static constexpr int WarpsInMma = cute::max(4, cute::size(typename GemmKernel::TiledMma{}) / 32); + static constexpr int WarpsInMmaM = 4; + static constexpr int WarpsInMmaN = cute::ceil_div(WarpsInMma, WarpsInMmaM); + using WarpCount = mctlass::gemm::GemmShape; + using WarpShape = mctlass::gemm::GemmShape< + cute::size<0>(typename CollectiveMainloop::TiledMma::TiledShape_MNK{}) / WarpsInMmaM, + cute::size<1>(typename CollectiveMainloop::TiledMma::TiledShape_MNK{}) / WarpsInMmaN, + cute::size<2>(typename CollectiveMainloop::TiledMma::TiledShape_MNK{})>; + + static int constexpr kStages = CollectiveMainloop::DispatchPolicy::Stages; + + // Inspect TiledCopy for A and B to compute the alignment size + static int constexpr kAlignmentA = gemm::detail::get_alignment_count_from_gmem_tiled_copy< + typename CollectiveMainloop::GmemTiledCopyA, ElementA>(); + static int constexpr kAlignmentB = gemm::detail::get_alignment_count_from_gmem_tiled_copy< + typename CollectiveMainloop::GmemTiledCopyB, ElementB>(); + static int constexpr kAlignmentC = gemm::detail::get_alignment_count_from_gmem_tiled_copy< + typename CollectiveEpilogue::GmemTiledCopyC, ElementC>(); + static int constexpr kAlignmentD = gemm::detail::get_alignment_count_from_gmem_tiled_copy< + typename CollectiveEpilogue::GmemTiledCopyD, ElementD>(); + + using EpilogueOutputOp = typename CollectiveEpilogue::ThreadEpilogueOp; + + // Split-K preserves splits that are 128b aligned + static int constexpr kSplitKAlignment = cute::max( + 128 / sizeof_bits::value, 128 / sizeof_bits::value); + + /// Argument structure: User API + using Arguments = typename GemmKernel::Arguments; + /// Argument structure: Kernel API + using Params = typename GemmKernel::Params; + +private: + + /// Kernel API parameters object + Params params_; + +public: + + /// Determines whether the GEMM can execute the given problem. + static Status + can_implement(Arguments const& args) { + if (GemmKernel::can_implement(args)) { + return Status::kSuccess; + } + else { + return Status::kInvalid; + } + } + + /// Gets the workspace size + static size_t + get_workspace_size(Arguments const& args) { + size_t workspace_bytes = 0; + if (args.mode == GemmUniversalMode::kGemmSplitKParallel) { + workspace_bytes += sizeof(int) * size_t(cute::size<0>(TileShape{})) * size_t(cute::size<1>(TileShape{})); + } + + MCTLASS_TRACE_HOST(" workspace_bytes: " << workspace_bytes); + + workspace_bytes += GemmKernel::get_workspace_size(args); + return workspace_bytes; + } + + /// Computes the grid shape + static dim3 + get_grid_shape(Arguments const& args, void* workspace = nullptr) { + auto tmp_params = GemmKernel::to_underlying_arguments(args, workspace); + return GemmKernel::get_grid_shape(tmp_params); + } + + /// Computes the grid shape + static dim3 + get_grid_shape(Params const& params) { + return GemmKernel::get_grid_shape(params); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int /* smem_capacity */ = -1) { + MCTLASS_TRACE_HOST("GemmUniversal::maximum_active_blocks()"); + int max_active_blocks = -1; + int smem_size = GemmKernel::SharedStorageSize; + + // first, account for dynamic smem capacity if needed + mcError_t result; + if (smem_size >= (48 << 10)) { + MCTLASS_TRACE_HOST(" Setting smem size to " << smem_size); + result = mcFuncSetAttribute( + device_kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + if (mcSuccess != result) { + result = mcGetLastError(); // to clear the error bit + MCTLASS_TRACE_HOST( + " mcFuncSetAttribute() returned error: " + << mcGetErrorString(result)); + return -1; + } + } + + // query occupancy after setting smem size + result = mcOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks, + device_kernel, + GemmKernel::MaxThreadsPerBlock, + smem_size); + + if (mcSuccess != result) { + result = mcGetLastError(); // to clear the error bit + MCTLASS_TRACE_HOST( + " mcOccupancyMaxActiveBlocksPerMultiprocessor() returned error: " + << mcGetErrorString(result)); + return -1; + } + + MCTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks); + return max_active_blocks; + } + + /// Initializes GEMM state from arguments. + Status + initialize(Arguments const& args, void* workspace = nullptr, mcStream_t stream = nullptr) { + MCTLASS_TRACE_HOST("GemmUniversal::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + size_t workspace_bytes = GemmKernel::get_workspace_size(args); + MCTLASS_TRACE_HOST(" workspace_bytes: " << workspace_bytes); + + if (workspace_bytes) { + if (!workspace) { + MCTLASS_TRACE_HOST(" error: device workspace must not be null"); + return Status::kErrorWorkspaceNull; + } + + if (args.mode == GemmUniversalMode::kGemm) { + MCTLASS_TRACE_HOST(" clearing device workspace"); + mcError_t result = mcMemsetAsync(workspace, 0, workspace_bytes, stream); + if (mcSuccess != result) { + result = mcGetLastError(); // to clear the error bit + MCTLASS_TRACE_HOST(" mcMemsetAsync() returned error " << mcGetErrorString(result)); + return Status::kErrorInternal; + } + } + } + + // Initialize the Params structure + params_ = GemmKernel::to_underlying_arguments(args, workspace); + + // account for dynamic smem capacity if needed + int smem_size = GemmKernel::SharedStorageSize; + if (smem_size >= (48 << 10)) { + MCTLASS_TRACE_HOST(" Setting smem size to " << smem_size); + mcError_t result = mcFuncSetAttribute( + device_kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + if (mcSuccess != result) { + result = mcGetLastError(); // to clear the error bit + MCTLASS_TRACE_HOST(" mcFuncSetAttribute() returned error: " << mcGetErrorString(result)); + return Status::kErrorInternal; + } + } + return Status::kSuccess; + } + + /// Update API is preserved in 3.0, but does not guarantee a lightweight update of params. + Status + update(Arguments const& args, void* workspace = nullptr) { + MCTLASS_TRACE_HOST("GemmUniversal()::update() - workspace: " << workspace); + + size_t workspace_bytes = get_workspace_size(args); + if (workspace_bytes > 0 && nullptr == workspace) { + return Status::kErrorWorkspaceNull; + } + + params_ = GemmKernel::to_underlying_arguments(args, workspace); + return Status::kSuccess; + } + + /// Primary run() entry point API that is static allowing users to create and manage their own params. + /// Supplied params struct must be construct by calling GemmKernel::to_underling_arguments() + static Status + run(Params& params, mcStream_t stream = nullptr) { + MCTLASS_TRACE_HOST("GemmUniversal::run()"); + dim3 const block = GemmKernel::get_block_shape(); + dim3 const grid = get_grid_shape(params); + + // configure smem size and carveout + int smem_size = GemmKernel::SharedStorageSize; + + Status launch_result; + // Use extended launch API only for mainloops that use it + if constexpr(GemmKernel::ArchTag::kMinComputeCapability >= 90) { + dim3 cluster(cute::size<0>(typename GemmKernel::DispatchPolicy::ClusterShape{}), + cute::size<1>(typename GemmKernel::DispatchPolicy::ClusterShape{}), + cute::size<2>(typename GemmKernel::DispatchPolicy::ClusterShape{})); + void const* kernel = (void const*) device_kernel; + void* kernel_params[] = {¶ms}; + launch_result = ClusterLauncher::launch(grid, cluster, block, smem_size, stream, kernel, kernel_params); + } + else { + launch_result = Status::kSuccess; + device_kernel<<>>(params); + } + + mcError_t result = mcGetLastError(); + if (mcSuccess == result && Status::kSuccess == launch_result) { + return Status::kSuccess; + } + else { + MCTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result); + return Status::kErrorInternal; + } + } + + // + // Non-static launch overloads that first create and set the internal params struct of this kernel handle. + // + + /// Launches the kernel after first constructing Params internal state from supplied arguments. + Status + run(Arguments const& args, void* workspace = nullptr, mcStream_t stream = nullptr) { + Status status = initialize(args, workspace, stream); + if (Status::kSuccess == status) { + status = run(params_, stream); + } + return status; + } + + /// Launches the kernel after first constructing Params internal state from supplied arguments. + Status + operator()(Arguments const& args, void* workspace = nullptr, mcStream_t stream = nullptr) { + return run(args, workspace, stream); + } + + /// Overload that allows a user to re-launch the same kernel without updating internal params struct. + Status + run(mcStream_t stream = nullptr) { + return run(params_, stream); + } + + /// Overload that allows a user to re-launch the same kernel without updating internal params struct. + Status + operator()(mcStream_t stream = nullptr) { + return run(params_, stream); + } +}; + +//////////////////////////////////////////////////////////////////////////////// +////////////////////////////// MCTLASS 2.x API ///////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +template +class GemmUniversalAdapter< + GemmKernel_, + cute::enable_if_t::value>> +{ +public: + + using GemmKernel = GemmKernel_; + + static bool const kInternalTranspose = + cute::is_same::value; + + using ThreadblockShape = typename GemmKernel::Mma::Shape; + using WarpShape = typename GemmKernel::WarpShape; + using InstructionShape = typename GemmKernel::InstructionShape; + + // warp-level, arch-level (instruction), math operator + using WarpMmaOperator = typename GemmKernel::Mma::Policy::Operator; + using ArchMmaOperator = typename WarpMmaOperator::ArchMmaOperator; + using MathOperator = typename WarpMmaOperator::MathOperator; + + // Operator class and arch tag extract bottom-up + // set it for top-level gemm device-level template + using OperatorClass = typename WarpMmaOperator::OperatorClass; + using ArchTag = typename WarpMmaOperator::ArchTag; + + // Type, layout, and complex transform deliberately exchanged with B + using MapArguments = kernel::detail::MapArguments< + typename GemmKernel::ElementA, + typename GemmKernel::LayoutA, + GemmKernel::kTransformA, + GemmKernel::kAlignmentA, + typename GemmKernel::ElementB, + typename GemmKernel::LayoutB, + GemmKernel::kTransformB, + GemmKernel::kAlignmentB, + typename GemmKernel::LayoutC, + kInternalTranspose + >; + + using ElementA = typename MapArguments::ElementA; + using LayoutA = typename MapArguments::LayoutA; + static ComplexTransform const kTransformA = MapArguments::kTransformA; + static int const kAlignmentA = MapArguments::kAlignmentA; + + using ElementB = typename MapArguments::ElementB; + using LayoutB = typename MapArguments::LayoutB; + static ComplexTransform const kTransformB = MapArguments::kTransformB; + static int const kAlignmentB = MapArguments::kAlignmentB; + + using ElementC = typename GemmKernel::ElementC; + using LayoutC = typename MapArguments::LayoutC; + static int const kAlignmentC = GemmKernel::kAlignmentC; + + // C and D same type for 2.x kernel + using ElementD = ElementC; + using LayoutD = LayoutC; + + using TensorRefA = TensorRef; + using TensorRefB = TensorRef; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + + static int const kStages = GemmKernel::Mma::kStages; + + using EpilogueOutputOp = typename GemmKernel::EpilogueOutputOp; + using ElementAccumulator = typename EpilogueOutputOp::ElementAccumulator; + using ThreadblockSwizzle = typename GemmKernel::ThreadblockSwizzle; + + using UnderlyingOperator = GemmUniversalBase; + using Arguments = typename UnderlyingOperator::Arguments; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the GEMM. + GemmUniversalAdapter() { } + + /// Helper to construct a transposed equivalent for the underying GEMM operator + static Arguments to_underlying_arguments(Arguments const &args) { + if (kInternalTranspose) { + return args.transposed_problem(); + } + else { + return args; + } + } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Computes the grid shape + static dim3 get_grid_shape(Arguments const &args) { + return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args)); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int smem_capacity = -1) { + return UnderlyingOperator::maximum_active_blocks(smem_capacity); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args) { + + return underlying_operator_.update(to_underlying_arguments(args)); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::device + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_universal_base.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_universal_base.h new file mode 100644 index 0000000..fa91458 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_universal_base.h @@ -0,0 +1,419 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief The universal GEMM accommodates streamk, batched strided, and batched array variants. +*/ + + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/gemm_universal.h" + +#include "mctlass/gemm/kernel/default_gemm_universal.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +#include "mctlass/trace.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + + +template +class GemmUniversalBase { +public: + + using GemmKernel = GemmKernel_; + using ThreadblockShape = typename GemmKernel::Mma::Shape; + + using ElementA = typename GemmKernel::ElementA; + using LayoutA = typename GemmKernel::LayoutA; + using TensorRefA = TensorRef; + static ComplexTransform const kTransformA = GemmKernel::kTransformA; + + using ElementB = typename GemmKernel::ElementB; + using LayoutB = typename GemmKernel::LayoutB; + using TensorRefB = TensorRef; + static ComplexTransform const kTransformB = GemmKernel::kTransformB; + + using ElementC = typename GemmKernel::ElementC; + using LayoutC = typename GemmKernel::LayoutC; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + + /// Numerical accumulation element type + using ElementAccumulator = typename GemmKernel::Mma::ElementC; + + using EpilogueOutputOp = typename GemmKernel::EpilogueOutputOp; + using ThreadblockSwizzle = typename GemmKernel::ThreadblockSwizzle; + using Operator = typename GemmKernel::Operator; + + /// Argument structure + using Arguments = typename GemmKernel::Arguments; + +protected: + + // + // Device properties (uniform across all instances of the current thread) + // + + // Device ordinal + thread_local static int device_ordinal_; + + /// Device SM count + thread_local static int device_sms_; + + /// Kernel SM occupancy (in thread blocks) + thread_local static int sm_occupancy_; + + /// Kernel dynamic shared memory allocation requirement + thread_local static int smem_size_; + + /// Initialize static thread-local members for the thread's current device, + /// if necessary. + static Status init_device_props() + { + MCTLASS_TRACE_HOST("GemmUniversalBase::init_device_props()"); + + mcError_t mcrt_result; + + // Get current device ordinal + int current_ordinal; + mcrt_result = mcGetDevice(¤t_ordinal); + if (mcrt_result != mcSuccess) { + MCTLASS_TRACE_HOST(" mcGetDevice() returned error " << mcGetErrorString(mcrt_result)); + return Status::kErrorInternal; + } + + // Done if matches the current static member + if (current_ordinal == device_ordinal_) { + // Already initialized + return Status::kSuccess; + } + + // Update SM count member + mcrt_result = mcDeviceGetAttribute (&device_sms_, mcDeviceAttributeMultiProcessorCount, current_ordinal); + if (mcrt_result != mcSuccess) { + MCTLASS_TRACE_HOST(" mcDeviceGetAttribute() returned error " << mcGetErrorString(mcrt_result)); + return Status::kErrorInternal; + } + + // Update the kernel function's shared memory configuration for the current device + smem_size_ = int(sizeof(typename GemmKernel::SharedStorage)); + + // If requires more than 48KB: configure for extended, dynamic shared memory + if (smem_size_ >= (48 << 10)) + { + mcrt_result = mcFuncSetAttribute( + Kernel2, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size_); + if (mcrt_result != mcSuccess) { + MCTLASS_TRACE_HOST(" mcFuncSetAttribute() returned error " << mcGetErrorString(mcrt_result)); + return Status::kErrorInternal; + } + + mcrt_result = mcFuncSetAttribute( + Kernel2, + mcFuncAttributePreferredSharedMemoryCarveout, 100); // 100% shared memory + if (mcrt_result != mcSuccess) { + MCTLASS_TRACE_HOST(" mcFuncSetAttribute() returned error " << mcGetErrorString(mcrt_result)); + return Status::kErrorInternal; + } + } + + // Update SM occupancy member + mcrt_result = mcOccupancyMaxActiveBlocksPerMultiprocessorWithFlags( + &sm_occupancy_, + Kernel2, + GemmKernel::kThreadCount, + smem_size_, + mcOccupancyDisableCachingOverride); + if (mcrt_result != mcSuccess) { + MCTLASS_TRACE_HOST(" mcOccupancyMaxActiveBlocksPerMultiprocessorWithFlags() returned error " << mcGetErrorString(mcrt_result)); + return Status::kErrorInternal; + } + + // Update device ordinal member on success + device_ordinal_ = current_ordinal; + + MCTLASS_TRACE_HOST(" " + "device_ordinal: (" << device_ordinal_ << "), " + "device_sms: (" << device_sms_ << "), " + "sm_occupancy: (" << sm_occupancy_ << ") " + "smem_size: (" << smem_size_ << ") " + "GemmKernel::kThreadCount: (" << GemmKernel::kThreadCount << ")"); + + return Status::kSuccess; + } + + +protected: + + // + // Instance data members + // + + /// Kernel parameters + typename GemmKernel::Params params_; + + + /// Initialize params member + Status init_params(Arguments const &args) + { + // Initialize static device properties, if necessary + Status result = init_device_props(); + if (result != Status::kSuccess) { + return result; + } + + // Initialize params member + params_ = typename GemmKernel::Params(args, device_sms_, sm_occupancy_); + return Status::kSuccess; + } + +public: + + //--------------------------------------------------------------------------------------------- + // Stateless API + //--------------------------------------------------------------------------------------------- + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) + { + MCTLASS_TRACE_HOST("GemmUniversalBase::can_implement()"); + + // Initialize static kernel and device properties, if necessary. + Status result = init_device_props(); + if (result != Status::kSuccess) { + return result; + } + + dim3 grid = get_grid_shape(args); + + if (!(grid.y <= std::numeric_limits::max() && + grid.z <= std::numeric_limits::max())) + { + return Status::kErrorInvalidProblem; + } + + return GemmKernel::can_implement(args); + } + + + /// Returns the workspace size (in bytes) needed for the problem + /// geometry expressed by these arguments + static size_t get_workspace_size(Arguments const &args) + { + MCTLASS_TRACE_HOST("GemmUniversalBase::get_workspace_size()"); + + // Initialize parameters from args + GemmUniversalBase base; + if (base.init_params(args) != Status::kSuccess) { + return 0; + } + + // Get size from parameters + size_t workspace_bytes = base.params_.get_workspace_size(); + + MCTLASS_TRACE_HOST(" workspace_bytes: " << workspace_bytes); + return workspace_bytes; + } + + + /// Returns the grid extents in thread blocks to launch + static dim3 get_grid_shape(Arguments const &args) + { + MCTLASS_TRACE_HOST("GemmUniversalBase::get_grid_shape()"); + + // Initialize parameters from args + GemmUniversalBase base; + if (base.init_params(args) != Status::kSuccess) { + return dim3(0,0,0); + } + + // Get dims from parameters + dim3 grid_dims = base.params_.get_grid_dims(); + + MCTLASS_TRACE_HOST( + " tiled_shape: " << base.params_.get_tiled_shape() << "\n" + << " grid_dims: {" << grid_dims << "}"); + + return grid_dims; + } + + + /// Returns the maximum number of active thread blocks per multiprocessor + static int maximum_active_blocks() + { + MCTLASS_TRACE_HOST("GemmUniversalBase::maximum_active_blocks()"); + + // Initialize static device properties, if necessary + if (init_device_props() != Status::kSuccess) { + return -1; + } + + MCTLASS_TRACE_HOST(" max_active_blocks: " << sm_occupancy_); + return sm_occupancy_; + } + + + //--------------------------------------------------------------------------------------------- + // Stateful API + //--------------------------------------------------------------------------------------------- + + /// Initializes GEMM state from arguments and workspace memory + Status initialize( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) + { + MCTLASS_TRACE_HOST("GemmUniversalBase::initialize() - workspace " + << workspace << ", stream: " << (stream ? "non-null" : "null")); + + // Initialize parameters from args + Status result = init_params(args); + if (result != Status::kSuccess) { + return result; + } + + // Assign and prepare workspace memory + return params_.init_workspace(workspace, stream); + } + + + /// Lightweight update given a subset of arguments. + Status update(Arguments const &args) + { + MCTLASS_TRACE_HOST("GemmUniversalBase()::update()"); + params_.update(args); + return Status::kSuccess; + } + + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) + { + MCTLASS_TRACE_HOST("GemmUniversalBase::run()"); + + // Configure grid and block dimensions + dim3 block(GemmKernel::kThreadCount, 1, 1); + dim3 grid = params_.get_grid_dims(); + + // Launch kernel + MCTLASS_TRACE_HOST(" " + "grid: (" << grid << "), " + "block: (" << block << "), " + "SMEM: (" << smem_size_ << ")"); + + Kernel2<<>>(params_); + + // Query for errors + mcError_t result = mcGetLastError(); + if (result != mcSuccess) { + MCTLASS_TRACE_HOST(" grid launch failed with error " << mcGetErrorString(result)); + return Status::kErrorInternal; + } + + return Status::kSuccess; + } + + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) + { + return run(stream); + } + + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) + { + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Static initializers +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Device ordinal +template +thread_local int GemmUniversalBase::device_ordinal_ = -1; + +/// Device SM count +template +thread_local int GemmUniversalBase::device_sms_ = -1; + +/// Kernel SM occupancy (in thread blocks) +template +thread_local int GemmUniversalBase::sm_occupancy_ = -1; + +/// Kernel dynamic shared memory allocation requirement +template +thread_local int GemmUniversalBase::smem_size_ = -1; + + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_universal_streamk_with_broadcast.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_universal_streamk_with_broadcast.h new file mode 100644 index 0000000..4a0ffda --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_universal_streamk_with_broadcast.h @@ -0,0 +1,386 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Template for a Stream-K GEMM kernel that can broadcast bias vector in the + epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/epilogue/thread/linear_combination_bias_elementwise.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/gemm_universal.h" + +#include "mctlass/gemm/kernel/default_gemm_universal.h" +#include "mctlass/gemm/kernel/default_gemm_streamk_with_broadcast.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" +#include "mctlass/gemm/device/gemm_universal_base.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/*! + The universal GEMM with a broadcast epilogue. + Supports +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassSimt, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_ = arch::Sm70, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator - must satisfy concept of 'EpilogueWithBroadcastOp' + typename EpilogueOutputOp_ = mctlass::epilogue::thread::LinearCombinationBiasElementwise< + ElementC_, ElementAccumulator_, ElementAccumulator_, + ElementC_, ElementC_, 128 / mctlass::sizeof_bits::value>, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// Operation performed by GEMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB = ComplexTransform::kNone +> +class GemmUniversalStreamkWithBroadcast : + public GemmUniversalBase< + typename kernel::DefaultGemmStreamkWithBroadcast< + ElementA_, + LayoutA_, + TransformA, + AlignmentA, + ElementB_, + LayoutB_, + TransformB, + AlignmentB, + ElementC_, + LayoutC_, + ElementAccumulator_, + OperatorClass_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + Operator_ + >::GemmKernel + > { + + public: + + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + + using Base = GemmUniversalBase< + typename kernel::DefaultGemmStreamkWithBroadcast< + ElementA_, + LayoutA_, + TransformA, + AlignmentA, + ElementB_, + LayoutB_, + TransformB, + AlignmentB, + ElementC_, + LayoutC_, + ElementAccumulator_, + OperatorClass_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + Operator_ + >::GemmKernel + >; + + using Arguments = typename Base::Arguments; + using GemmKernel = typename Base::GemmKernel; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output exchanges problem size and operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB, + /// Operation performed by GEMM + typename Operator_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB> +class GemmUniversalStreamkWithBroadcast { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + + using UnderlyingOperator = typename GemmUniversalStreamkWithBroadcast< + ElementB, + typename layout::LayoutTranspose::type, + ElementA, + typename layout::LayoutTranspose::type, + ElementC, + layout::RowMajor, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + kAlignmentB, + kAlignmentA, + Operator, + kTransformB, + kTransformA + >::Base; + + using GemmKernel = typename UnderlyingOperator::GemmKernel; + static int const kAlignmentC = EpilogueOutputOp::kCount; + + /// Argument structure + using Arguments = typename UnderlyingOperator::Arguments; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the GEMM. + GemmUniversalStreamkWithBroadcast() { } + + /// Helper to construct a transposed equivalent for the underying GEMM operator + static Arguments to_underlying_arguments(Arguments const &args) { + return args.transposed_problem(); + } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Computes the grid shape + static dim3 get_grid_shape(Arguments const &args) { + return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args)); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int smem_capacity = -1) { + return UnderlyingOperator::maximum_active_blocks(smem_capacity); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_universal_with_broadcast.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_universal_with_broadcast.h new file mode 100644 index 0000000..b804150 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_universal_with_broadcast.h @@ -0,0 +1,386 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Template for a GEMM kernel that can broadcast bias vector in the + epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/epilogue/thread/linear_combination_bias_elementwise.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/gemm_universal.h" + +#include "mctlass/gemm/kernel/default_gemm_universal.h" +#include "mctlass/gemm/kernel/default_gemm_with_broadcast.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" +#include "mctlass/gemm/device/gemm_universal_base.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/*! + The universal GEMM with a broadcast epilogue. + Supports +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassSimt, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_ = arch::Sm70, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator - must satisfy concept of 'EpilogueWithBroadcastOp' + typename EpilogueOutputOp_ = mctlass::epilogue::thread::LinearCombinationBiasElementwise< + ElementC_, ElementAccumulator_, ElementAccumulator_, + ElementC_, ElementC_, 128 / mctlass::sizeof_bits::value>, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// Operation performed by GEMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB = ComplexTransform::kNone +> +class GemmUniversalWithBroadcast : + public GemmUniversalBase< + typename kernel::DefaultGemmWithBroadcast< + ElementA_, + LayoutA_, + TransformA, + AlignmentA, + ElementB_, + LayoutB_, + TransformB, + AlignmentB, + ElementC_, + LayoutC_, + ElementAccumulator_, + OperatorClass_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + Operator_ + >::GemmKernel + > { + + public: + + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + + using Base = GemmUniversalBase< + typename kernel::DefaultGemmWithBroadcast< + ElementA_, + LayoutA_, + TransformA, + AlignmentA, + ElementB_, + LayoutB_, + TransformB, + AlignmentB, + ElementC_, + LayoutC_, + ElementAccumulator_, + OperatorClass_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + Operator_ + >::GemmKernel + >; + + using Arguments = typename Base::Arguments; + using GemmKernel = typename Base::GemmKernel; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output exchanges problem size and operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB, + /// Operation performed by GEMM + typename Operator_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB> +class GemmUniversalWithBroadcast { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + + using UnderlyingOperator = typename GemmUniversalWithBroadcast< + ElementB, + typename layout::LayoutTranspose::type, + ElementA, + typename layout::LayoutTranspose::type, + ElementC, + layout::RowMajor, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + kAlignmentB, + kAlignmentA, + Operator, + kTransformB, + kTransformA + >::Base; + + using GemmKernel = typename UnderlyingOperator::GemmKernel; + static int const kAlignmentC = EpilogueOutputOp::kCount; + + /// Argument structure + using Arguments = typename UnderlyingOperator::Arguments; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the GEMM. + GemmUniversalWithBroadcast() { } + + /// Helper to construct a transposed equivalent for the underying GEMM operator + static Arguments to_underlying_arguments(Arguments const &args) { + return args.transposed_problem(); + } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Computes the grid shape + static dim3 get_grid_shape(Arguments const &args) { + return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args)); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int smem_capacity = -1) { + return UnderlyingOperator::maximum_active_blocks(smem_capacity); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemm_with_k_reduction.h b/csrc/mctlass/include/mctlass/gemm/device/gemm_with_k_reduction.h new file mode 100644 index 0000000..981b5c5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemm_with_k_reduction.h @@ -0,0 +1,415 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a GEMM kernel that can reduce one of the input matrix + into a vector along the K dimension. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/gemm_with_k_reduction.h" + +#include "mctlass/gemm/kernel/default_gemm_with_k_reduction.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" +#include "mctlass/gemm/device/gemm_universal_base.h" + +#include "mctlass/layout/permute.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/*! + The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and + batched array variants. +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassSimt, + /// Reduce A or B operand along the K dimension + bool ReduceKForA_ = true, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_ = arch::Sm70, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// Operation performed by GEMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB = ComplexTransform::kNone, + /// Gather operand A by using an index array + bool GatherA = false, + /// Gather operand B by using an index array + bool GatherB = false, + /// Scatter result D by using an index array + bool ScatterD = false, + /// Permute result D + typename PermuteDLayout = layout::NoPermute +> +class GemmWithKReduction : + public GemmUniversalBase< + typename kernel::DefaultGemmWithKReduction< + ElementA_, + LayoutA_, + TransformA, + AlignmentA, + ElementB_, + LayoutB_, + TransformB, + AlignmentB, + ElementC_, + LayoutC_, + ElementAccumulator_, + OperatorClass_, + ReduceKForA_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + Operator_, + SharedMemoryClearOption::kNone + >::GemmKernel + > { + + public: + + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static constexpr int kStages = Stages; + static constexpr int kAlignmentA = AlignmentA; + static constexpr int kAlignmentB = AlignmentB; + static constexpr int kAlignmentC = EpilogueOutputOp::kCount; + static constexpr ComplexTransform kTransformA = TransformA; + static constexpr ComplexTransform kTransformB = TransformB; + + using Base = GemmUniversalBase< + typename kernel::DefaultGemmWithKReduction< + ElementA_, + LayoutA_, + TransformA, + AlignmentA, + ElementB_, + LayoutB_, + TransformB, + AlignmentB, + ElementC_, + LayoutC_, + ElementAccumulator_, + OperatorClass_, + ReduceKForA_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + EpilogueOutputOp_, + ThreadblockSwizzle_, + Stages, + Operator_, + SharedMemoryClearOption::kNone + >::GemmKernel + >; + + using Arguments = typename Base::Arguments; + using GemmKernel = typename Base::GemmKernel; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output exchanges problem size and operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Reduce A or B operand along the K dimension + bool ReduceKForA_, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB, + /// Operation performed by GEMM + typename Operator_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Scatter result D by using an index array + bool ScatterD, + /// Permute result D + typename PermuteDLayout +> +class GemmWithKReduction { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + + using UnderlyingOperator = typename GemmWithKReduction< + ElementB, + typename layout::LayoutTranspose::type, + ElementA, + typename layout::LayoutTranspose::type, + ElementC, + layout::RowMajor, + ElementAccumulator, + OperatorClass, + !ReduceKForA_, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + kAlignmentB, + kAlignmentA, + Operator, + kTransformB, + kTransformA, + GatherB, + GatherA, + ScatterD, + PermuteDLayout + >::Base; + + using GemmKernel = typename UnderlyingOperator::GemmKernel; + static int const kAlignmentC = EpilogueOutputOp::kCount; + + /// Argument structure + using Arguments = typename UnderlyingOperator::Arguments; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the GEMM. + GemmWithKReduction() = default; + + /// Helper to construct a transposed equivalent for the underying GEMM operator + static Arguments to_underlying_arguments(Arguments const &args) { + return args.transposed_problem(); + } + + /// Determines whether the GEMM can execute the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Computes the grid shape + static dim3 get_grid_shape(Arguments const &args) { + return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args)); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int smem_capacity = -1) { + return UnderlyingOperator::maximum_active_blocks(smem_capacity); + } + + /// Initializes GEMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/gemv.h b/csrc/mctlass/include/mctlass/gemm/device/gemv.h new file mode 100644 index 0000000..3424f3f --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/gemv.h @@ -0,0 +1,186 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/gemm_universal.h" + +#include "mctlass/gemm/kernel/default_gemm_universal.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" +#include "mctlass/gemm/device/gemm_universal_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +class Gemv { +public: + + using GemvKernel = GemvKernel_; + + + using ElementA = typename GemvKernel::ElementA; + using LayoutA = typename GemvKernel::LayoutA; + using ElementB = typename GemvKernel::ElementB; + using ElementC = typename GemvKernel::ElementC; + + using ElementAccumulator = typename GemvKernel::ElementAccumulator; + using EpilogueOutputOp = typename GemvKernel::EpilogueOutputOp; + + static ComplexTransform const kTransformA = GemvKernel::kTransformA; + static ComplexTransform const kTransformB = GemvKernel::kTransformB; + + static int const kThreadCount = GemvKernel::kThreadCount; + static int const kThreadsPerRow = GemvKernel::kThreadsPerRow; + static int const kStages = GemvKernel::kStages; + + static int const kAlignmentA = GemvKernel::kAlignmentA; + static int const kAlignmentB = GemvKernel::kAlignmentB; + static int const kAlignmentC = GemvKernel::kAlignmentC; + + using Arguments = typename GemvKernel::Arguments; + using Params = typename GemvKernel::Params; + +private: + + Params params_; + +public: + + /// Constructs the Gemv. + Gemv() { } + + /// Determines whether the Gemv can execute the given problem. + static Status can_implement(Arguments const &args) { + + return GemvKernel::can_implement(args); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return 0; + } + + /// Computes the grid shape + static dim3 get_grid_shape(Arguments const &args, dim3 const &block) { + if(platform::is_same::value) { + return dim3((args.problem_size.row() + (block.x - 1)) / block.x, 1, args.batch_count % 65536); + } + else { + return dim3((args.problem_size.row() + (block.y - 1)) / block.y, 1, args.batch_count % 65536); + } + } + + /// Computes the block shape + static dim3 get_block_shape() { + if(platform::is_same::value) { + return dim3(kThreadCount, 1, 1); + } + else { + return dim3(kThreadsPerRow, kThreadCount / kThreadsPerRow, 1); + } + } + + /// Initializes Gemv state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + params_ = Params(args); + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + return params_.update(args); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + dim3 block = get_block_shape(); + dim3 grid = get_grid_shape(params_, block); + + int smem_size = int(sizeof(typename GemvKernel::SharedStorage)); + + // Launch + mctlass::Kernel<<>>(params_); + + // + // Query for errors + // + mcError_t result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/rank_2k.h b/csrc/mctlass/include/mctlass/gemm/device/rank_2k.h new file mode 100644 index 0000000..20dd793 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/rank_2k.h @@ -0,0 +1,547 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined Rank2K kernel. Does not compute batching or support split-K. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/rank_2k_universal.h" + +#include "mctlass/gemm/kernel/default_rank_2k_universal.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassTensorOp, + /// Tag indicating architecture to tune for + typename ArchTag_ = arch::Sm80, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementA_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementA_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementA_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementA_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = + typename threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial = false, + /// Operation performed by SYRK + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator, + /// Complex elementwise transformation + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex elementwise transformation + ComplexTransform TransformB = ComplexTransform::kNone, + /// Blas3 computation mode (symmetric/hermitian) + BlasMode BlasMode_ = BlasMode::kSymmetric> +class Rank2K { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static FillMode const kFillModeC = FillModeC; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = SplitKSerial; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + static BlasMode const kBlasMode = BlasMode_; + static int const kUpdateRank = 2; + + // static asserts for rank 2k update kernel + static_assert(platform::is_same::value, + "Rank 2K update operator support same layouts for operandA and B"); + + /// Define the kernel + using Rank2Kkernel = typename kernel::DefaultRank2KUniversal< + ElementA, + LayoutA, + kTransformA, + kAlignmentA, + ElementB, + LayoutB, + kTransformB, + kAlignmentB, + ElementC, + LayoutC, + kFillModeC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + kSplitKSerial, + Operator, + kBlasMode + >::Rank2Kkernel; + + using Arguments = typename Rank2Kkernel::Arguments; + +private: + + /// Kernel parameters object + typename Rank2Kkernel::Params params_; +public: + + /// Constructs the SYRK. + Rank2K() { } + + /// Determines whether the SYRK can execute the given problem. + static Status can_implement(Arguments const &args) { + + if (!kSplitKSerial && args.batch_count > 1) { + return Status::kErrorInvalidProblem; + } + + Status status = Rank2Kkernel::can_implement(args); + + if (FillModeC != FillMode::kLower && FillModeC != FillMode::kUpper) { + return Status::kErrorInvalidProblem; + } + + if (status != Status::kSuccess) { + return status; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + size_t bytes = 0; + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count); + + if (kSplitKSerial && args.batch_count > 1) { + + bytes += sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n()); + } + + return bytes; + } + + /// Initializes SYRK state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count); + + if (kSplitKSerial) { + if (args.batch_count > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + size_t bytes = get_workspace_size(args); + + mcError_t result = mcMemsetAsync(workspace, 0, bytes, stream); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + } + else { + + if (args.batch_count > 1) { + return Status::kErrorInvalidProblem; + } + } + + int gemm_k_size = args.problem_size.k(); + + // Initialize the Params structure + params_ = typename Rank2Kkernel::Params{ + args, + grid_tiled_shape, + gemm_k_size, + static_cast(workspace) + }; + + int smem_size = int(sizeof(typename Rank2Kkernel::SharedStorage)); + + if (smem_size >= (64 << 10)) { + mcError_t result = mcFuncSetAttribute(Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + if (kSplitKSerial && args.batch_count > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + } + + size_t workspace_bytes = get_workspace_size(args); + + if (workspace_bytes && !workspace) { + return Status::kErrorWorkspaceNull; + } + + params_.update(args, workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(Rank2Kkernel::kThreadCount, 1, 1); + + int smem_size = int(sizeof(typename Rank2Kkernel::SharedStorage)); + + mctlass::Kernel<<>>(params_); + + mcError_t result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output exchange operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB, + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial, + /// Operation performed by Rank2K update kernel + typename Operator_, + /// Complex elementwise transformation + ComplexTransform TransformA, + /// Complex elementwise transformation + ComplexTransform TransformB, + /// Blas3 computation mode (symmetric/hermitian) + BlasMode BlasMode_ + > +class Rank2K { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static FillMode const kFillModeC = FillModeC; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = SplitKSerial; + static BlasMode const kBlasMode = BlasMode_; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + static int const kUpdateRank = 2; + + /// Define the kernel + using UnderlyingOperator = typename mctlass::gemm::device::Rank2K< + ElementB, + LayoutB, + ElementA, + LayoutA, + ElementC, + layout::RowMajor, + InvertFillMode::mode, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + kAlignmentB, + kAlignmentA, + kSplitKSerial, + Operator, + kTransformA, + kTransformB, + kBlasMode + >; + + + /// Argument structure + using Arguments = typename UnderlyingOperator::Arguments; + using Rank2Kkernel = typename UnderlyingOperator::Rank2Kkernel; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the Rank2K. + Rank2K() { } + + /// Helper to construct a transposed equivalent for the underying Rank2K operator + static Arguments to_underlying_arguments(Arguments const &args) { + return args.transposed_problem(); + } + + /// Determines whether the Rank2K can execute the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Computes the grid shape + static dim3 get_grid_shape(Arguments const &args) { + return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args)); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int smem_capacity = -1) { + return UnderlyingOperator::maximum_active_blocks(smem_capacity); + } + + /// Initializes Rank2K state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace Rank2K +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/rank_2k_grouped.h b/csrc/mctlass/include/mctlass/gemm/device/rank_2k_grouped.h new file mode 100644 index 0000000..025059c --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/rank_2k_grouped.h @@ -0,0 +1,63 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Device-level grouped Rank2K. +*/ + +#pragma once + +#include "mctlass/gemm/device/base_grouped.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Rank2K Grouped +template +class Rank2KGrouped : public BaseGrouped { +public: + using Rank2Kkernel = Rank2Kkernel_; + static const mctlass::FillMode kFillModeC = Rank2Kkernel::kFillModeC; + static const mctlass::BlasMode kBlasMode = Rank2Kkernel::kBlasMode; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/rank_k.h b/csrc/mctlass/include/mctlass/gemm/device/rank_k.h new file mode 100644 index 0000000..c610225 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/rank_k.h @@ -0,0 +1,509 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined RankK kernel. Does not compute batching or support split-K. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/rank_k_universal.h" + +#include "mctlass/gemm/kernel/default_rank_k_universal.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassTensorOp, + /// Tag indicating architecture to tune for + typename ArchTag_ = arch::Sm80, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementA_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementA_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementA_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementA_, ElementC_, + ElementAccumulator_>::EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = + typename threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial = false, + /// Operation performed by SYRK + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementA_, ElementC_, + ElementAccumulator_>::Operator, + /// Complex elementwise transformation + ComplexTransform TransformA = ComplexTransform::kNone, + /// Blas3 computation mode (symmetric/hermitian) + BlasMode BlasMode_ = BlasMode::kSymmetric> +class RankK { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static FillMode const kFillModeC = FillModeC; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = SplitKSerial; + static ComplexTransform const kTransformA = TransformA; + static BlasMode const kBlasMode = BlasMode_; + static int const kUpdateRank = 1; + + /// Define the kernel + using RankKkernel = typename kernel::DefaultRankKUniversal< + ElementA, + LayoutA, + kTransformA, + kAlignmentA, + ElementC, + LayoutC, + kFillModeC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + kSplitKSerial, + Operator, + kBlasMode + >::RankKkernel; + + using Arguments = typename RankKkernel::Arguments; + +private: + + /// Kernel parameters object + typename RankKkernel::Params params_; +public: + + /// Constructs the SYRK. + RankK() { } + + /// Determines whether the SYRK can execute the given problem. + static Status can_implement(Arguments const &args) { + + if (!kSplitKSerial && args.batch_count > 1) { + return Status::kErrorInvalidProblem; + } + + Status status = RankKkernel::can_implement(args); + + if (FillModeC != FillMode::kLower && FillModeC != FillMode::kUpper) { + return Status::kErrorInvalidProblem; + } + + if (status != Status::kSuccess) { + return status; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + size_t bytes = 0; + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count); + + if (kSplitKSerial && args.batch_count > 1) { + + bytes += sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n()); + } + + return bytes; + } + + /// Initializes SYRK state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count); + + if (kSplitKSerial) { + if (args.batch_count > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + size_t bytes = get_workspace_size(args); + + mcError_t result = mcMemsetAsync(workspace, 0, bytes, stream); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + } + else { + + if (args.batch_count > 1) { + return Status::kErrorInvalidProblem; + } + } + + int gemm_k_size = args.problem_size.k(); + + // Initialize the Params structure + params_ = typename RankKkernel::Params{ + args, + grid_tiled_shape, + gemm_k_size, + static_cast(workspace) + }; + + int smem_size = int(sizeof(typename RankKkernel::SharedStorage)); + + if (smem_size >= (64 << 10)) { + mcError_t result = mcFuncSetAttribute(Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + if (kSplitKSerial && args.batch_count > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + } + + size_t workspace_bytes = get_workspace_size(args); + + if (workspace_bytes && !workspace) { + return Status::kErrorWorkspaceNull; + } + + params_.update(args, workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(RankKkernel::kThreadCount, 1, 1); + + int smem_size = int(sizeof(typename RankKkernel::SharedStorage)); + + mctlass::Kernel<<>>(params_); + + mcError_t result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for column-major output exchange operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial, + /// Operation performed by RankK update kernel + typename Operator_, + /// Complex elementwise transformation + ComplexTransform TransformA, + /// Blas3 computation mode (symmetric/hermitian) + BlasMode BlasMode_ + > +class RankK { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static FillMode const kFillModeC = FillModeC; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = SplitKSerial; + static BlasMode const kBlasMode = BlasMode_; + static int const kUpdateRank = 1; + + // Complex transform for input A matrices (function on input layout) + static ComplexTransform const kTransformA = TransformA; + + /// Define the kernel + using UnderlyingOperator = typename mctlass::gemm::device::RankK< + ElementA, + LayoutA, + ElementC, + layout::RowMajor, + InvertFillMode::mode, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + kAlignmentA, + kSplitKSerial, + Operator, + kTransformA, + kBlasMode + >; + + + /// Argument structure + using Arguments = typename UnderlyingOperator::Arguments; + using RankKkernel = typename UnderlyingOperator::RankKkernel; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the RankK. + RankK() { } + + /// Helper to construct a transposed equivalent for the underying RankK operator + static Arguments to_underlying_arguments(Arguments const &args) { + return args; + } + + /// Determines whether the RankK can execute the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Computes the grid shape + static dim3 get_grid_shape(Arguments const &args) { + return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args)); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int smem_capacity = -1) { + return UnderlyingOperator::maximum_active_blocks(smem_capacity); + } + + /// Initializes RankK state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace RankK +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/symm.h b/csrc/mctlass/include/mctlass/gemm/device/symm.h new file mode 100644 index 0000000..c72b6b1 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/symm.h @@ -0,0 +1,602 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the domcmentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined SYMM and HEMM kernels. Does not compute batching or support split-K. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/symm_universal.h" + +#include "mctlass/gemm/kernel/default_symm_universal.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Side Mode for A (kLeft or kRight) + SideMode SideModeA, + /// Fill Mode for A (kLower or kUpper) + FillMode FillModeA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal acmcmulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassTensorOp, + /// Tag indicating architecture to tune for + typename ArchTag_ = arch::Sm80, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = epilogue::thread::LinearCombination< + ElementC_, + 128 / sizeof_bits::value, + ElementAccumulator_, + ElementAccumulator_, + epilogue::thread::ScaleType::OnlyAlphaScaling + >, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial = false, + /// Operation performed by SYMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator, + /// Blas3 computation mode (symmetric/hermitian) + BlasMode BlasMode_ = BlasMode::kSymmetric> +class Symm { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using ElementAKernel = typename platform::conditional<(SideModeA == SideMode::kRight), ElementB_, ElementA_>::type; + using LayoutAKernel = typename platform::conditional<(SideModeA == SideMode::kRight), LayoutB_, LayoutA_>::type; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using ElementBKernel = typename platform::conditional<(SideModeA == SideMode::kRight), ElementA_, ElementB_>::type; + using LayoutBKernel = typename platform::conditional<(SideModeA == SideMode::kRight), LayoutA_, LayoutB_>::type; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static SideMode const kSideModeA = SideModeA; + static FillMode const kFillModeA = FillModeA; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentAKernel = (SideModeA == SideMode::kRight) ? AlignmentB : AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentBKernel = (SideModeA == SideMode::kRight) ? AlignmentA : AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = SplitKSerial; + static BlasMode const kBlasMode = BlasMode_; + + // static asserts for symm update kernel + static_assert(platform::is_same::value, + "SYMM update operator support same layouts for operand A and B"); + + /// Define the kernel + using SymmKernel = typename kernel::DefaultSymmUniversal< + ElementAKernel, + LayoutAKernel, + kSideModeA, + kFillModeA, + kAlignmentAKernel, + ElementBKernel, + LayoutBKernel, + kAlignmentBKernel, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + kSplitKSerial, + Operator, + kBlasMode + >::SymmKernel; + + using Arguments = typename SymmKernel::Arguments; + +private: + + /// Kernel parameters object + typename SymmKernel::Params params_; +public: + + /// Constructs the SYMM. + Symm() { } + + /// Determines whether the SYMM can exemcte the given problem. + static Status can_implement(Arguments const &args) { + + if (!kSplitKSerial && args.batch_count > 1) { + return Status::kErrorInvalidProblem; + } + + Status status = SymmKernel::can_implement(args); + + if (SideModeA == SideMode::kInvalid) { + return Status::kErrorInvalidProblem; + } + + if (FillModeA != FillMode::kLower && FillModeA != FillMode::kUpper) { + return Status::kErrorInvalidProblem; + } + + if (status != Status::kSuccess) { + return status; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + size_t bytes = 0; + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count); + + if (kSplitKSerial && args.batch_count > 1) { + + bytes += sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n()); + } + + return bytes; + } + + /// Initializes SYMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count); + + if (kSplitKSerial) { + if (args.batch_count > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + size_t bytes = get_workspace_size(args); + + mcError_t result = mcMemsetAsync(workspace, 0, bytes, stream); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + } + else { + + if (args.batch_count > 1) { + return Status::kErrorInvalidProblem; + } + } + + int gemm_k_size = args.problem_size.k(); + + // Swapping argument for A and B, if A was on the right side (problem size doesn't need to change here). + if (kSideModeA == SideMode::kRight) { + // Initialize the Params structure + params_ = typename SymmKernel::Params{ + args.swapped_matrices(), + grid_tiled_shape, + gemm_k_size, + static_cast(workspace) + }; + + return Status::kSuccess; + } + + // Initialize the Params structure + params_ = typename SymmKernel::Params{ + args, + grid_tiled_shape, + gemm_k_size, + static_cast(workspace) + }; + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + if (kSplitKSerial && args.batch_count > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + } + + size_t workspace_bytes = get_workspace_size(args); + + if (workspace_bytes && !workspace) { + return Status::kErrorWorkspaceNull; + } + + params_.update(args, workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(SymmKernel::kThreadCount, 1, 1); + + int smem_size = int(sizeof(typename SymmKernel::SharedStorage)); + + if (smem_size >= (64 << 10)) { + mcError_t result = mcFuncSetAttribute(Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + mctlass::Kernel<<>>(params_); + + mcError_t result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; +//////////////////////////////////////////////////////////////////////////////// + +/******************************************************************************************************** + SYMM/HEMM has 4 combinations based on Layouts {RowMajor, ColumnMajor} x Side mode {LeftSide, RightSide} + In templates and arguments to mctlass kernel, `matrix A` is always symmetric/hermitian, and `matrix B` is rectangular. + (adhering to the mcBLAS convention) + + Although, mcBLAS SYMM/HEMM only supports ColumnMajor layouts for all matrices (A, B, C/D). + + For the mainloop and symm kernel, `A` and `B` points to left-side and right-side matrices, respectively. + + Thus, for LeftSide mode `A` and `B` points to `matrix A` and `matrix B`, respectively. While for + the RightSide mode `A` and `B` points to `matrix B` and `matrix A`, respectively. + + Additionally, MCTLASS GEMM epilogue is always RowMajor, and ColumnMajor output is achieved by + transposing the GEMM problem. Thus, ColumnMajor output layout for SYMM/HEMM requires: + - Transposing `matrix A` and `matrix B` layouts + - Swapping problem size m and n values + - Swapping LeftSide and RightSide mode + + RowMajor output: D = matrix A x matrix B + ColumnMajor output: D = matrix A x matrix B -> Transpose (D) = Transpose(matrix B) x Transpose(matrix A) + + {RowMajor, ColumnMajor} x Side Mode {LeftSide, RightSide} 4 cases: + 1. LeftSide mode and RowMajor output (default template) + 2. LeftSide mode and ColumnMajor output + 3. RightSide mode and RowMajor output + 4. RightSide mode and ColumnMajor output + + Mapping ColumnMajor output layout cases 2 and 4 to RowMajor efficient epilogue implementation: + + Case 2 -> Case 3: + D_col = matrix A x matrix B (LeftSide mode) + => Transpose(D_col) = Transpose(matrix B) x Transpose(matrix A) (RightSide mode) + + swap pointers for `A` and `B` call GEMM mainloop with RowMajor efficient-epilogue + + Case 4 -> Case 1: + D_col = matrix B x matrix A (RightSide mode) + => Transpose(D_col) = Transpose(matrix A) x Transpose(matrix B) (LeftSide mode) + + call GEMM mainloop for with RowMajor efficient-epilogue +********************************************************************************************************/ + +/// Partial specialization for column-major output exchanges problem size and operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Side Mode for A (kLeft or kRight) + SideMode SideModeA, + /// Fill Mode for A (kLower or kUpper) + FillMode FillModeA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal acmcmulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB, + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial, + /// Operation performed by Symm update kernel + typename Operator_, + /// Blas3 computation mode (symmetric/hermitian) + BlasMode BlasMode_ + > +class Symm { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static SideMode const kSideModeA = SideModeA; + static FillMode const kFillModeA = FillModeA; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = SplitKSerial; + static BlasMode const kBlasMode = BlasMode_; + + /// Define the kernel + using UnderlyingOperator = typename mctlass::gemm::device::Symm< + ElementA, + typename layout::LayoutTranspose::type, + InvertSideMode::mode, + InvertFillMode::mode, + ElementB, + typename layout::LayoutTranspose::type, + ElementC, + layout::RowMajor, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + kAlignmentA, + kAlignmentB, + kSplitKSerial, + Operator, + kBlasMode + >; + + + /// Argument structure + using Arguments = typename UnderlyingOperator::Arguments; + using SymmKernel = typename UnderlyingOperator::SymmKernel; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the Symm. + Symm() { } + + /// Helper to construct a transposed equivalent for the underying SYMM operator + static Arguments to_underlying_arguments(Arguments const &args) { + return args.transposed_problem_size(); + } + + /// Determines whether the Symm can exemcte the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Computes the grid shape + static dim3 get_grid_shape(Arguments const &args) { + return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args)); + } + + /// Computes the maximum number of active blocks per multiprocessor + static int maximum_active_blocks(int smem_capacity = -1) { + return UnderlyingOperator::maximum_active_blocks(smem_capacity); + } + + /// Initializes Symm state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace Symm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/device/trmm.h b/csrc/mctlass/include/mctlass/gemm/device/trmm.h new file mode 100644 index 0000000..3aadef4 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/device/trmm.h @@ -0,0 +1,758 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the domcmentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a TRMM kernel. Does not compute batching or support split-K. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/arch/arch.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/gemm/kernel/trmm_universal.h" + +#include "mctlass/gemm/kernel/default_trmm_universal.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/*! Trmm device-level operator. This is an interface to efficient MCTLASS TRMM kernels that may + be invoked from host code. + + The contributions of this class are: + + 1. At compile time, it maps data types and high-level structural parameters onto + specific MCTLASS components. + + 2. At runtime, it maps logical arguments to TRMM problems to kernel parameters. + + 3. At runtime, it launches kernels on the device. + + The intent is to provide a convenient mechanism for interacting with most plausible TRMM + configurations for each supported architecture. Consequently, not all parameters are exposed + to the top-level interface. Rather, sensible defaults at each level of the MCTLASS hierarchy + are selected to tradeoff simplicity of the interface with flexibility. We expect + most configurations to be specified at this level. Applications with more exotic requirements + may construct their kernels of interest using MCTLASS components at the threadblock, warp, + and thread levels of abstraction. + + MCTLASS exposes computations using the functor design pattern in which objects compose some + internal state with an overloaded function call operator. This enables decoupling of + initialization from exemction, possibly reducing overhead during steady state phases of + application exemction. + + MCTLASS device-level operators expose an Arguments structure encompassing each logical + input to the computation. This is distinct from the kernel-level Params structure pattern + which contains application-specific precomputed state needed by the device code. + + Example of a MCTLASS TRMM operator implementing the functionality of mcBLAS's STRMM NN + is as follows: + + // + // Instantiate the MCTLASS TRMM operator. + // + + mctlass::gemm::device::Trmm< + float, + mctlass::layout::ColumnMajor, + mctlass::SideMode::kLeft, + mctlass::FillMode::kLower, + mctlass::DiagType::kNonUnit, + float, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::ColumnMajor, + > trmm_op; + + // + // Launch the TRMM operation on the device + // + + mctlass::Status status = trmm_op({ + mctlass::gemm::GemmUniversalMode, // Trmm Problem Mode + {m, n, m/n}, // GemmCoord problem_size (k is based on left- or right-side mode) + batch_count, + {alpha}, // EpilogueOutputOp::Params epilogue_op_params + void const * ptr_A, + void const * ptr_B, + void const * ptr_C, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C, + int lda, + int ldb, + int ldc + }); + + A simplified view of the template is listed below. + + template < + /// Element type for A matrix operand + typename ElementA, + + /// Layout type for A matrix operand + typename LayoutA, + + /// Side Mode for A (kLeft or kRight) + SideMode SideModeA, + + /// Fill Mode for A (kLower or kUpper) + FillMode FillModeA, + + /// DiagType for A (kNonUnit or kUnit) + DiagType DiagTypeA, + + /// Element type for B matrix operand + typename ElementB, + + /// Layout type for B matrix operand + typename LayoutB, + + /// Element type for C and D matrix operands + typename ElementC, + + /// Layout type for C and D matrix operands + typename LayoutC, + + /// Element type for internal acmcmulation + typename ElementAccumulator, + + /// Operator class tag + typename OperatorClass, + + /// Tag indicating architecture to tune for. This is the minimum SM that + /// supports the intended feature. The device kernel can be built + /// targeting any SM larger than this number. + typename ArchTag, + + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + + /// Epilogue output operator + typename EpilogueOutputOp, + + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + + /// Number of stages used in the pipelined mainloop + int Stages, + + /// Access granularity of A matrix in units of elements + int AlignmentA, + + /// Access granularity of B matrix in units of elements + int AlignmentB, + + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial, + + /// Operation performed by TRMM + typename Operator, + + /// Complex elementwise transformation on A operand + ComplexTransform TransformA + > + class Trmm; +*/ +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Side Mode for A + SideMode SideModeA, + /// Fill Mode for A + FillMode FillModeA, + /// DiagType for A + DiagType DiagTypeA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal acmcmulation + typename ElementAccumulator_ = ElementC_, + /// Operator class tag + typename OperatorClass_ = arch::OpClassTensorOp, + /// Tag indicating architecture to tune for + typename ArchTag_ = arch::Sm80, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp_ = epilogue::thread::LinearCombination< + ElementC_, + 128 / sizeof_bits::value, + ElementAccumulator_, + ElementAccumulator_, + epilogue::thread::ScaleType::OnlyAlphaScaling + >, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>, + /// Number of stages used in the pipelined mainloop + int Stages = + DefaultGemmConfiguration::kStages, + /// Access granularity of A matrix in units of elements + int AlignmentA = + DefaultGemmConfiguration::kAlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB = + DefaultGemmConfiguration::kAlignmentB, + /// If true, kernel supports split-K with serial reduction + bool SplitKSerial = false, + /// Operation performed by TRMM + typename Operator_ = typename DefaultGemmConfiguration< + OperatorClass_, ArchTag_, ElementA_, ElementB_, ElementC_, + ElementAccumulator_>::Operator, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA = ComplexTransform::kNone> +class Trmm { + public: + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementAKernel = typename platform::conditional<(SideModeA == SideMode::kRight), ElementB_, ElementA_>::type; + using LayoutAKernel = typename platform::conditional<(SideModeA == SideMode::kRight), LayoutB_, LayoutA_>::type; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementBKernel = typename platform::conditional<(SideModeA == SideMode::kRight), ElementA_, ElementB_>::type; + using LayoutBKernel = typename platform::conditional<(SideModeA == SideMode::kRight), LayoutA_, LayoutB_>::type; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static SideMode const kSideMode = SideModeA; + static FillMode const kFillMode = FillModeA; + static DiagType const kDiagType = DiagTypeA; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentAKernel = (SideModeA == SideMode::kRight) ? AlignmentB : AlignmentA; + static int const kAlignmentB = AlignmentB; + static int const kAlignmentBKernel = (SideModeA == SideMode::kRight) ? AlignmentA : AlignmentB; + static int const kAlignmentC = EpilogueOutputOp::kCount; + static bool const kSplitKSerial = SplitKSerial; + // Complex Transform don't appply to B + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + static ComplexTransform const kTransformAKernel = (SideModeA == SideMode::kRight) ? + ComplexTransform::kNone : TransformA; + static ComplexTransform const kTransformBKernel = (SideModeA == SideMode::kRight) ? + TransformA : ComplexTransform::kNone; + + /// Define the kernel + using TrmmKernel = typename kernel::DefaultTrmmUniversal< + ElementAKernel, + LayoutAKernel, + kTransformAKernel, + kAlignmentAKernel, + ElementBKernel, + LayoutBKernel, + kTransformBKernel, + kAlignmentBKernel, + kSideMode, + kFillMode, + kDiagType, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + kSplitKSerial, + Operator + >::TrmmKernel; + + using Arguments = typename TrmmKernel::Arguments; + +private: + + /// Kernel parameters object + typename TrmmKernel::Params params_; +public: + + /// Constructs the TRMM. + Trmm() { } + + /// Determines whether the TRMM can exemcte the given problem. + static Status can_implement(Arguments const &args) { + + if (!kSplitKSerial && args.batch_count > 1) { + return Status::kErrorInvalidProblem; + } + + Status status = TrmmKernel::can_implement(args); + + if (SideModeA == SideMode::kInvalid) { + return Status::kErrorInvalidProblem; + } + + if (FillModeA == FillMode::kInvalid) { + return Status::kErrorInvalidProblem; + } + + if (DiagTypeA == DiagType::kInvalid) { + return Status::kErrorInvalidProblem; + } + + if (status != Status::kSuccess) { + return status; + } + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + size_t bytes = 0; + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count); + + if (kSplitKSerial && args.batch_count > 1) { + + bytes += sizeof(int) * size_t(tiled_shape.m()) * size_t(tiled_shape.n()); + } + + return bytes; + } + + /// Initializes TRMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + // Determine grid shape + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord grid_tiled_shape = threadblock_swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count); + + if (kSplitKSerial) { + if (args.batch_count > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + + size_t bytes = get_workspace_size(args); + + mcError_t result = mcMemsetAsync(workspace, 0, bytes, stream); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + } + else { + + if (args.batch_count > 1) { + return Status::kErrorInvalidProblem; + } + } + + int gemm_k_size = args.problem_size.k(); + + // Swapping argument for A and B, if A was on the right side (problem size doesn't need to change here). + if (kSideMode == SideMode::kRight) { + // Initialize the Params structure + params_ = typename TrmmKernel::Params{ + args.swapped_matrices(), + grid_tiled_shape, + gemm_k_size, + static_cast(workspace) + }; + + return Status::kSuccess; + } + + // Initialize the Params structure + params_ = typename TrmmKernel::Params{ + args, + grid_tiled_shape, + gemm_k_size, + static_cast(workspace) + }; + + return Status::kSuccess; + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + if (kSplitKSerial && args.batch_count > 1) { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + } + + size_t workspace_bytes = get_workspace_size(args); + + if (workspace_bytes && !workspace) { + return Status::kErrorWorkspaceNull; + } + + params_.update(args, workspace); + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + ThreadblockSwizzle threadblock_swizzle; + + dim3 grid = threadblock_swizzle.get_grid_shape(params_.grid_tiled_shape); + dim3 block(TrmmKernel::kThreadCount, 1, 1); + + int smem_size = int(sizeof(typename TrmmKernel::SharedStorage)); + + if (smem_size >= (64 << 10)) { + mcError_t result = mcFuncSetAttribute(Kernel, + mcFuncAttributeMaxDynamicSharedMemorySize, + smem_size); + + if (result != mcSuccess) { + return Status::kErrorInternal; + } + } + + mctlass::Kernel<<>>(params_); + + mcError_t result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +/******************************************************************************************************** + TRMM has 4 combinations based on Layouts {RowMajor, ColumnMajor} x Side mode {LeftSide, RightSide} + In templates and arguments to mctlass kernel, `matrix A` is always triangular, and `matrix B` is rectangular. + (adhering to the mcBLAS convention) + +For the mainloop and trmm kernel, `A` and `B` points to left-side and right-side matrices, respectively. + + Thus, for LeftSide mode `A` and `B` points to `matrix A` and `matrix B`, respectively. While for + the RightSide mode `A` and `B` points to `matrix B` and `matrix A`, respectively. + + Additionally, MCTLASS GEMM epilogue is always RowMajor, and ColumnMajor output is achieved by + transposing the GEMM problem. Thus, ColumnMajor output layout for TRMM requires: + - Transposing `matrix A` and `matrix B` layouts + - Swapping problem size m and n values + - Swapping LeftSide and RightSide mode + + RowMajor output: D = matrix A x matrix B + ColumnMajor output: D = matrix A x matrix B -> Transpose (D) = Transpose(matrix B) x Transpose(matrix A) + + {RowMajor, ColumnMajor} x Side Mode {LeftSide, RightSide} 4 cases: + 1. LeftSide mode and RowMajor output (default template) + 2. LeftSide mode and ColumnMajor output + 3. RightSide mode and RowMajor output + 4. RightSide mode and ColumnMajor output + + Mapping ColumnMajor output layout cases 2 and 4 to RowMajor efficient epilogue implementation: + + Case 2 -> Case 3: + D_col = matrix A x matrix B (LeftSide mode) + => Transpose(D_col) = Transpose(matrix B) x Transpose(matrix A) (RightSide mode) + + swap pointers for `A` and `B` call GEMM mainloop with RowMajor efficient-epilogue + + Case 4 -> Case 1: + D_col = matrix B x matrix A (RightSide mode) + => Transpose(D_col) = Transpose(matrix A) x Transpose(matrix B) (LeftSide mode) + + call GEMM mainloop for with RowMajor efficient-epilogue +********************************************************************************************************/ + +/// Partial specialization for column-major output exchanges problem size and operand. +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Side Mode for A + SideMode SideModeA, + /// Fill Mode for A + FillMode FillModeA, + /// DiagType for A + DiagType DiagTypeA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Element type for internal acmcmulation + typename ElementAccumulator_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Epilogue output operator + typename EpilogueOutputOp_, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Access granularity of A matrix in units of elements + int AlignmentA, + /// Access granularity of B matrix in units of elements + int AlignmentB, + /// If true, kernel supports split-K as a serial reduction + bool SplitKSerial, + /// Operation performed by TRMM + typename Operator_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA> +class Trmm { + public: + + using ElementA = ElementA_; + using LayoutA = LayoutA_; + using TensorRefA = TensorRef; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + using TensorRefB = TensorRef; + using ElementC = ElementC_; + using LayoutC = layout::ColumnMajor; + using TensorRefC = TensorRef; + using TensorRefD = TensorRef; + using ElementAccumulator = ElementAccumulator_; + using OperatorClass = OperatorClass_; + using ArchTag = ArchTag_; + using ThreadblockShape = ThreadblockShape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using EpilogueOutputOp = EpilogueOutputOp_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + using Operator = Operator_; + static SideMode const kSideMode = SideModeA; + static FillMode const kFillMode = FillModeA; + static DiagType const kDiagType = DiagTypeA; + // Changing SideMode as we change the layout + static SideMode const kSideModeT = (SideModeA == SideMode::kLeft) ? + SideMode::kRight : SideMode::kLeft; + // Changing FillMode as we change the layout + static FillMode const kFillModeT = (FillModeA == FillMode::kLower) ? + FillMode::kUpper : FillMode::kLower; + static int const kStages = Stages; + static int const kAlignmentA = AlignmentA; + static int const kAlignmentB = AlignmentB; + static ComplexTransform const kTransformA = TransformA; + // Complex Transform don't appply to B + static ComplexTransform const kTransformB = ComplexTransform::kNone; + static bool const kSplitKSerial = SplitKSerial; + + using UnderlyingOperator = Trmm< + ElementA, + typename layout::LayoutTranspose::type, + kSideModeT, + kFillModeT, + kDiagType, + ElementB, + typename layout::LayoutTranspose::type, + ElementC, + layout::RowMajor, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + kStages, + kAlignmentA, + kAlignmentB, + kSplitKSerial, + Operator, + TransformA + >; + + using Arguments = typename UnderlyingOperator::Arguments; + using TrmmKernel = typename UnderlyingOperator::TrmmKernel; + static int const kAlignmentC = UnderlyingOperator::kAlignmentC; + +private: + + UnderlyingOperator underlying_operator_; + +public: + + /// Constructs the TRMM. + Trmm() { } + + /// Helper to construct a transposed equivalent for the underying TRMM operator which is identical + static Arguments to_underlying_arguments(Arguments const &args) { + return args.transposed_problem_size(); + } + + /// Determines whether the TRMM can exemcte the given problem. + static Status can_implement(Arguments const &args) { + + return UnderlyingOperator::can_implement(to_underlying_arguments(args)); + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + + return UnderlyingOperator::get_workspace_size(to_underlying_arguments(args)); + } + + /// Initializes TRMM state from arguments. + Status initialize(Arguments const &args, void *workspace = nullptr, mcStream_t stream = nullptr) { + + return underlying_operator_.initialize(to_underlying_arguments(args), workspace, stream); + } + + /// Lightweight update given a subset of arguments + Status update(Arguments const &args, void *workspace = nullptr) { + + return underlying_operator_.update(to_underlying_arguments(args), workspace); + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + return underlying_operator_.run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/dispatch_policy.hpp b/csrc/mctlass/include/mctlass/gemm/dispatch_policy.hpp new file mode 100644 index 0000000..23f8e8c --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/dispatch_policy.hpp @@ -0,0 +1,170 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/arch/arch.h" + +#include "cute/layout.hpp" +#include "cute/numeric/integral_constant.hpp" + +////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm { +using namespace cute; + +////////////////////////////////////////////////////////////////////////////// + +// +// Policies for categorical dispatch of mainloop against kernel grid schedules +// +struct KernelMultistage { }; +struct KernelTma { }; +struct KernelTmaWarpSpecialized { }; +struct KernelTmaWarpSpecializedPingpong { }; +struct KernelTmaWarpSpecializedCooperative { }; + +// Policies for dispatch of epilogue +struct EpilogueDefault { }; +struct EpilogueTransposed { }; + +// +// Collective Mainloop Policies +// + +// 2 stage pipeline through 1 stage in smem, 1 in rmem, WITHOUT predicated gmem loads +struct MainloopSm70TwoStageUnpredicated { + constexpr static int Stages = 2; + using ArchTag = arch::Sm70; + using Schedule = KernelMultistage; + using ClusterShape = Shape<_1,_1,_1>; +}; + +// 2 stage pipeline through 1 stage in smem, 1 in rmem, with predicated gmem loads +struct MainloopSm70TwoStage { + constexpr static int Stages = 2; + using ArchTag = arch::Sm70; + using Schedule = KernelMultistage; + using ClusterShape = Shape<_1,_1,_1>; +}; + +// n-buffer in smem (cp.async), pipelined with registers, WITHOUT predicated gmem loads +template +struct MainloopSm80CpAsyncUnpredicated { + constexpr static int Stages = Stages_; + using ArchTag = arch::Sm80; + using Schedule = KernelMultistage; + using ClusterShape = Shape<_1,_1,_1>; +}; + +// n-buffer in smem (cp.async), pipelined with registers, with predicated gmem loads +template +struct MainloopSm80CpAsync { + constexpr static int Stages = Stages_; + using ArchTag = arch::Sm80; + using Schedule = KernelMultistage; + using ClusterShape = Shape<_1,_1,_1>; +}; + +// n-buffer in smem (cp.async), pipelined with Hopper GMMA, WITHOUT predicated gmem loads +template< + int Stages_, + class ClusterShape_ = Shape<_1,_1,_1> +> +struct MainloopSm90CpAsyncGmmaUnpredicated { + constexpr static int Stages = Stages_; + using ClusterShape = ClusterShape_; + using ArchTag = arch::Sm90; + using Schedule = KernelMultistage; +}; + +// n-buffer in smem (cp.async), pipelined with Hopper GMMA, with predicated gmem loads +template< + int Stages_, + class ClusterShape_ = Shape<_1,_1,_1> +> +struct MainloopSm90CpAsyncGmma { + constexpr static int Stages = Stages_; + using ClusterShape = ClusterShape_; + using ArchTag = arch::Sm90; + using Schedule = KernelMultistage; +}; + +// n-buffer in smem (Hopper TMA), pipelined with Hopper GMMA and TMA, static schedule between TMA and GMMA +template< + int Stages_, + class ClusterShape_ = Shape<_1,_1,_1>, + int PipelineAsyncMmaStages_ = 1 +> +struct MainloopSm90TmaGmma { + constexpr static int Stages = Stages_; + using ClusterShape = ClusterShape_; + constexpr static int PipelineAsyncMmaStages = PipelineAsyncMmaStages_; + using ArchTag = arch::Sm90; + using Schedule = KernelTma; +}; + +// n-buffer in smem (Hopper TMA), pipelined with Hopper GMMA and TMA, Warp specialized dynamic schedule +template< + int Stages_, + class ClusterShape_ = Shape<_1,_1,_1>, + class KernelSchedule = KernelTmaWarpSpecializedCooperative +> +struct MainloopSm90TmaGmmaWarpSpecialized { + constexpr static int Stages = Stages_; + using ClusterShape = ClusterShape_; + using ArchTag = arch::Sm90; + using Schedule = KernelSchedule; +}; + +// n-buffer in smem (Hopper TMA), pipelined with Hopper GMMA and TMA, Warp specialized dynamic schedule +// With GMMA's A data from registers. +template< + int Stages_, + class ClusterShape_ = Shape<_1,_1,_1>, + class KernelSchedule = KernelTmaWarpSpecialized, + int PipelineAsyncMmaStages_ = 0 +> +struct MainloopSm90TmaGmmaRmemAWarpSpecialized { + constexpr static int Stages = Stages_; + using ClusterShape = ClusterShape_; + constexpr static int PipelineAsyncMmaStages = PipelineAsyncMmaStages_; + using ArchTag = arch::Sm90; + using Schedule = KernelSchedule; + static_assert( + cute::is_same_v || + cute::is_same_v || + cute::is_same_v, + "KernelSchedule must be one of the warp specialized policies"); +}; + +////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm diff --git a/csrc/mctlass/include/mctlass/gemm/gemm.h b/csrc/mctlass/include/mctlass/gemm/gemm.h new file mode 100644 index 0000000..ed5beb1 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/gemm.h @@ -0,0 +1,623 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines common types used for all GEMM-like operators. +*/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/coord.h" +#include "mctlass/layout/matrix.h" +#include "cute/layout.hpp" +// #include "cute/arch/copy_sm90_tma.hpp" + +namespace mctlass { +namespace gemm { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// GEMM operand enumeration: D = A * B + C +enum class Operand { + kA, /// A multiplicand + kB, /// B multiplicand + kC, /// Source accumulator + kD /// Destination accumulator +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Shape of a matrix multiply-add operation +template < + /// Rows of matrix product + int M = 1, + /// Columns of matrix product + int N = 1, + /// Inner dimension of matrix product + int K = 1 +> +struct GemmShape { + static int const kM = M; + static int const kN = N; + static int const kK = K; + + static int const kMN = M * N; + static int const kMK = M * K; + static int const kKN = N * K; + static int const kMNK = M * N * K; + + static int const kCount = kMNK; + + // + // Static member functions + // + + /// Returns a Coord object + MCTLASS_HOST_DEVICE + static Coord<3> toCoord() { + return make_Coord(kM, kN, kK); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Type alias of the transpose of a GemmShape +template < + /// concept: GemmShape + typename Shape +> +using GemmShapeTranspose = GemmShape; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// GemmCoord is a structure derived from Coord<3> that specifies a location within the +/// coordinate space of a GEMM problem. +struct GemmCoord : public Coord<3, int> { + + /// Integer-valued index + typedef int Index; + + /// Base type is a Coord of rank=3 + typedef Coord<3, Index> Base; + + /// GEMM M dimension - rows of the output C matrix + static int const kM = 0; + + /// GEMM N dimension - columns of the output C matrix + static int const kN = 1; + + /// GEMM K dimension - inner dimension of the GEMM problem + static int const kK = 2; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + GemmCoord() { } + + /// Constructs from Coord<3> and a batch + MCTLASS_HOST_DEVICE + GemmCoord(Coord<3, Index> const &coord): Base(make_Coord(coord[0], coord[1], coord[2])) { } + + /// Helper to construct from a K, N, M, batch variables + MCTLASS_HOST_DEVICE + GemmCoord(Index m, Index n, Index k): Base(make_Coord(m, n, k)) { } + + /// Returns the GEMM M coordinate + MCTLASS_HOST_DEVICE + Index const & m() const { return this->at(kM); } + + /// Returns reference to the GEMM M coordinate + MCTLASS_HOST_DEVICE + Index & m() { return this->at(kM); } + + /// Returns the GEMM N coordinate + MCTLASS_HOST_DEVICE + Index const & n() const { return this->at(kN); } + + /// Returns reference to the GEMM N coordinate + MCTLASS_HOST_DEVICE + Index & n() { return this->at(kN); } + + /// Returns the GEMM K coordinate + MCTLASS_HOST_DEVICE + Index const & k() const { return this->at(kK); } + + /// Returns reference to the GEMM K coordinate + MCTLASS_HOST_DEVICE + Index & k() { return this->at(kK); } + + /// Obtains a Coord<3> from GemmCoord + MCTLASS_HOST_DEVICE + Coord<3> mnk() const { + return make_Coord(m(), n(), k()); + } + + /// Obtains a Coord<3> from GemmCoord + MCTLASS_HOST_DEVICE + Coord<3> knm() const { + return make_Coord(k(), n(), m()); + } + + /// Obtains a Coord<2> from GemmCoord + MCTLASS_HOST_DEVICE + Coord<2> nm() const { + return make_Coord(n(), m()); + } + + /// Obtains a Coord<2> from GemmCoord + MCTLASS_HOST_DEVICE + Coord<2> mn() const { + return make_Coord(m(), n()); + } + + /// Obtains a Coord<2> from GemmCoord + MCTLASS_HOST_DEVICE + Coord<2> mk() const { + return make_Coord(m(), k()); + } + + /// Obtains a Coord<2> from GemmCoord + MCTLASS_HOST_DEVICE + Coord<2> km() const { + return make_Coord(k(), m()); + } + + /// Obtains a Coord<2> from GemmCoord + MCTLASS_HOST_DEVICE + Coord<2> nk() const { + return make_Coord(n(), k()); + } + + /// Obtains a Coord<2> from GemmCoord + MCTLASS_HOST_DEVICE + Coord<2> kn() const { + return make_Coord(k(), n()); + } + + // + // Coord operators + // + + /// Element-wise addition + MCTLASS_HOST_DEVICE + GemmCoord operator+(Base const& b) const { + return GemmCoord(Base::operator+(b)); + } + + /// Element-wise subtraction + MCTLASS_HOST_DEVICE + GemmCoord operator-(Base const& b) const { + return GemmCoord(Base::operator-(b)); + } + + /// Element-wise multiplication + MCTLASS_HOST_DEVICE + GemmCoord operator*(Base const& b) const { + return GemmCoord(Base::operator*(b)); + } + + /// Element-wise division + MCTLASS_HOST_DEVICE + GemmCoord operator/(Base const& b) const { + return GemmCoord(Base::operator/(b)); + } + + /// In-place addition + MCTLASS_HOST_DEVICE + GemmCoord& operator+=(Base const& b) { + Base::operator+=(b); + return *this; + } + + /// In-place subtraction + MCTLASS_HOST_DEVICE + GemmCoord& operator-=(Base const& b) { + Base::operator-=(b); + return *this; + } + + /// In-place multiplication + MCTLASS_HOST_DEVICE + GemmCoord& operator*=(Base const& b) { + Base::operator*=(b); + return *this; + } + + /// In-place division + MCTLASS_HOST_DEVICE + GemmCoord& operator/=(Base const& b) { + Base::operator/=(b); + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// BatchedGemmCoord is a structure derived from Coord<4> that specifies a location within the +/// coordinate space of a batched GEMM problem. +struct BatchedGemmCoord : public Coord<4, int> { + + /// Integer-valued index + typedef int Index; + + /// Base type is a Coord of rank=4 + typedef Coord<4, Index> Base; + + /// GEMM M dimension - rows of the output C matrix + static int const kM = 0; + + /// GEMM N dimension - columns of the output C matrix + static int const kN = 1; + + /// GEMM K dimension - inner dimension of the GEMM problem + static int const kK = 2; + + /// GEMM Batch dimension - inner dimension of the GEMM problem + static int const kBatch = 3; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + BatchedGemmCoord() { } + + /// Constructs from Coord<4> + MCTLASS_HOST_DEVICE + BatchedGemmCoord(Base const &coord): Base(coord) { } + + /// Helper to construct from a K, N, M, and batch variables + MCTLASS_HOST_DEVICE + BatchedGemmCoord(Index m, Index n, Index k, Index b): Base(make_Coord(m, n, k, b)) { } + + /// Returns the GEMM M coordinate + MCTLASS_HOST_DEVICE + Index const & m() const { return this->at(kM); } + + /// Returns reference to the GEMM M coordinate + MCTLASS_HOST_DEVICE + Index & m() { return this->at(kM); } + + /// Returns the GEMM N coordinate + MCTLASS_HOST_DEVICE + Index const & n() const { return this->at(kN); } + + /// Returns reference to the GEMM N coordinate + MCTLASS_HOST_DEVICE + Index & n() { return this->at(kN); } + + /// Returns the GEMM K coordinate + MCTLASS_HOST_DEVICE + Index const & k() const { return this->at(kK); } + + /// Returns reference to the GEMM K coordinate + MCTLASS_HOST_DEVICE + Index & k() { return this->at(kK); } + + /// Returns the GEMM batch coordinate + MCTLASS_HOST_DEVICE + Index const & batch() const { return this->at(kBatch); } + + /// Returns reference to the GEMM batch coordinate + MCTLASS_HOST_DEVICE + Index & batch() { return this->at(kBatch); } + + /// Obtains a GemmCoord from BatchedGemmCoord + MCTLASS_HOST_DEVICE + GemmCoord mnk() const { + return GemmCoord(m(), n(), k()); + } + + /// Obtains a Coord<4> from BatchedGemmCoord + MCTLASS_HOST_DEVICE + Coord<4> mnkb() const { + return make_Coord(m(), n(), k(), batch()); + } + + // + // Coord operators + // + + /// Element-wise addition + MCTLASS_HOST_DEVICE + BatchedGemmCoord operator+(Base const& b) const { + return BatchedGemmCoord(Base::operator+(b)); + } + + /// Element-wise subtraction + MCTLASS_HOST_DEVICE + BatchedGemmCoord operator-(Base const& b) const { + return BatchedGemmCoord(Base::operator-(b)); + } + + /// Element-wise multiplication + MCTLASS_HOST_DEVICE + BatchedGemmCoord operator*(Base const& b) const { + return BatchedGemmCoord(Base::operator*(b)); + } + + /// Element-wise division + MCTLASS_HOST_DEVICE + BatchedGemmCoord operator/(Base const& b) const { + return BatchedGemmCoord(Base::operator/(b)); + } + + /// In-place addition + MCTLASS_HOST_DEVICE + BatchedGemmCoord& operator+=(Base const& b) { + Base::operator+=(b); + return *this; + } + + /// In-place subtraction + MCTLASS_HOST_DEVICE + BatchedGemmCoord& operator-=(Base const& b) { + Base::operator-=(b); + return *this; + } + + /// In-place multiplication + MCTLASS_HOST_DEVICE + BatchedGemmCoord& operator*=(Base const& b) { + Base::operator*=(b); + return *this; + } + + /// In-place division + MCTLASS_HOST_DEVICE + BatchedGemmCoord& operator/=(Base const& b) { + Base::operator/=(b); + return *this; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +enum class GemmUniversalMode { + kGemm, + kGemmSplitKParallel, + kBatched, + kArray, + kInvalid +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Some options for clearing shared memory +enum class SharedMemoryClearOption { + kNone, ///< SMEM is in don't-care state + kZfill, ///< Kernels fill out of bounds accesses with zeros + kClearLastStage ///< Last SMEM stage is explicitly cleared. Mainloop uses 'kNone' +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// For each mctlass::layout, provides its corresponding cute stride types, 64b by default +template +struct TagToStrideA { + using type = L; +}; + +// Maps to modes [M, K, L] +template <> +struct TagToStrideA { + using type = cute::Stride, int64_t>; + using tag = layout::RowMajor; +}; + +// Maps to modes [M, K, L] +template <> +struct TagToStrideA { + using type = cute::Stride, int64_t, int64_t>; + using tag = layout::ColumnMajor; +}; + +template +struct TagToStrideB { + using type = L; +}; + +// Maps to modes [N, K, L] +template <> +struct TagToStrideB { + using type = cute::Stride, int64_t, int64_t>; + using tag = layout::RowMajor; +}; + +// Maps to modes [N, K, L] +template <> +struct TagToStrideB { + using type = cute::Stride, int64_t>; + using tag = layout::ColumnMajor; +}; + +// Maps to modes [N, N, L] +template +struct TagToStrideC : TagToStrideA { }; + +// Convenience aliases +template +using TagToStrideA_t = typename TagToStrideA::type; + +template +using TagToStrideB_t = typename TagToStrideB::type; + +template +using TagToStrideC_t = typename TagToStrideC::type; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// For 2.x compatibility APIs, provide stride->layout tag mappers + +namespace detail { + +template +constexpr bool +is_mn_major() { + // Account for stride types with and without batch mode and batch modes with static zero stride + return cute::is_constant<1, decltype(cute::size<0,0>(Stride{}))>::value; +} + +// Note : This method can be used for deducing the Layout Tag of A, C, D Matrices +template +constexpr +auto +stride_to_layout_tag_A() { + if constexpr (is_mn_major()) { // M major + return layout::ColumnMajor{}; + } + else { // K major + return layout::RowMajor{}; + } + + CUTE_GCC_UNREACHABLE; +} + +template +constexpr +auto +stride_to_layout_tag_B() { + if constexpr (is_mn_major()) { // N major + return layout::RowMajor{}; + } + else { // K major + return layout::ColumnMajor{}; + } + + CUTE_GCC_UNREACHABLE; +} + +// Inspects a TiledCopy and returns its alignment in terms of element count +template +constexpr int +get_alignment_count_from_gmem_tiled_copy() { + if constexpr (cute::is_void_v) { + return 1; + } + + // Account for ElementC = void kernels + else if constexpr (cute::is_void_v) { + return 0; + } + + else { + // For TMA tiled copies, we know the alignment has to be 128 bits + // if constexpr ( cute::is_base_of_v + // || cute::is_base_of_v + // || cute::is_base_of_v + // ) { + if (0) { + return 128 / sizeof_bits::value; + } + else { + // For non-TMA tiled copies, TiledCopy holds the alignment count directly in its TiledShape_MN + return GmemTiledCopy::NumValSrc; + } + } +} + +// Utilities to map Stride back on to their corresponding layout tags +template +struct StrideToLayoutTagA { + using type = decltype(detail::stride_to_layout_tag_A()); +}; + +template +struct StrideToLayoutTagB { + using type = decltype(detail::stride_to_layout_tag_B()); +}; + +// Maps to modes [N, N, L] +template +struct StrideToLayoutTagC : StrideToLayoutTagA { }; + +// Convenience aliases +template +using StrideToLayoutTagA_t = typename StrideToLayoutTagA::type; + +template +using StrideToLayoutTagB_t = typename StrideToLayoutTagB::type; + +template +using StrideToLayoutTagC_t = typename StrideToLayoutTagC::type; + +template +constexpr +bool +is_k_major() { + return ! is_mn_major(); +} + +template +constexpr bool +is_mn_major_A() { + return is_mn_major>(); +} + +template +constexpr bool +is_mn_major_B() { + return is_mn_major>(); +} + +template +constexpr bool +is_k_major_A() { + return is_k_major>(); +} + +template +constexpr bool +is_k_major_B() { + return is_k_major>(); +} + +/////////////////////////////////////////////////////////////////////////////// + +// The following two metafunctions are used to detect whether a `kernel::Gemm` or `kernel::GemmUniversal` +// is implementing the MCTLASS 3.x API or not, by checking if the problem shape type is aliased within or not. +template +struct IsCutlass3GemmKernel : cute::false_type { }; + +template +struct IsCutlass3GemmKernel> + : cute::true_type { }; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace detail + +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_ell_gemm.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_ell_gemm.h new file mode 100644 index 0000000..239c33b --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_ell_gemm.h @@ -0,0 +1,837 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Default kernel-level Blocked-Ell sparse gemm operators. + This operator combines threadblock-scoped ELL MMA + with the appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/gemm.h" +#include "mctlass/gemm/kernel/gemm_pipelined.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_simt.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + +#include "mctlass/gemm/kernel/ell_gemm.h" +#include "mctlass/gemm/threadblock/default_ell_mma.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Sparse matrix is A or not + bool IsASparse> +struct DefaultEllGemm; + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Sparse matrix is A or not + bool IsASparse +> +struct DefaultEllGemm { + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultEllMma< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount>::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::EllGemm; +}; +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Turing Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Sparse matrix is A or not + bool IsASparse +> +struct DefaultEllGemm< + ElementA, LayoutA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + ElementC, layout::RowMajor, + ElementAccumulator, + arch::OpClassTensorOp, + arch::Sm75, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + SplitKSerial, + Operator, + IsASparse +> { + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultEllMma< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + layout::RowMajor, + arch::OpClassTensorOp, + arch::Sm75, + ThreadblockShape, + WarpShape, + InstructionShape, + 2, + Operator + >::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + typename Mma::Operator, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::EllGemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Integer Matrix Multiply Interleaved layout +template < + /// Element type for A matrix operand + typename ElementA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Number of Interleaved k + int InterleavedK, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Sparse matrix is A or not + bool IsASparse> +struct DefaultEllGemm< + ElementA, layout::ColumnMajorInterleaved, kAlignmentA, + ElementB, layout::RowMajorInterleaved, kAlignmentB, ElementC, + layout::ColumnMajorInterleaved, int32_t, + arch::OpClassTensorOp, arch::Sm80, ThreadblockShape, WarpShape, + InstructionShape, EpilogueOutputOp, ThreadblockSwizzle, Stages, + SplitKSerial, Operator, IsASparse> { + using LayoutA = layout::ColumnMajorInterleaved; + using LayoutB = layout::RowMajorInterleaved; + using LayoutC = layout::ColumnMajorInterleaved; + + using ElementAccumulator = int32_t; + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultEllMma< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + ElementAccumulator, LayoutC, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, Operator, + true>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock:: + DefaultInterleavedEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + 64 / sizeof_bits::value, InterleavedK>::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::EllGemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Turing Integer Matrix Multiply Interleaved layout +template < + /// Element type for A matrix operand + typename ElementA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of Interleaved k + int InterleavedK, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Sparse matrix is A or not + bool IsASparse> +struct DefaultEllGemm, + kAlignmentA, ElementB, + layout::RowMajorInterleaved, kAlignmentB, + ElementC, layout::ColumnMajorInterleaved, + int32_t, arch::OpClassTensorOp, arch::Sm75, ThreadblockShape, + WarpShape, InstructionShape, EpilogueOutputOp, + ThreadblockSwizzle, 2, SplitKSerial, Operator, IsASparse> { + using LayoutA = layout::ColumnMajorInterleaved; + using LayoutB = layout::RowMajorInterleaved; + using LayoutC = layout::ColumnMajorInterleaved; + + using ElementAccumulator = int32_t; + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultEllMma< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, ElementAccumulator, LayoutC, + arch::OpClassTensorOp, arch::Sm75, ThreadblockShape, WarpShape, + InstructionShape, 2, Operator, true>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock:: + DefaultInterleavedEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + 64 / sizeof_bits::value, InterleavedK>::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::EllGemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + + +/// Partial specialization for Volta architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Sparse matrix is A or not + bool IsASparse +> +struct DefaultEllGemm< + ElementA, LayoutA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + ElementC, layout::RowMajor, + ElementAccumulator, + arch::OpClassTensorOp, + arch::Sm70, + ThreadblockShape, + WarpShape, + GemmShape<8, 8, 4>, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + SplitKSerial, + Operator, + IsASparse +> { + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultEllMma< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + layout::RowMajor, + arch::OpClassTensorOp, + arch::Sm70, + ThreadblockShape, + WarpShape, + GemmShape<8, 8, 4>, + 2, + Operator + >::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueVoltaTensorOp< + ThreadblockShape, + typename Mma::Operator, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::EllGemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for SIMT +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Sparse matrix is A or not + bool IsASparse + > +struct DefaultEllGemm< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementC, + layout::RowMajor, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + GemmShape<1, 1, 1>, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + SplitKSerial, + Operator, + IsASparse> { + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultEllMma< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + layout::RowMajor, + arch::OpClassSimt, + arch::Sm50, + ThreadblockShape, + WarpShape, + GemmShape<1, 1, 1>, + 2, + Operator>::ThreadblockMma; + + static int const kEpilogueElementsPerAccess = EpilogueOutputOp::kCount; + static_assert(kEpilogueElementsPerAccess == 1, "simt epilogue must operate on scalars"); + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + typename Mma::Operator, + EpilogueOutputOp, + kEpilogueElementsPerAccess + >::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::EllGemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages + int Stages, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Sparse matrix is A or not + bool IsASparse + > +struct DefaultEllGemm, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + IsASparse> { + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultEllMma< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + ElementAccumulator, layout::RowMajor, arch::OpClassSimt, arch::Sm80, + ThreadblockShape, WarpShape, GemmShape<1, 1, 1>, Stages, + Operator>::ThreadblockMma; + + static int const kEpilogueElementsPerAccess = EpilogueOutputOp::kCount; + static_assert(kEpilogueElementsPerAccess == 1, "simt epilogue must operate on scalars"); + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + typename Mma::Operator, + EpilogueOutputOp, + kEpilogueElementsPerAccess + >::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::EllGemm; +}; + +//////////////////////////////////////////////////////////////////////////////// +/// Partial specialization for SIMT DP4A + +template < + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Layout type for C matrix operand + typename LayoutC, + /// Element type for C and D matrix operands + typename ElementC, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Sparse matrix is A or not + bool IsASparse + > +struct DefaultEllGemm, + EpilogueOutputOp, ThreadblockSwizzle, 2, SplitKSerial, + Operator, IsASparse> { + using InstructionShape = GemmShape<1, 1, 4>; + using ElementA = int8_t; + using ElementB = int8_t; + + using OperatorClass = arch::OpClassSimt; + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultEllMma::ThreadblockMma; + + static int const kEpilogueElementsPerAccess = EpilogueOutputOp::kCount; + static_assert(kEpilogueElementsPerAccess == 1, "simt epilogue must operate on scalars"); + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + typename Mma::Operator, + EpilogueOutputOp, + kEpilogueElementsPerAccess + >::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::EllGemm; +}; + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +//////////////////////////////////////////////////////////////////////////////// +/// Partial specialization for Wmma Gemm Kernel +template < + ///< Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Sparse matrix is A or not + bool IsASparse + > +struct DefaultEllGemm< + ElementA, LayoutA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + ElementC, LayoutC, + ElementAccumulator, + arch::OpClassWmmaTensorOp, + ArchTag, + ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + IsASparse> { + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultEllMma< + ElementA, LayoutA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + ElementAccumulator, LayoutC, + arch::OpClassWmmaTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + Stages, + Operator>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueWmmaTensorOp< + ThreadblockShape, + typename Mma::Operator, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::EllGemm; +}; +//////////////////////////////////////////////////////////////////////////////// +#endif //MCTLASS_ARCH_WMMA_ENABLED + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm.h new file mode 100644 index 0000000..e663c59 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm.h @@ -0,0 +1,1111 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. Partial + specializations here choose 'device::GemmTransposed' to implement this functionality. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/gemm.h" +#include "mctlass/gemm/kernel/gemm_pipelined.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_simt.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +#include "mctlass/layout/permute.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// Gather operand A by using an index array + bool GatherA = false, + /// Gather operand B by using an index array + bool GatherB = false, + /// Scatter result D by using an index array + bool ScatterD = false, + /// Permute result D + typename PermuteDLayout = layout::NoPermute, + /// Permute operand A + typename PermuteALayout = layout::NoPermute, + /// Permute operand B + typename PermuteBLayout = layout::NoPermute, + /// + typename Enable = void +> +struct DefaultGemm; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Hopper Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Scatter result D by using an index array + bool ScatterD, + /// Permute result D + typename PermuteDLayout, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout +> +struct DefaultGemm { + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator, false, SharedMemoryClear, GatherA, GatherB, + PermuteALayout, PermuteBLayout>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount, ScatterD, PermuteDLayout>::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operand + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Scatter result D by using an index array + bool ScatterD, + /// Permute result D + typename PermuteDLayout, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout +> +struct DefaultGemm { + + static_assert((platform::is_same::value + || platform::is_same>::value), + "Epilogue in the kernel level must be row major"); + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + ElementAccumulator, LayoutC, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator, false, SharedMemoryClear, GatherA, GatherB, + PermuteALayout, PermuteBLayout>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using RegularEpilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount, ScatterD, PermuteDLayout>::Epilogue; + + using Affine2Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOpAffineRankN< + 2, ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount>::Epilogue; + + using Epilogue = typename platform::conditional::value, + RegularEpilogue, + Affine2Epilogue>::type; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Turing Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Scatter result D by using an index array + bool ScatterD, + /// Permute result D + typename PermuteDLayout, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout +> +struct DefaultGemm< + ElementA, LayoutA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + ElementC, layout::RowMajor, + ElementAccumulator, + arch::OpClassTensorOp, + arch::Sm75, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + SplitKSerial, + Operator, + SharedMemoryClear, + GatherA, + GatherB, + ScatterD, + PermuteDLayout, + PermuteALayout, + PermuteBLayout +> { + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + layout::RowMajor, + arch::OpClassTensorOp, + arch::Sm75, + ThreadblockShape, + WarpShape, + InstructionShape, + 2, + Operator, + false, + SharedMemoryClear, + GatherA, + GatherB, + PermuteALayout, + PermuteBLayout + >::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, + typename Mma::Operator, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount, + ScatterD, + PermuteDLayout + >::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Integer Matrix Multiply Interleaved layout +template < + /// Element type for A matrix operand + typename ElementA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Number of Interleaved k + int InterleavedK, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear> +struct DefaultGemm< + ElementA, layout::ColumnMajorInterleaved, kAlignmentA, + ElementB, layout::RowMajorInterleaved, kAlignmentB, ElementC, + layout::ColumnMajorInterleaved, int32_t, + arch::OpClassTensorOp, arch::Sm80, ThreadblockShape, WarpShape, + InstructionShape, EpilogueOutputOp, ThreadblockSwizzle, Stages, + SplitKSerial, Operator, SharedMemoryClear, false, false, false> { + + using LayoutA = layout::ColumnMajorInterleaved; + using LayoutB = layout::RowMajorInterleaved; + using LayoutC = layout::ColumnMajorInterleaved; + + using ElementAccumulator = int32_t; + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + ElementAccumulator, LayoutC, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, Operator, + true, SharedMemoryClear>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock:: + DefaultInterleavedEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + 64 / sizeof_bits::value, InterleavedK>::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Turing Integer Matrix Multiply Interleaved layout +template < + /// Element type for A matrix operand + typename ElementA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of Interleaved k + int InterleavedK, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear> +struct DefaultGemm, + kAlignmentA, ElementB, + layout::RowMajorInterleaved, kAlignmentB, + ElementC, layout::ColumnMajorInterleaved, + int32_t, arch::OpClassTensorOp, arch::Sm75, ThreadblockShape, + WarpShape, InstructionShape, EpilogueOutputOp, + ThreadblockSwizzle, 2, SplitKSerial, Operator, SharedMemoryClear, + false, false, false> { + + using LayoutA = layout::ColumnMajorInterleaved; + using LayoutB = layout::RowMajorInterleaved; + using LayoutC = layout::ColumnMajorInterleaved; + + using ElementAccumulator = int32_t; + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, ElementAccumulator, LayoutC, + arch::OpClassTensorOp, arch::Sm75, ThreadblockShape, WarpShape, + InstructionShape, 2, Operator, true>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock:: + DefaultInterleavedEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + 64 / sizeof_bits::value, InterleavedK>::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Volta architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Scatter result D by using an index array + bool ScatterD, + /// Permute result D + typename PermuteDLayout, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout +> +struct DefaultGemm< + ElementA, LayoutA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + ElementC, layout::RowMajor, + ElementAccumulator, + arch::OpClassTensorOp, + arch::Sm70, + ThreadblockShape, + WarpShape, + GemmShape<8, 8, 4>, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + SplitKSerial, + Operator, + SharedMemoryClear, + GatherA, + GatherB, + ScatterD, + PermuteDLayout, + PermuteALayout, + PermuteBLayout +> { + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + layout::RowMajor, + arch::OpClassTensorOp, + arch::Sm70, + ThreadblockShape, + WarpShape, + GemmShape<8, 8, 4>, + 2, + Operator, + false, + SharedMemoryClear, + GatherA, + GatherB, + PermuteALayout, + PermuteBLayout + >::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueVoltaTensorOp< + ThreadblockShape, + typename Mma::Operator, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount, + ScatterD, + PermuteDLayout + >::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for SIMT +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operand + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Scatter result D by using an index array + bool ScatterD, + /// Permute result D + typename PermuteDLayout, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout + > +struct DefaultGemm< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + arch::OpClassSimt, + ArchTag, + ThreadblockShape, + WarpShape, + GemmShape<1, 1, 1>, + EpilogueOutputOp, + ThreadblockSwizzle, + 2, + SplitKSerial, + Operator, + SharedMemoryClear, + GatherA, + GatherB, + ScatterD, + PermuteDLayout, + PermuteALayout, + PermuteBLayout, + typename platform::enable_if< ! platform::is_same::value >::type > { + + static_assert((platform::is_same::value + || platform::is_same>::value), + "Epilogue in the kernel level must be row major"); + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + LayoutC, + arch::OpClassSimt, + arch::Sm50, + ThreadblockShape, + WarpShape, + GemmShape<1, 1, 1>, + 2, + Operator, + false, + SharedMemoryClear, + GatherA, + GatherB, + PermuteALayout, + PermuteBLayout>::ThreadblockMma; + + static int const kEpilogueElementsPerAccess = EpilogueOutputOp::kCount; + static_assert(kEpilogueElementsPerAccess == 1, "simt epilogue must operate on scalars"); + + /// Define the epilogue + using RegularEpilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + typename Mma::Operator, + EpilogueOutputOp, + kEpilogueElementsPerAccess, + ScatterD, + PermuteDLayout + >::Epilogue; + + using Affine2Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueSimtAffineRankN< + 2, + ThreadblockShape, + typename Mma::Operator, + EpilogueOutputOp, + kEpilogueElementsPerAccess + >::Epilogue; + + using Epilogue = typename platform::conditional::value, + RegularEpilogue, + Affine2Epilogue>::type; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operand + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages + int Stages, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Scatter result D by using an index array + bool ScatterD, + /// Permute result D + typename PermuteDLayout, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout +> +struct DefaultGemm, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + SharedMemoryClear, + GatherA, + GatherB, + ScatterD, + PermuteDLayout, + PermuteALayout, + PermuteBLayout> { + + static_assert((platform::is_same::value + || platform::is_same>::value), + "Epilogue in the kernel level must be row major"); + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + ElementAccumulator, LayoutC, arch::OpClassSimt, arch::Sm80, + ThreadblockShape, WarpShape, GemmShape<1, 1, 1>, Stages, + Operator, false, SharedMemoryClear, GatherA, GatherB, + PermuteALayout, PermuteBLayout>::ThreadblockMma; + + static int const kEpilogueElementsPerAccess = EpilogueOutputOp::kCount; + static_assert(kEpilogueElementsPerAccess == 1, "simt epilogue must operate on scalars"); + + /// Define the epilogue + using RegularEpilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + typename Mma::Operator, + EpilogueOutputOp, + kEpilogueElementsPerAccess, + ScatterD, + PermuteDLayout + >::Epilogue; + + using Affine2Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueSimtAffineRankN< + 2, + ThreadblockShape, + typename Mma::Operator, + EpilogueOutputOp, + kEpilogueElementsPerAccess + >::Epilogue; + + using Epilogue = typename platform::conditional::value, + RegularEpilogue, + Affine2Epilogue>::type; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; + +//////////////////////////////////////////////////////////////////////////////// +/// Partial specialization for SIMT DP4A + +template < + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Layout type for C matrix operand + typename LayoutC, + /// Element type for C and D matrix operands + typename ElementC, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear +> +struct DefaultGemm, + EpilogueOutputOp, ThreadblockSwizzle, 2, SplitKSerial, + Operator, SharedMemoryClear, false, false, false, + layout::NoPermute, layout::NoPermute> { + using InstructionShape = GemmShape<1, 1, 4>; + using ElementA = int8_t; + using ElementB = int8_t; + + using OperatorClass = arch::OpClassSimt; + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + LayoutC, + arch::OpClassSimt, + arch::Sm50, + ThreadblockShape, + WarpShape, + InstructionShape, + 2, + Operator + >::ThreadblockMma; + + static int const kEpilogueElementsPerAccess = EpilogueOutputOp::kCount; + static_assert(kEpilogueElementsPerAccess == 1, "simt epilogue must operate on scalars"); + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + typename Mma::Operator, + EpilogueOutputOp, + kEpilogueElementsPerAccess + >::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +//////////////////////////////////////////////////////////////////////////////// +/// Partial specialization for Wmma Gemm Kernel +template < + ///< Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear +> +struct DefaultGemm< + ElementA, LayoutA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + ElementC, LayoutC, + ElementAccumulator, + arch::OpClassWmmaTensorOp, + ArchTag, + ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + SharedMemoryClear, + false, + false, + false, + layout::NoPermute, + layout::NoPermute +> { + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, LayoutA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + ElementAccumulator, LayoutC, + arch::OpClassWmmaTensorOp, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + Stages, + Operator>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueWmmaTensorOp< + ThreadblockShape, + typename Mma::Operator, + kPartitionsK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; +//////////////////////////////////////////////////////////////////////////////// + +#endif //MCTLASS_ARCH_WMMA_ENABLED + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_complex.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_complex.h new file mode 100644 index 0000000..700e420 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_complex.h @@ -0,0 +1,404 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. Partial + specializations here choose 'device::GemmTransposed' to implement this functionality. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/gemm.h" +#include "mctlass/gemm/kernel/gemm_pipelined.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/default_multistage_mma_complex.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/epilogue/threadblock/default_epilogue_complex_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_simt.h" + +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Multiply-add operator + // (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial +> +struct DefaultGemmComplex; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Hopper Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Multiply-add operator + // (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial + > +struct DefaultGemmComplex< + ElementA, LayoutA, ElementB, LayoutB, ElementC, + layout::RowMajor, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm90, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, TransformA, TransformB, Operator, SplitKSerial> { + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplex< + ElementA, LayoutA, ElementB, LayoutB, ElementAccumulator, + layout::RowMajor, arch::OpClassTensorOp, arch::Sm90, ThreadblockShape, + WarpShape, InstructionShape, Stages, TransformA, TransformB, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp< + ThreadblockShape, typename Mma::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator>::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Multiply-add operator + // (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial + > +struct DefaultGemmComplex< + ElementA, LayoutA, ElementB, LayoutB, ElementC, + layout::RowMajor, ElementAccumulator, arch::OpClassSimt, + arch::Sm50, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, TransformA, TransformB, Operator, SplitKSerial> { + + /// Define the threadblock-scoped matrix multiply-accumulate + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, + WarpShape, + InstructionShape, + ElementA, LayoutA, + ElementB, LayoutB, + ElementAccumulator, layout::RowMajor, + arch::OpClassSimt, + Stages, + Operator, + false, + mctlass::arch::CacheOperation::Global, + mctlass::arch::CacheOperation::Global, + TransformA, + TransformB + >; + + // Define iterators over tiles from the A operand + using IteratorA = + mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, + typename MmaCore::IteratorThreadMapA>; + + // Define iterators over tiles from the B operand + using IteratorB = + mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, + typename MmaCore::IteratorThreadMapB>; + + // Define the threadblock-scoped pipelined matrix multiply + using Mma = mctlass::gemm::threadblock::MmaPipelined< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, + layout::RowMajor, typename MmaCore::MmaPolicy>; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + typename Mma::Operator, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Multiply-add operator + // (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial + > +struct DefaultGemmComplex< + ElementA, LayoutA, ElementB, LayoutB, ElementC, + layout::RowMajor, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm80, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, TransformA, TransformB, Operator, SplitKSerial> { + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplex< + ElementA, LayoutA, ElementB, LayoutB, ElementAccumulator, + layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, ThreadblockShape, + WarpShape, InstructionShape, Stages, TransformA, TransformB, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp< + ThreadblockShape, typename Mma::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator>::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Multiply-add operator + // (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial + > +struct DefaultGemmComplex< + ElementA, LayoutA, ElementB, LayoutB, ElementC, + layout::RowMajor, ElementAccumulator, arch::OpClassSimt, + arch::Sm80, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, TransformA, TransformB, Operator, SplitKSerial> { + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplex< + ElementA, LayoutA, ElementB, LayoutB, ElementAccumulator, + layout::RowMajor, arch::OpClassSimt, arch::Sm80, ThreadblockShape, + WarpShape, InstructionShape, Stages, TransformA, TransformB, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueSimt< + ThreadblockShape, + typename Mma::Operator, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::Gemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_grouped.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_grouped.h new file mode 100644 index 0000000..4a92d16 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_grouped.h @@ -0,0 +1,384 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. Partial + specializations here choose 'device::GemmTransposed' to implement this functionality. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/complex.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" + +#include "mctlass/gemm/kernel/gemm_grouped.h" +#include "mctlass/gemm/kernel/gemm_transpose_operands.h" +#include "mctlass/gemm/kernel/default_gemm.h" +#include "mctlass/gemm/kernel/default_gemm_complex.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" + +#include "mctlass/layout/permute.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Whether the schedule of problems to visit has been precomputed + GroupScheduleMode GroupScheduleMode_ = GroupScheduleMode::kDeviceOnly, + /// Operation performed by GEMM + typename Operator = typename device::DefaultGemmConfiguration< + OperatorClass, ArchTag, ElementA_, ElementB_, ElementC_, + ElementAccumulator>::Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// Permute result D + typename PermuteDLayout = layout::NoPermute, + /// + typename Enable = void + > +struct DefaultGemmGrouped; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Real-valued GEMM kernels +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Whether the schedule of problems to visit has been precomputed + GroupScheduleMode GroupScheduleMode_, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + /// Permute result D + typename PermuteDLayout +> +struct DefaultGemmGrouped< + ElementA, + LayoutA, + ComplexTransform::kNone, // transform A + kAlignmentA, + ElementB, + LayoutB, + ComplexTransform::kNone, // transform B + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + GroupScheduleMode_, + Operator, + SharedMemoryClear, + PermuteDLayout, + typename platform::enable_if< ! mctlass::is_complex::value>::type +> { + + // If true, we must construct a 'transposed-and-exchanged' Mma operator. + static bool const kInternalTranspose = platform::is_same::value; + + using MapArguments = kernel::detail::MapArguments< + ElementA, + LayoutA, + ComplexTransform::kNone, + kAlignmentA, + ElementB, + LayoutB, + ComplexTransform::kNone, + kAlignmentB, + LayoutC, + kInternalTranspose + >; + + // Define the default GEMM kernel + using DefaultGemmKernel = typename kernel::DefaultGemm< + typename MapArguments::ElementA, + typename MapArguments::LayoutA, + MapArguments::kAlignmentA, + typename MapArguments::ElementB, + typename MapArguments::LayoutB, + MapArguments::kAlignmentB, + ElementC, + typename MapArguments::LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + true, + Operator, + SharedMemoryClear, + false, /*GatherA*/ + false, /*GatherB*/ + false, /*ScatterD*/ + PermuteDLayout + >::GemmKernel; + + /// Define the kernel in terms of the default kernel + using GemmKernel = kernel::GemmGrouped< + typename DefaultGemmKernel::Mma, + typename DefaultGemmKernel::Epilogue, + ThreadblockSwizzle, + GroupScheduleMode_, + kInternalTranspose + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Complex-valued GEMM kernels +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Whether the schedule of problems to visit has been precomputed + GroupScheduleMode GroupScheduleMode_, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear + > +struct DefaultGemmGrouped< + ElementA, + LayoutA, + TransformA, + kAlignmentA, + ElementB, + LayoutB, + TransformB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + GroupScheduleMode_, + Operator, + SharedMemoryClear, + layout::NoPermute, /*PermuteDLayout*/ + typename platform::enable_if::value>::type +> { + + // If true, we must construct a 'transposed-and-exchanged' Mma operator. + static bool const kInternalTranspose = platform::is_same::value; + + using MapArguments = kernel::detail::MapArguments< + ElementA, + LayoutA, + TransformA, + kAlignmentA, + ElementB, + LayoutB, + TransformB, + kAlignmentB, + LayoutC, + kInternalTranspose + >; + + using DefaultGemmKernel = typename kernel::DefaultGemmComplex< + typename MapArguments::ElementA, + typename MapArguments::LayoutA, + typename MapArguments::ElementB, + typename MapArguments::LayoutB, + ElementC, + typename MapArguments::LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MapArguments::kTransformA, + MapArguments::kTransformB, + Operator, + false + >::GemmKernel; + + /// Define the kernel in terms of the default kernel + using GemmKernel = kernel::GemmGrouped< + typename DefaultGemmKernel::Mma, + typename DefaultGemmKernel::Epilogue, + ThreadblockSwizzle, + GroupScheduleMode_, + kInternalTranspose + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_grouped_softmax_mainloop_fusion.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_grouped_softmax_mainloop_fusion.h new file mode 100644 index 0000000..e0accf6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_grouped_softmax_mainloop_fusion.h @@ -0,0 +1,164 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level softmax-grouped-GEMM +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/complex.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" + +#include "mctlass/gemm/kernel/gemm_grouped_softmax_mainloop_fusion.h" +#include "mctlass/gemm/kernel/gemm_transpose_operands.h" +#include "mctlass/gemm/kernel/default_gemm.h" +#include "mctlass/gemm/kernel/default_gemm_complex.h" +#include "mctlass/gemm/device/default_gemm_configuration.h" +#include "mctlass/gemm/threadblock/default_mma_softmax_mainloop_fusion.h" + +#include "mctlass/layout/permute.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for Scale/Bias vectors + typename ElementScaleBias_, + /// Layout type for Scale/Bias vectors + typename LayoutScaleBias_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Whether the schedule of problems to visit has been precomputed + GroupScheduleMode GroupScheduleMode_ = GroupScheduleMode::kDeviceOnly, + /// Operation performed by GEMM + typename Operator = typename device::DefaultGemmConfiguration< + OperatorClass, ArchTag, ElementA_, ElementB_, ElementC_, + ElementAccumulator>::Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone + > +struct DefaultGemmGroupedSoftmaxMainloopFusion { + // If true, we must construct a 'transposed-and-exchanged' Mma operator. + static bool const kInternalTranspose = platform::is_same::value; + + using MapArguments = kernel::detail::MapArguments< + ElementA_, + LayoutA_, + ComplexTransform::kNone, + kAlignmentA, + ElementB_, + LayoutB_, + ComplexTransform::kNone, + kAlignmentB, + LayoutC_, + kInternalTranspose + >; + +private: + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMmaSoftmaxMainloopFusion< + typename MapArguments::ElementA, typename MapArguments::LayoutA, MapArguments::kAlignmentA, + typename MapArguments::ElementB, typename MapArguments::LayoutB, MapArguments::kAlignmentB, + ElementScaleBias_, LayoutScaleBias_, ElementAccumulator, layout::RowMajor, OperatorClass, ArchTag, + ThreadblockShape, WarpShape, InstructionShape, Stages, kInternalTranspose, + Operator, false, SharedMemoryClear>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount>::Epilogue; + +public: + using GemmKernel = kernel::GemmGroupedSoftmaxMainloopFusion< + Mma, + Epilogue, + ThreadblockSwizzle, + GroupScheduleMode_, + kInternalTranspose + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_layernorm_mainloop_fusion.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_layernorm_mainloop_fusion.h new file mode 100644 index 0000000..f699f8f --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_layernorm_mainloop_fusion.h @@ -0,0 +1,137 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. Partial + specializations here choose 'device::GemmTransposed' to implement this functionality. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/gemm_layernorm_mainloop_fusion.h" +#include "mctlass/gemm/threadblock/default_mma_layernorm_mainloop_fusion.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for Scale/Bias vectors + typename ElementScaleBias, + /// Layout type for Scale/Bias vectors + typename LayoutScaleBias, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone> +struct DefaultGemmLayernormMainloopFusion { + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMmaLayernormMainloopFusion< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + ElementScaleBias, LayoutScaleBias, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator, false, SharedMemoryClear>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount>::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::GemmLayernormMainloopFusion; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_planar_complex_universal.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_planar_complex_universal.h new file mode 100644 index 0000000..b10e788 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_planar_complex_universal.h @@ -0,0 +1,352 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. Partial + specializations here choose 'device::GemmTransposed' to implement this functionality. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/complex.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" + +#include "mctlass/gemm/kernel/gemm_planar_complex.h" +#include "mctlass/gemm/kernel/gemm_planar_complex_array.h" +#include "mctlass/gemm/kernel/default_gemm.h" +#include "mctlass/gemm/kernel/default_gemm_complex.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_planar_complex.h" +#include "mctlass/gemm/threadblock/default_mma_planar_complex_pipelined.h" +#include "mctlass/gemm/threadblock/default_mma_planar_complex_multistage.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Math operation performed by GEMM (e.g. arch::OpMultiplyAdd) + typename Operator, + /// Conditional enabling to switch between stages + typename Enable = void + > +struct DefaultGemmPlanarComplexUniversal; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for pipelined mainloop +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator + > +struct DefaultGemmPlanarComplexUniversal< + ElementA, + LayoutA, + TransformA, + kAlignmentA, + ElementB, + LayoutB, + TransformB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + Operator, + typename platform::enable_if<(Stages <= 2)>::type +> { + + /// Define planar complex valued variants instead + using Mma = typename gemm::threadblock::DefaultMmaPlanarComplexPipelined< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + LayoutC, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + Stages, + TransformA, + TransformB, + Operator + >::ThreadblockMma; + + /// Planar complex epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpiloguePlanarComplex< + ThreadblockShape, + typename Mma::Policy::Operator, + OperatorClass, + ArchTag, + ThreadblockShape::kK / WarpShape::kK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + /// Define the kernel in terms of the default kernel + using GemmKernel = kernel::GemmPlanarComplex< + Mma, + Epilogue, + ThreadblockSwizzle + >; + + // Array variant + using GemmArrayKernel = kernel::GemmPlanarComplexArray< + Mma, + Epilogue, + ThreadblockSwizzle + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for multiple pipeline stages. +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator + > +struct DefaultGemmPlanarComplexUniversal< + ElementA, + LayoutA, + TransformA, + kAlignmentA, + ElementB, + LayoutB, + TransformB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + Operator, + typename platform::enable_if<(Stages > 2)>::type +> { + + /// Define planar complex valued variants instead + using Mma = typename gemm::threadblock::DefaultMmaPlanarComplexMultistage< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementAccumulator, + LayoutC, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + Stages, + TransformA, + TransformB, + Operator + >::ThreadblockMma; + + /// Planar complex epilogue + using Epilogue = typename epilogue::threadblock::DefaultEpiloguePlanarComplex< + ThreadblockShape, + typename Mma::Policy::Operator, + OperatorClass, + ArchTag, + ThreadblockShape::kK / WarpShape::kK, + EpilogueOutputOp, + EpilogueOutputOp::kCount + >::Epilogue; + + /// Define the kernel in terms of the default kernel + using GemmKernel = kernel::GemmPlanarComplex< + Mma, + Epilogue, + ThreadblockSwizzle + >; + + // Array variant + using GemmArrayKernel = kernel::GemmPlanarComplexArray< + Mma, + Epilogue, + ThreadblockSwizzle + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_sparse.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_sparse.h new file mode 100644 index 0000000..c15e835 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_sparse.h @@ -0,0 +1,191 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief + Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. Partial + specializations here choose 'device::GemmTransposed' to implement this functionality. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/gemm.h" +#include "mctlass/gemm/kernel/sparse_gemm.h" +#include "mctlass/gemm/kernel/gemm_pipelined.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma_core_sparse_sm80.h" +#include "mctlass/gemm/threadblock/default_sparse_mma.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_simt.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator> +struct DefaultSparseGemm; + +//////////////////////////////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator> +struct DefaultSparseGemm { + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultSparseMma< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount>::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::SparseGemm; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_sparse_row_broadcast.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_sparse_row_broadcast.h new file mode 100644 index 0000000..3628ded --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_sparse_row_broadcast.h @@ -0,0 +1,191 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief + Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. Partial + specializations here choose 'device::GemmTransposed' to implement this functionality. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/gemm.h" +#include "mctlass/gemm/kernel/sparse_gemm_row_broadcast.h" +#include "mctlass/gemm/kernel/gemm_pipelined.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma_core_sparse_sm80.h" +#include "mctlass/gemm/threadblock/default_sparse_mma.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op_row_broadcast.h" +#include "mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_simt.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator> +struct DefaultSparseGemmRowBroadcast; + +//////////////////////////////////////////////////////////////////////////////// +/////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator> +struct DefaultSparseGemmRowBroadcast { + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultSparseMma< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOpRowBroadcast< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount>::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::SparseGemmRowBroadcast; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_splitk_parallel.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_splitk_parallel.h new file mode 100644 index 0000000..043bce0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_splitk_parallel.h @@ -0,0 +1,136 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. Partial + specializations here choose 'device::GemmTransposed' to implement this functionality. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/kernel/default_gemm.h" +#include "mctlass/gemm/kernel/gemm_splitk_parallel.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator +> +struct DefaultGemmSplitKParallel { + + /// Define the threadblock-scoped matrix multiply-accumulate using the basic GEMM's + /// mainloop. + using Default = DefaultGemm< + ElementA_, + LayoutA_, + kAlignmentA, + ElementB_, + LayoutB_, + kAlignmentB, + ElementAccumulator, + LayoutC_, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + false, + Operator + >; + + /// Define the matrix multiply operator + using Mma = typename Default::Mma; + + /// Define the epilogue + using Epilogue = typename Default::Epilogue; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::GemmSplitKParallel; +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +/////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_streamk_with_broadcast.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_streamk_with_broadcast.h new file mode 100644 index 0000000..2a946fb --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_streamk_with_broadcast.h @@ -0,0 +1,146 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Defines a Stream-K GEMM that can broadcast a bias vector in the epilogue. + Similar structure to DefaultGemmWithBroadcast, but uses its own epilogue + (DefaultStreamkEpilogueWithBroadcastTensorOp) and its own GEMM kernel + (GemmStreamkWithFusedEpilogue). + +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/gemm/kernel/gemm_streamk_with_fused_epilogue.h" +#include "mctlass/gemm/kernel/default_gemm_universal.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_with_broadcast.h" +#include "mctlass/epilogue/threadblock/epilogue_with_broadcast.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator - must satisfy concept of 'EpilogueWithBroadcastOp' + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// + typename Enable = void +> +struct DefaultGemmStreamkWithBroadcast { + + using GemmBase = typename DefaultGemmUniversal< + ElementA_, LayoutA_, TransformA, kAlignmentA, + ElementB_, LayoutB_, TransformB, kAlignmentB, + ElementC_, LayoutC_, ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + Operator + >::GemmKernel; + + // Replace epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultStreamkEpilogueWithBroadcastTensorOp< + typename GemmBase::Epilogue::Shape, + typename GemmBase::Epilogue::WarpMmaOperator, + GemmBase::Epilogue::kPartitionsK, + ElementC_, + typename EpilogueOutputOp::ElementT, + typename EpilogueOutputOp::ElementVector, + EpilogueOutputOp, + GemmBase::Epilogue::kElementsPerAccess + >::Epilogue; + + // Compose the GEMM kernel + using GemmKernel = GemmStreamkWithFusedEpilogue< + typename GemmBase::Mma, + Epilogue, + ThreadblockSwizzle + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_universal.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_universal.h new file mode 100644 index 0000000..f79fc9a --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_universal.h @@ -0,0 +1,396 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. Partial + specializations here choose 'device::GemmTransposed' to implement this functionality. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/complex.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" + +#include "mctlass/gemm/kernel/gemm_universal.h" +#include "mctlass/gemm/kernel/gemm_universal_streamk.h" +#include "mctlass/gemm/kernel/default_gemm.h" +#include "mctlass/gemm/kernel/default_gemm_complex.h" + +#include "mctlass/layout/permute.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// Gather operand A by using an index array + bool GatherA = false, + /// Gather operand B by using an index array + bool GatherB = false, + /// Scatter result D by using an index array + bool ScatterD = false, + /// Permute result D + typename PermuteDLayout = layout::NoPermute, + /// Permute operand A + typename PermuteALayout_ = layout::NoPermute, + /// Permute operand B + typename PermuteBLayout_ = layout::NoPermute, + /// + typename Enable = void + > +struct DefaultGemmUniversal; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Real-valued GEMM kernels +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Scatter result D by using an index array + bool ScatterD, + /// Permute result D + typename PermuteDLayout, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout +> +struct DefaultGemmUniversal< + ElementA, + LayoutA, + ComplexTransform::kNone, // transform A + kAlignmentA, + ElementB, + LayoutB, + ComplexTransform::kNone, // transform B + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + Operator, + SharedMemoryClear, + GatherA, + GatherB, + ScatterD, + PermuteDLayout, + PermuteALayout, + PermuteBLayout, + typename platform::enable_if< ! mctlass::is_complex::value>::type +> { + + using DefaultGemmKernel = typename kernel::DefaultGemm< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + true, + Operator, + SharedMemoryClear, + GatherA, + GatherB, + ScatterD, + PermuteDLayout, + PermuteALayout, + PermuteBLayout + >::GemmKernel; + + /// Universal kernel without StreamkFeature member type + template + class SelectBase : + public kernel::GemmUniversal< + typename DefaultGemmKernel::Mma, + typename DefaultGemmKernel::Epilogue, + SwizzleT> + {}; + + /// Universal kernel with StreamkFeature member type + template + class SelectBase : + public kernel::GemmUniversalStreamk< + typename DefaultGemmKernel::Mma, + typename DefaultGemmKernel::Epilogue, + SwizzleT> + {}; + + /// Select kernel by ThreadblockSwizzle's support for StreamkFeature + using GemmKernel = SelectBase; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Complex-valued GEMM kernels +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear + > +struct DefaultGemmUniversal< + ElementA, + LayoutA, + TransformA, + kAlignmentA, + ElementB, + LayoutB, + TransformB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + Operator, + SharedMemoryClear, + false, + false, + false, + layout::NoPermute, + layout::NoPermute, + layout::NoPermute, + typename platform::enable_if::value>::type +> { + + using DefaultGemmKernel = typename kernel::DefaultGemmComplex< + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + TransformA, + TransformB, + Operator, + false + >::GemmKernel; + + /// Universal kernel without StreamkFeature member type + template + class SelectBase : + public kernel::GemmUniversal< + typename DefaultGemmKernel::Mma, + typename DefaultGemmKernel::Epilogue, + SwizzleT> + {}; + + /// Universal kernel with StreamkFeature member type + template + class SelectBase : + public kernel::GemmUniversalStreamk< + typename DefaultGemmKernel::Mma, + typename DefaultGemmKernel::Epilogue, + SwizzleT> + {}; + + /// Select kernel by ThreadblockSwizzle's support for StreamkFeature + using GemmKernel = SelectBase; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_with_broadcast.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_with_broadcast.h new file mode 100644 index 0000000..e1101c2 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_with_broadcast.h @@ -0,0 +1,243 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Defines a GEMM with Reduction based on an existing UniversalGemm kernel. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/gemm/kernel/gemm_with_fused_epilogue.h" +#include "mctlass/gemm/kernel/default_gemm_universal.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_with_broadcast.h" +#include "mctlass/epilogue/threadblock/epilogue_with_broadcast.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator - must satisfy concept of 'EpilogueWithBroadcastOp' + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// + typename Enable = void +> +struct DefaultGemmWithBroadcast { + + using GemmBase = typename DefaultGemmUniversal< + ElementA_, LayoutA_, TransformA, kAlignmentA, + ElementB_, LayoutB_, TransformB, kAlignmentB, + ElementC_, LayoutC_, ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + Operator + >::GemmKernel; + + // Define epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueWithBroadcastTensorOp< + typename GemmBase::Epilogue::Shape, + typename GemmBase::Epilogue::WarpMmaOperator, + GemmBase::Epilogue::kPartitionsK, + ElementC_, + typename EpilogueOutputOp::ElementT, + typename EpilogueOutputOp::ElementVector, + EpilogueOutputOp, + GemmBase::Epilogue::kElementsPerAccess + >::Epilogue; + + // Compose the GEMM kernel + using GemmKernel = GemmWithFusedEpilogue< + typename GemmBase::Mma, + Epilogue, + ThreadblockSwizzle + >; +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: ArchTag = mctlass::arch::Sm70 +/// +/// +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator - must satisfy concept of 'EpilogueWithBroadcastOp' + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// + typename Enable +> +struct DefaultGemmWithBroadcast< + ElementA_, LayoutA_, TransformA, kAlignmentA, + ElementB_, LayoutB_, TransformB, kAlignmentB, + ElementC_, LayoutC_, + ElementAccumulator, + OperatorClass, + mctlass::arch::Sm70, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + Operator, + Enable + > { + + using GemmBase = typename DefaultGemmUniversal< + ElementA_, LayoutA_, TransformA, kAlignmentA, + ElementB_, LayoutB_, TransformB, kAlignmentB, + ElementC_, LayoutC_, ElementAccumulator, + OperatorClass, + mctlass::arch::Sm70, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + Operator + >::GemmKernel; + + // Define epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueWithBroadcastVoltaTensorOp< + typename GemmBase::Epilogue::Shape, + typename GemmBase::Epilogue::WarpMmaOperator, + GemmBase::Epilogue::kPartitionsK, + ElementC_, + typename EpilogueOutputOp::ElementT, + typename EpilogueOutputOp::ElementVector, + EpilogueOutputOp, + GemmBase::Epilogue::kElementsPerAccess + >::Epilogue; + + // Compose the GEMM kernel + using GemmKernel = GemmWithFusedEpilogue< + typename GemmBase::Mma, + Epilogue, + ThreadblockSwizzle + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_with_k_reduction.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_with_k_reduction.h new file mode 100644 index 0000000..25f26f9 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_with_k_reduction.h @@ -0,0 +1,150 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. Partial + specializations here choose 'device::GemmTransposed' to implement this functionality. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/gemm_with_k_reduction.h" +#include "mctlass/gemm/threadblock/default_mma_with_reduction.h" +#include "mctlass/gemm/threadblock/default_mma_core_with_reduction.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "mctlass/epilogue/threadblock/epilogue_gemm_k_reduction.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Reduce A or B along the K dimension + bool ReduceKForA_, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// + typename Enable = void> +struct DefaultGemmWithKReduction { + + static const bool kReduceKForA = (platform::is_same::value) ? ReduceKForA_ : !ReduceKForA_; + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMmaWithReduction< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, kReduceKForA, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator, false, SharedMemoryClear>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount>::Epilogue; + + /// Define the epilogue of the reduction vector + using EpilogueGemmKReduction = + typename mctlass::epilogue::threadblock::EpilogueGemmKReduction< + ElementAccumulator, ElementC, ThreadblockShape, typename Mma::Operator, kReduceKForA>; + + /// Define the kernel-level GEMM operator. + using GemmKernel = kernel::GemmWithKReduction; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_with_reduction.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_with_reduction.h new file mode 100644 index 0000000..5c19bbb --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemm_with_reduction.h @@ -0,0 +1,246 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Defines a GEMM with Reduction based on an existing UniversalGemm kernel. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/gemm/kernel/gemm_with_fused_epilogue.h" +#include "mctlass/gemm/kernel/default_gemm_universal.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_with_reduction.h" +#include "mctlass/epilogue/threadblock/epilogue_with_reduction.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Epilogue reduction operator + typename EpilogueReductionOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// + typename Enable = void +> +struct DefaultGemmWithReduction { + + using GemmBase = typename DefaultGemmUniversal< + ElementA_, LayoutA_, TransformA, kAlignmentA, + ElementB_, LayoutB_, TransformB, kAlignmentB, + ElementC_, LayoutC_, ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + Operator, + SharedMemoryClearOption::kClearLastStage + >::GemmKernel; + + // Define epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueWithReductionTensorOp< + typename GemmBase::Epilogue::Shape, + typename GemmBase::Epilogue::WarpMmaOperator, + GemmBase::Epilogue::kPartitionsK, + ElementC_, + EpilogueOutputOp, + EpilogueReductionOp, + GemmBase::Epilogue::kElementsPerAccess + >::Epilogue; + + // Compose the GEMM kernel + using GemmKernel = GemmWithFusedEpilogue< + typename GemmBase::Mma, + Epilogue, + ThreadblockSwizzle + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: ArchTag = mctlass::arch::Sm70 +/// +/// +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Epilogue reduction operator + typename EpilogueReductionOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// + typename Enable +> +struct DefaultGemmWithReduction< + ElementA_, LayoutA_, TransformA, kAlignmentA, + ElementB_, LayoutB_, TransformB, kAlignmentB, + ElementC_, LayoutC_, + ElementAccumulator, + OperatorClass, + mctlass::arch::Sm70, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + EpilogueReductionOp, + ThreadblockSwizzle, + Stages, + Operator, + Enable + > { + + using GemmBase = typename DefaultGemmUniversal< + ElementA_, LayoutA_, TransformA, kAlignmentA, + ElementB_, LayoutB_, TransformB, kAlignmentB, + ElementC_, LayoutC_, ElementAccumulator, + OperatorClass, + mctlass::arch::Sm70, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + Operator + >::GemmKernel; + + // Define epilogue + using Epilogue = typename mctlass::epilogue::threadblock::DefaultEpilogueWithReductionVoltaTensorOp< + typename GemmBase::Epilogue::Shape, + typename GemmBase::Epilogue::WarpMmaOperator, + GemmBase::Epilogue::kPartitionsK, + ElementC_, + EpilogueOutputOp, + EpilogueReductionOp, + GemmBase::Epilogue::kElementsPerAccess + >::Epilogue; + + // Compose the GEMM kernel + using GemmKernel = GemmWithFusedEpilogue< + typename GemmBase::Mma, + Epilogue, + ThreadblockSwizzle + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_gemv.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemv.h new file mode 100644 index 0000000..40d09a6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_gemv.h @@ -0,0 +1,132 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include "mctlass/gemm/threadblock/gemv.h" +#include "mctlass/gemm/threadblock/default_gemv_core.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the ThreadBlock tile - concept: gemm::GemmShape<> + typename ThreadBlockShape_, + /// Size of the per-thread shape - concept: gemm::GemmShape<> + typename ThreadShape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C/D matrix + typename ElementCD_, + /// Layout of C/D matrix (concept: MatrixLayout) + typename LayoutCD_, + /// Data type of the accumulator + typename ElementAccumulator_ = ElementCD_> +struct DefaultGemv { + + /// Shape of Threadblock-level matrix operation (concept: GemmShape) + using ThreadBlockShape = ThreadBlockShape_; + + /// Shape of warp-level matrix operation (concept: GemmShape) + using ThreadShape = ThreadShape_; + + /// Data type of multiplicand A + using ElementA = ElementA_; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of multiplicand B + using ElementB = ElementB_; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of accumulators + using ElementAccumulator = ElementAccumulator_; + + /// Data type of accumulators (same as C/D) + using LayoutAccumulator = LayoutCD_; + + /// Data type of input/output matrix C/D + using ElementCD = ElementCD_; + + /// Layout of input/output matrix C/D + using LayoutCD = LayoutCD_; + + // Define the core components + using Core = typename mctlass::gemm::threadblock::DefaultGemvCore< + ThreadBlockShape, ThreadShape, ElementA, LayoutA, ElementB, LayoutB, + ElementAccumulator, LayoutAccumulator>; + + // Define the threadblock-scoped gemv + using ThreadBlockGemv = mctlass::gemm::threadblock::Gemv; + + // Iterator for multiplicand A + using IteratorA = typename ThreadBlockGemv::IteratorA; + + // Iterator for multiplicand B + using IteratorB = typename ThreadBlockGemv::IteratorB; + + /// Policy for the iterator that reads/writes C/D + using IteratorPolicyCD = typename platform::conditional< + platform::is_same::value, + mctlass::transform::PitchLinearTilePolicyStripminedThreadContiguous< + layout::PitchLinearShape, Core::kThreadsPerN, ThreadShape::kN>, + mctlass::transform::PitchLinearTilePolicyStripminedThreadStrided< + layout::PitchLinearShape, Core::kThreadsPerN, ThreadShape::kM>>::type; + + /// Iterator that reads/writes C/D + using IteratorCD = mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, ElementCD, LayoutCD, 0, IteratorPolicyCD>; + + /// Fragment storage for C/D + using FragmentCD = typename IteratorCD::Fragment; + + // Define the threadblock swizzle + using ThreadBlockSwizzle = mctlass::gemm::threadblock::GemvBatchedStridedThreadblockDefaultSwizzle; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k.h new file mode 100644 index 0000000..52dad99 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k.h @@ -0,0 +1,285 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level Rank2K definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/rank_2k_universal.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op_blas3.h" +#include "mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_simt.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Blas3 computation mode + BlasMode BlasMode_ = BlasMode::kSymmetric> +struct DefaultRank2K; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Hopper Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator> +struct DefaultRank2K< + ElementA, LayoutA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + ElementC,layout::RowMajor, FillModeC, + ElementAccumulator, arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, SplitKSerial, + Operator> { + /// Define the threadblock-scoped matrix multiply-accumulate (A x BT) + using Mma1 = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, LayoutA, + kAlignmentA, + ElementB, typename layout::LayoutTranspose::type, + kAlignmentB, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator>::ThreadblockMma; + + /// Define the threadblock-scoped matrix multiply-accumulate (B x AT) + using Mma2 = typename mctlass::gemm::threadblock::DefaultMma< + ElementB, LayoutB, + kAlignmentB, + ElementA, typename layout::LayoutTranspose::type, + kAlignmentA, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOpBlas3< + ThreadblockShape, typename Mma1::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount, BlasMode::kSymmetric>::Epilogue; + + /// Define the kernel-level Rank2K operator. + using Rank2Kkernel = kernel::Rank2KUniversal; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator> +struct DefaultRank2K< + ElementA, LayoutA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + ElementC,layout::RowMajor, FillModeC, + ElementAccumulator, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, SplitKSerial, + Operator> { + /// Define the threadblock-scoped matrix multiply-accumulate (A x BT) + using Mma1 = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, LayoutA, + kAlignmentA, + ElementB, typename layout::LayoutTranspose::type, + kAlignmentB, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator>::ThreadblockMma; + + /// Define the threadblock-scoped matrix multiply-accumulate (B x AT) + using Mma2 = typename mctlass::gemm::threadblock::DefaultMma< + ElementB, LayoutB, + kAlignmentB, + ElementA, typename layout::LayoutTranspose::type, + kAlignmentA, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOpBlas3< + ThreadblockShape, typename Mma1::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount, BlasMode::kSymmetric>::Epilogue; + + /// Define the kernel-level Rank2K operator. + using Rank2Kkernel = kernel::Rank2KUniversal; +}; +//////////////////////////////////////////////////////////////////////////////// + + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k_complex.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k_complex.h new file mode 100644 index 0000000..6694cac --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k_complex.h @@ -0,0 +1,334 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level Rank2K definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/rank_2k_universal.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/default_multistage_mma_complex.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_complex_tensor_op_blas3.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Operation performed by GEMM + typename Operator, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Blas3 computation mode + BlasMode BlasMode_ = BlasMode::kSymmetric> +struct DefaultRank2KComplex; + + +//////////////////////////////////////////////////////////////////////////////// +namespace detail { + +template < + /// Layout type for A matrix operand + typename LayoutA_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Complex elementwise transformation + ComplexTransform TransformA, + /// Complex elementwise transformation + ComplexTransform TransformB, + /// Blas3 computation mode (symmetric/hermitian) + BlasMode BlasMode_ + > struct Rank2KTransposedComplexTransform { + + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + +}; + + // partial specializations for HER2K CUBLAS_OP_N layout (ColumMajor) +template <> + struct Rank2KTransposedComplexTransform < + layout::ColumnMajor, layout::ColumnMajor, + ComplexTransform::kNone, ComplexTransform::kNone, + BlasMode::kHermitian> { + + static ComplexTransform const kTransformA = ComplexTransform::kConjugate; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + +}; + + // partial specializations for HER2K CUBLAS_OP_C layout (RowMajor + Complex conjugate) +template <> + struct Rank2KTransposedComplexTransform < + layout::RowMajor, layout::RowMajor, + ComplexTransform::kConjugate, ComplexTransform::kConjugate, + BlasMode::kHermitian> { + + static ComplexTransform const kTransformA = ComplexTransform::kNone; + static ComplexTransform const kTransformB = ComplexTransform::kConjugate; + +}; + +} +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture complex datatype (symmetric) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Element type for C and D matrix operands + typename ElementC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Operation performed by GEMM + typename Operator, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial> +struct DefaultRank2KComplex< + ElementA, LayoutA, ElementB, LayoutB, ElementC, + layout::RowMajor, FillModeC, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm80, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, + TransformA, TransformB, Operator, SplitKSerial, BlasMode::kSymmetric> { + + static BlasMode const kBlasMode = BlasMode::kSymmetric; + + /// Define the threadblock-scoped matrix multiply-accumulate (A x B^T) + using Mma1 = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplex< + ElementA, LayoutA, + ElementB, typename layout::LayoutTranspose::type, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + TransformA, TransformB, Operator>::ThreadblockMma; + + /// Define the threadblock-scoped matrix multiply-accumulate (B x A^T) + using Mma2 = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplex< + ElementB, LayoutB, + ElementA, typename layout::LayoutTranspose::type, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + TransformA, TransformB, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOpBlas3< + ThreadblockShape, typename Mma1::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator, kBlasMode>::Epilogue; + + /// Define the kernel-level Rank2K operator. + using Rank2Kkernel = kernel::Rank2KUniversal; + +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture complex datatype (hermitian) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Element type for C and D matrix operands + typename ElementC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Operation performed by GEMM + typename Operator, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial> +struct DefaultRank2KComplex< + ElementA, LayoutA, ElementB, LayoutB, ElementC, + layout::RowMajor, FillModeC, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm80, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, + TransformA, TransformB, Operator, SplitKSerial, BlasMode::kHermitian> { + + static BlasMode const kBlasMode = BlasMode::kHermitian; + + // Complex transform for input A and B matrices (function on input layout) + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + + using TransposedComplexTransform = detail::Rank2KTransposedComplexTransform< + LayoutA, LayoutB, + TransformA, TransformB, + kBlasMode>; + + // Complex transform on operandA and operandB (function of blas3 computation) + static ComplexTransform const kTransformOperandA = TransposedComplexTransform::kTransformA; + static ComplexTransform const kTransformOperandB = TransposedComplexTransform::kTransformB; + + /// Define the threadblock-scoped matrix multiply-accumulate (A x B^H) + using Mma1 = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplex< + ElementA, LayoutA, + ElementB, typename layout::LayoutTranspose::type, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + kTransformOperandA, kTransformOperandB, Operator>::ThreadblockMma; + + /// Define the threadblock-scoped matrix multiply-accumulate (B x A^H) + using Mma2 = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplex< + ElementB, LayoutB, + ElementA, typename layout::LayoutTranspose::type, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + kTransformOperandA, kTransformOperandB, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOpBlas3< + ThreadblockShape, typename Mma1::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator, kBlasMode>::Epilogue; + + /// Define the kernel-level Rank2K operator. + using Rank2Kkernel = kernel::Rank2KUniversal; + +}; + +//////////////////////////////////////////////////////////////////////////////// + + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k_grouped.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k_grouped.h new file mode 100644 index 0000000..6003dbf --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k_grouped.h @@ -0,0 +1,355 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level grouped Rank2K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/complex.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/numeric_types.h" + +#include "mctlass/gemm/kernel/rank_2k_transpose_operands.h" +#include "mctlass/gemm/kernel/default_rank_2k.h" +#include "mctlass/gemm/kernel/default_rank_2k_complex.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// Blas3 computation mode + BlasMode BlasMode_ = BlasMode::kSymmetric, + /// Whether the schedule of problems to visit has been precomputed + GroupScheduleMode GroupScheduleMode_ = GroupScheduleMode::kDeviceOnly, + /// + typename Enable = void + > +struct DefaultRank2KGrouped; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Real-valued grouped Rank2K +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// Blas3 computation mode + BlasMode BlasMode_, + /// Whether the schedule of problems to visit has been precomputed + GroupScheduleMode GroupScheduleMode_ + > +struct DefaultRank2KGrouped::value>::type +> { + // If true, we must construct a 'transposed-and-exchanged' Rank2K operator. + static bool const kInternalTranspose = platform::is_same::value; + + using MapArguments = kernel::detail::Rank2KMapArguments< + ElementA, + LayoutA, + TransformA, + kAlignmentA, + ElementB, + LayoutB, + TransformB, + kAlignmentB, + LayoutC, + FillModeC, + kInternalTranspose + >; + + // Define the default grouped Rank2K kernel + using DefaultRank2Kkernel = typename kernel::DefaultRank2K< + typename MapArguments::ElementA, + typename MapArguments::LayoutA, + MapArguments::kAlignmentA, + typename MapArguments::ElementB, + typename MapArguments::LayoutB, + MapArguments::kAlignmentB, + ElementC, + typename MapArguments::LayoutC, + MapArguments::kFillModeC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + false, // SplitKSerial + Operator, + BlasMode_ + >::Rank2Kkernel; + + /// Define the kernel in terms of the default kernel + using Rank2Kkernel = kernel::Rank2KGrouped< + typename DefaultRank2Kkernel::Mma1, + typename DefaultRank2Kkernel::Mma2, + typename DefaultRank2Kkernel::Epilogue, + ThreadblockSwizzle, + TransformA, + TransformB, + DefaultRank2Kkernel::kFillModeC, + DefaultRank2Kkernel::kBlasMode, + GroupScheduleMode_, + kInternalTranspose + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Complex-valued grouped Rank2K +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// Blas3 computation mode + BlasMode BlasMode_, + /// Whether the schedule of problems to visit has been precomputed + GroupScheduleMode GroupScheduleMode_ + > +struct DefaultRank2KGrouped::value>::type +> { + // If true, we must construct a 'transposed-and-exchanged' Rank2K operator. + static bool const kInternalTranspose = platform::is_same::value; + + using MapArguments = kernel::detail::Rank2KMapArguments< + ElementA, + LayoutA, + TransformA, + kAlignmentA, + ElementB, + LayoutB, + TransformB, + kAlignmentB, + LayoutC, + FillModeC, + kInternalTranspose + >; + + // Define the default grouped Rank2K kernel + using DefaultRank2Kkernel = typename kernel::DefaultRank2KComplex< + typename MapArguments::ElementA, + typename MapArguments::LayoutA, + typename MapArguments::ElementB, + typename MapArguments::LayoutB, + ElementC, + typename MapArguments::LayoutC, + MapArguments::kFillModeC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + MapArguments::kTransformA, + MapArguments::kTransformB, + Operator, + false, // SplitKSerial + BlasMode_ + >::Rank2Kkernel; + + /// Define the kernel in terms of the default kernel + /// Pass through the user-provided TransformA and TransformB so as to + /// correctly set public-facing TransformA and TransformB in kernel::Rank2KGrouped. + /// This is needed because kernel::DefaultRank2KComplex may change TransformA and + /// TransformB that become template arguments to Mma1 and Mma2. + using Rank2Kkernel = kernel::Rank2KGrouped< + typename DefaultRank2Kkernel::Mma1, + typename DefaultRank2Kkernel::Mma2, + typename DefaultRank2Kkernel::Epilogue, + ThreadblockSwizzle, + TransformA, + TransformB, + DefaultRank2Kkernel::kFillModeC, + DefaultRank2Kkernel::kBlasMode, + GroupScheduleMode_, + kInternalTranspose + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k_universal.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k_universal.h new file mode 100644 index 0000000..1fcef52 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_2k_universal.h @@ -0,0 +1,346 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level Rank 2k definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" + +#include "mctlass/complex.h" +#include "mctlass/layout/matrix.h" + +#include "mctlass/gemm/kernel/rank_2k_universal.h" +#include "mctlass/gemm/kernel/default_rank_2k.h" +#include "mctlass/gemm/kernel/default_rank_2k_complex.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by SYRK + typename Operator, + /// Blas3 computation mode (symmetric/hermitian) + BlasMode BlasMode_ = BlasMode::kSymmetric, + /// + typename Enable = void + > +struct DefaultRank2KUniversal; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Real-valued Rank 2k update kernels +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by Rank2k + typename Operator> +struct DefaultRank2KUniversal< + ElementA, + LayoutA, + ComplexTransform::kNone, // transform A + kAlignmentA, + ElementB, + LayoutB, + ComplexTransform::kNone, // transform B + kAlignmentB, + ElementC, + LayoutC, + FillModeC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + BlasMode::kSymmetric, + typename std::enable_if< ! mctlass::is_complex::value>::type +> { + + using DefaultRank2Kkernel = typename kernel::DefaultRank2K< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementC, + LayoutC, + FillModeC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + BlasMode::kSymmetric + >::Rank2Kkernel; + + /// Define the kernel in terms of the default kernel + using Rank2Kkernel = kernel::Rank2KUniversal< + typename DefaultRank2Kkernel::Mma1, + typename DefaultRank2Kkernel::Mma2, + typename DefaultRank2Kkernel::Epilogue, + ThreadblockSwizzle, + FillModeC, + BlasMode::kSymmetric + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Complex-valued Rank 2K update kernels +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by SYRK + typename Operator, + // BlasMode + BlasMode kBlasMode + > + +struct DefaultRank2KUniversal< + ElementA, + LayoutA, + TransformA, + kAlignmentA, + ElementB, + LayoutB, + TransformB, + kAlignmentB, + ElementC, + LayoutC, + FillModeC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + kBlasMode, + typename std::enable_if::value>::type +> { + + using DefaultRank2Kkernel = typename kernel::DefaultRank2KComplex< + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + FillModeC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + TransformA, + TransformB, + Operator, + SplitKSerial, + kBlasMode + >::Rank2Kkernel; + + /// Define the kernel in terms of the default kernel + using Rank2Kkernel = kernel::Rank2KUniversal< + typename DefaultRank2Kkernel::Mma1, + typename DefaultRank2Kkernel::Mma2, + typename DefaultRank2Kkernel::Epilogue, + ThreadblockSwizzle, + FillModeC, + kBlasMode + >; +}; + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_k.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_k.h new file mode 100644 index 0000000..183028a --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_k.h @@ -0,0 +1,247 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level RankK definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/rank_k_universal.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op_blas3.h" +#include "mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_simt.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Blas3 computation mode + BlasMode BlasMode_ = BlasMode::kSymmetric> +struct DefaultRankK; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Hopper Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for C and D matrix operands + typename ElementC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator> +struct DefaultRankK< + ElementA, LayoutA, kAlignmentA, + ElementC,layout::RowMajor, FillModeC, + ElementAccumulator, arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, SplitKSerial, + Operator> { + /// Define the threadblock-scoped matrix multiply-accumulate (A x AT) + using Mma = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, LayoutA, + kAlignmentA, + ElementA, typename layout::LayoutTranspose::type, + kAlignmentA, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator>::ThreadblockMma; + + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOpBlas3< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount, BlasMode::kSymmetric>::Epilogue; + + /// Define the kernel-level Rank2 operator. + using RankKkernel = kernel::RankKUniversal; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for C and D matrix operands + typename ElementC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator> +struct DefaultRankK< + ElementA, LayoutA, kAlignmentA, + ElementC,layout::RowMajor, FillModeC, + ElementAccumulator, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, SplitKSerial, + Operator> { + /// Define the threadblock-scoped matrix multiply-accumulate (A x AT) + using Mma = typename mctlass::gemm::threadblock::DefaultMma< + ElementA, LayoutA, + kAlignmentA, + ElementA, typename layout::LayoutTranspose::type, + kAlignmentA, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator>::ThreadblockMma; + + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOpBlas3< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount, BlasMode::kSymmetric>::Epilogue; + + /// Define the kernel-level Rank2 operator. + using RankKkernel = kernel::RankKUniversal; +}; +//////////////////////////////////////////////////////////////////////////////// + + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_k_complex.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_k_complex.h new file mode 100644 index 0000000..991af5c --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_k_complex.h @@ -0,0 +1,429 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level RankK definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/rank_k_universal.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/default_multistage_mma_complex.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_complex_tensor_op_blas3.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h" +#endif //CUTLASS_ARCH_WMMA_ENABLED + + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Operation performed by GEMM + typename Operator, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Blas3 computation mode + BlasMode BlasMode_ = BlasMode::kSymmetric> +struct DefaultRankKComplex; + + +//////////////////////////////////////////////////////////////////////////////// +namespace detail { + +template < + /// Layout type for A matrix operand + typename LayoutA_, + /// Complex elementwise transformation + ComplexTransform TransformA, + /// Blas3 computation mode (symmetric/hermitian) + BlasMode BlasMode_ + > struct RankKTransposedComplexTransform { + + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformA; + +}; + + // partial specializations for HERK CUBLAS_OP_N layout (ColumMajor) +template <> + struct RankKTransposedComplexTransform < + layout::ColumnMajor, + ComplexTransform::kNone, + BlasMode::kHermitian> { + + static ComplexTransform const kTransformA = ComplexTransform::kConjugate; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + +}; + + // partial specializations for HERK CUBLAS_OP_C layout (RowMajor + Complex conjugate) +template <> + struct RankKTransposedComplexTransform < + layout::RowMajor, + ComplexTransform::kConjugate, + BlasMode::kHermitian> { + + static ComplexTransform const kTransformA = ComplexTransform::kNone; + static ComplexTransform const kTransformB = ComplexTransform::kConjugate; + +}; + +} +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Hopper Architecture complex datatype (symmetric) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for C and D matrix operands + typename ElementC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Operation performed by GEMM + typename Operator, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial> +struct DefaultRankKComplex< + ElementA, LayoutA, ElementC, + layout::RowMajor, FillModeC, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm90, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, + TransformA, Operator, SplitKSerial, BlasMode::kSymmetric> { + + static BlasMode const kBlasMode = BlasMode::kSymmetric; + + /// Define the threadblock-scoped matrix multiply-accumulate (A x B^T) + using Mma = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplex< + ElementA, LayoutA, + ElementA, typename layout::LayoutTranspose::type, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, Stages, + TransformA, TransformA, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOpBlas3< + ThreadblockShape, typename Mma::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator, kBlasMode>::Epilogue; + + /// Define the kernel-level RankK operator. + using RankKkernel = kernel::RankKUniversal; + +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Hopper Architecture complex datatype (hermitian) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for C and D matrix operands + typename ElementC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Operation performed by GEMM + typename Operator, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial> +struct DefaultRankKComplex< + ElementA, LayoutA, ElementC, + layout::RowMajor, FillModeC, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm90, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, + TransformA, Operator, SplitKSerial, BlasMode::kHermitian> { + + static BlasMode const kBlasMode = BlasMode::kHermitian; + + // Complex transform for input A and B matrices (function on input layout) + static ComplexTransform const kTransformA = TransformA; + + using TransposedComplexTransform = detail::RankKTransposedComplexTransform< + LayoutA, + TransformA, + kBlasMode>; + + // Complex transform on operandA and operandB (function of blas3 computation) + static ComplexTransform const kTransformOperandA = TransposedComplexTransform::kTransformA; + static ComplexTransform const kTransformOperandB = TransposedComplexTransform::kTransformB; + + /// Define the threadblock-scoped matrix multiply-accumulate (A x A^H) + using Mma = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplex< + ElementA, LayoutA, + ElementA, typename layout::LayoutTranspose::type, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, Stages, + kTransformOperandA, kTransformOperandB, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOpBlas3< + ThreadblockShape, typename Mma::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator, kBlasMode>::Epilogue; + + /// Define the kernel-level RankK operator. + using RankKkernel = kernel::RankKUniversal; + +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture complex datatype (symmetric) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for C and D matrix operands + typename ElementC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Operation performed by GEMM + typename Operator, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial> +struct DefaultRankKComplex< + ElementA, LayoutA, ElementC, + layout::RowMajor, FillModeC, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm80, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, + TransformA, Operator, SplitKSerial, BlasMode::kSymmetric> { + + static BlasMode const kBlasMode = BlasMode::kSymmetric; + + /// Define the threadblock-scoped matrix multiply-accumulate (A x B^T) + using Mma = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplex< + ElementA, LayoutA, + ElementA, typename layout::LayoutTranspose::type, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + TransformA, TransformA, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOpBlas3< + ThreadblockShape, typename Mma::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator, kBlasMode>::Epilogue; + + /// Define the kernel-level RankK operator. + using RankKkernel = kernel::RankKUniversal; + +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture complex datatype (hermitian) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for C and D matrix operands + typename ElementC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Operation performed by GEMM + typename Operator, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial> +struct DefaultRankKComplex< + ElementA, LayoutA, ElementC, + layout::RowMajor, FillModeC, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm80, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, + TransformA, Operator, SplitKSerial, BlasMode::kHermitian> { + + static BlasMode const kBlasMode = BlasMode::kHermitian; + + // Complex transform for input A and B matrices (function on input layout) + static ComplexTransform const kTransformA = TransformA; + + using TransposedComplexTransform = detail::RankKTransposedComplexTransform< + LayoutA, + TransformA, + kBlasMode>; + + // Complex transform on operandA and operandB (function of blas3 computation) + static ComplexTransform const kTransformOperandA = TransposedComplexTransform::kTransformA; + static ComplexTransform const kTransformOperandB = TransposedComplexTransform::kTransformB; + + /// Define the threadblock-scoped matrix multiply-accumulate (A x A^H) + using Mma = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplex< + ElementA, LayoutA, + ElementA, typename layout::LayoutTranspose::type, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + kTransformOperandA, kTransformOperandB, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOpBlas3< + ThreadblockShape, typename Mma::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator, kBlasMode>::Epilogue; + + /// Define the kernel-level RankK operator. + using RankKkernel = kernel::RankKUniversal; + +}; + +//////////////////////////////////////////////////////////////////////////////// + + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_k_universal.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_k_universal.h new file mode 100644 index 0000000..63caf9c --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_rank_k_universal.h @@ -0,0 +1,305 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level Rank k definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" + +#include "mctlass/complex.h" +#include "mctlass/layout/matrix.h" + +#include "mctlass/gemm/kernel/rank_k_universal.h" +#include "mctlass/gemm/kernel/default_rank_k.h" +#include "mctlass/gemm/kernel/default_rank_k_complex.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by SYRK + typename Operator, + /// Blas3 computation mode (symmetric/hermitian) + BlasMode BlasMode_ = BlasMode::kSymmetric, + /// + typename Enable = void + > +struct DefaultRankKUniversal; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Real-valued Rank k update kernels +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by Rank2k + typename Operator> +struct DefaultRankKUniversal< + ElementA, + LayoutA, + ComplexTransform::kNone, // transform A + kAlignmentA, + ElementC, + LayoutC, + FillModeC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + BlasMode::kSymmetric, + typename std::enable_if< ! mctlass::is_complex::value>::type +> { + + using DefaultRankKkernel = typename kernel::DefaultRankK< + ElementA, + LayoutA, + kAlignmentA, + ElementC, + LayoutC, + FillModeC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + BlasMode::kSymmetric + >::RankKkernel; + + /// Define the kernel in terms of the default kernel + using RankKkernel = kernel::RankKUniversal< + typename DefaultRankKkernel::Mma, + typename DefaultRankKkernel::Epilogue, + ThreadblockSwizzle, + FillModeC + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Complex-valued Rank 2K update kernels +// +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Fill Mode for C (kLower or kUpper) + FillMode FillModeC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by SYRK + typename Operator, + // BlasMode + BlasMode kBlasMode + > + +struct DefaultRankKUniversal< + ElementA, + LayoutA, + TransformA, + kAlignmentA, + ElementC, + LayoutC, + FillModeC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + kBlasMode, + typename std::enable_if::value>::type +> { + + using DefaultRankKkernel = typename kernel::DefaultRankKComplex< + ElementA, + LayoutA, + ElementC, + LayoutC, + FillModeC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + TransformA, + Operator, + SplitKSerial, + kBlasMode + >::RankKkernel; + + /// Define the kernel in terms of the default kernel + using RankKkernel = kernel::RankKUniversal< + typename DefaultRankKkernel::Mma, + typename DefaultRankKkernel::Epilogue, + ThreadblockSwizzle, + FillModeC + >; +}; + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_symm.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_symm.h new file mode 100644 index 0000000..1d3ece8 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_symm.h @@ -0,0 +1,322 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level SYMM/HEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/symm_universal.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/gemm/threadblock/default_trmm.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_simt.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Side Mode for A (kLeft or kRight) + SideMode kSideModeA, + /// Fill Mode for A (kLower or kUpper) + FillMode kFillModeA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator, + /// Blas3 computation mode + BlasMode BlasMode_ = BlasMode::kSymmetric> +struct DefaultSymm; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Hopper Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Side Mode for A (kLeft or kRight) + SideMode kSideModeA, + /// Fill Mode for A (kLower or kUpper) + FillMode kFillModeA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator> +struct DefaultSymm< + ElementA, LayoutA, kSideModeA, kFillModeA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + ElementC,layout::RowMajor, + ElementAccumulator, arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, SplitKSerial, + Operator> { + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + /// TRMM - with diagonal: alpha * A * B or alpha * B * A + static const DiagType kDiagTypeMma1 = DiagType::kNonUnit; + using Mma1 = typename mctlass::gemm::threadblock::DefaultTrmm< + ElementA, LayoutA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + kSideModeA, kFillModeA, kDiagTypeMma1, + ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, + Stages, Operator>::ThreadblockMma; + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + /// TRMM - withOUT diagonal: alpha * AT * B or alpha * B * AT + static const DiagType kDiagTypeMma2 = DiagType::kZero; + using LayoutAMma2 = typename platform::conditional< + (kSideModeA == SideMode::kLeft), + typename layout::LayoutTranspose::type, + LayoutA + >::type; + using LayoutBMma2 = typename platform::conditional< + (kSideModeA == SideMode::kLeft), + LayoutB, + typename layout::LayoutTranspose::type + >::type; + using Mma2 = typename mctlass::gemm::threadblock::DefaultTrmm< + ElementA, LayoutAMma2, kAlignmentA, + ElementB, LayoutBMma2, kAlignmentB, + kSideModeA, InvertFillMode::mode, kDiagTypeMma2, + ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, + Stages, Operator>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, typename Mma1::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount>::Epilogue; + + /// Define the kernel-level SYMM/HEMM operator. + using SymmKernel = kernel::SymmUniversal; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Side Mode for A (kLeft or kRight) + SideMode kSideModeA, + /// Fill Mode for A (kLower or kUpper) + FillMode kFillModeA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator> +struct DefaultSymm< + ElementA, LayoutA, kSideModeA, kFillModeA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + ElementC,layout::RowMajor, + ElementAccumulator, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, SplitKSerial, + Operator> { + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + /// TRMM - with diagonal: alpha * A * B or alpha * B * A + static const DiagType kDiagTypeMma1 = DiagType::kNonUnit; + using Mma1 = typename mctlass::gemm::threadblock::DefaultTrmm< + ElementA, LayoutA, kAlignmentA, + ElementB, LayoutB, kAlignmentB, + kSideModeA, kFillModeA, kDiagTypeMma1, + ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, + Stages, Operator>::ThreadblockMma; + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + /// TRMM - withOUT diagonal: alpha * AT * B or alpha * B * AT + static const DiagType kDiagTypeMma2 = DiagType::kZero; + using LayoutAMma2 = typename platform::conditional< + (kSideModeA == SideMode::kLeft), + typename layout::LayoutTranspose::type, + LayoutA + >::type; + using LayoutBMma2 = typename platform::conditional< + (kSideModeA == SideMode::kLeft), + LayoutB, + typename layout::LayoutTranspose::type + >::type; + using Mma2 = typename mctlass::gemm::threadblock::DefaultTrmm< + ElementA, LayoutAMma2, kAlignmentA, + ElementB, LayoutBMma2, kAlignmentB, + kSideModeA, InvertFillMode::mode, kDiagTypeMma2, + ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, + Stages, Operator>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, typename Mma1::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount>::Epilogue; + + /// Define the kernel-level SYMM/HEMM operator. + using SymmKernel = kernel::SymmUniversal; +}; +//////////////////////////////////////////////////////////////////////////////// + + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_symm_complex.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_symm_complex.h new file mode 100644 index 0000000..3f42432 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_symm_complex.h @@ -0,0 +1,509 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level SYMM/HEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/symm_universal.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/default_multistage_trmm_complex.h" +#include "mctlass/gemm/threadblock/default_multistage_mma_complex.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_complex_tensor_op.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Side Mode for A (kLeft or kRight) + SideMode kSideModeA, + /// Fill Mode for A (kLower or kUpper) + FillMode kFillModeA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Blas3 computation mode + BlasMode BlasMode_ = BlasMode::kSymmetric> +struct DefaultSymmComplex; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture complex datatype (symmetric) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Side Mode for A (kLeft or kRight) + SideMode kSideModeA, + /// Fill Mode for A (kLower or kUpper) + FillMode kFillModeA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial> +struct DefaultSymmComplex< + ElementA, LayoutA, kSideModeA, kFillModeA, ElementB, LayoutB, ElementC, + layout::RowMajor, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm80, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, + Operator, SplitKSerial, BlasMode::kSymmetric> { + + static BlasMode const kBlasMode = BlasMode::kSymmetric; + // Complex Transform don't appply to A or B for SYMM + static ComplexTransform const TransformA = ComplexTransform::kNone; + static ComplexTransform const TransformB = ComplexTransform::kNone; + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + /// TRMM - with diagonal: alpha * A * B or alpha * B * A + static const DiagType kDiagTypeMma1 = DiagType::kNonUnit; + using Mma1 = typename mctlass::gemm::threadblock::DefaultMultistageTrmmComplex< + ElementA, LayoutA, + ElementB, LayoutB, + kSideModeA, kFillModeA, kDiagTypeMma1, + ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, + Stages, TransformA, TransformB, Operator>::ThreadblockMma; + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + /// TRMM - withOUT diagonal: alpha * AT * B or alpha * B * AT + static const DiagType kDiagTypeMma2 = DiagType::kZero; + using LayoutAMma2 = typename platform::conditional< + (kSideModeA == SideMode::kLeft), + typename layout::LayoutTranspose::type, + LayoutA + >::type; + using LayoutBMma2 = typename platform::conditional< + (kSideModeA == SideMode::kLeft), + LayoutB, + typename layout::LayoutTranspose::type + >::type; + using Mma2 = typename mctlass::gemm::threadblock::DefaultMultistageTrmmComplex< + ElementA, LayoutAMma2, + ElementB, LayoutBMma2, + kSideModeA, InvertFillMode::mode, kDiagTypeMma2, + ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, + Stages, TransformA, TransformB, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp< + ThreadblockShape, typename Mma1::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator>::Epilogue; + + /// Define the kernel-level Symm operator. + using SymmKernel = kernel::SymmUniversal; + +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Hopper Architecture complex datatype (symmetric) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Side Mode for A (kLeft or kRight) + SideMode kSideModeA, + /// Fill Mode for A (kLower or kUpper) + FillMode kFillModeA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial> +struct DefaultSymmComplex< + ElementA, LayoutA, kSideModeA, kFillModeA, ElementB, LayoutB, ElementC, + layout::RowMajor, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm90, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, + Operator, SplitKSerial, BlasMode::kSymmetric> { + + static BlasMode const kBlasMode = BlasMode::kSymmetric; + // Complex Transform don't appply to A or B for SYMM + static ComplexTransform const TransformA = ComplexTransform::kNone; + static ComplexTransform const TransformB = ComplexTransform::kNone; + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + /// TRMM - with diagonal: alpha * A * B or alpha * B * A + static const DiagType kDiagTypeMma1 = DiagType::kNonUnit; + using Mma1 = typename mctlass::gemm::threadblock::DefaultMultistageTrmmComplex< + ElementA, LayoutA, + ElementB, LayoutB, + kSideModeA, kFillModeA, kDiagTypeMma1, + ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, + Stages, TransformA, TransformB, Operator>::ThreadblockMma; + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + /// TRMM - withOUT diagonal: alpha * AT * B or alpha * B * AT + static const DiagType kDiagTypeMma2 = DiagType::kZero; + using LayoutAMma2 = typename platform::conditional< + (kSideModeA == SideMode::kLeft), + typename layout::LayoutTranspose::type, + LayoutA + >::type; + using LayoutBMma2 = typename platform::conditional< + (kSideModeA == SideMode::kLeft), + LayoutB, + typename layout::LayoutTranspose::type + >::type; + using Mma2 = typename mctlass::gemm::threadblock::DefaultMultistageTrmmComplex< + ElementA, LayoutAMma2, + ElementB, LayoutBMma2, + kSideModeA, InvertFillMode::mode, kDiagTypeMma2, + ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, + Stages, TransformA, TransformB, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp< + ThreadblockShape, typename Mma1::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator>::Epilogue; + + /// Define the kernel-level Symm operator. + using SymmKernel = kernel::SymmUniversal; + +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Hopper Architecture complex datatype (hermitian) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Side Mode for A (kLeft or kRight) + SideMode kSideModeA, + /// Fill Mode for A (kLower or kUpper) + FillMode kFillModeA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial> +struct DefaultSymmComplex< + ElementA, LayoutA, kSideModeA, kFillModeA, ElementB, LayoutB, ElementC, + layout::RowMajor, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm90, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, + Operator, SplitKSerial, BlasMode::kHermitian> { + + static BlasMode const kBlasMode = BlasMode::kHermitian; + + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + /// TRMM - with diagonal: alpha * A * B or alpha * B * A + static const DiagType kDiagTypeMma1 = DiagType::kNonUnit; + static ComplexTransform const TransformAMma1 = ComplexTransform::kNone; + static ComplexTransform const TransformBMma1 = ComplexTransform::kNone; + using Mma1 = typename mctlass::gemm::threadblock::DefaultMultistageTrmmComplex< + ElementA, LayoutA, + ElementB, LayoutB, + kSideModeA, kFillModeA, kDiagTypeMma1, + ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, + Stages, TransformAMma1, TransformBMma1, Operator, BlasMode::kHermitian>::ThreadblockMma; + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + /// TRMM - withOUT diagonal - with conjugate transpose: alpha * AT * B or alpha * B * AT + static const DiagType kDiagTypeMma2 = DiagType::kZero; + using LayoutAMma2 = typename platform::conditional< + (kSideModeA == SideMode::kLeft), + typename layout::LayoutTranspose::type, + LayoutA + >::type; + using LayoutBMma2 = typename platform::conditional< + (kSideModeA == SideMode::kLeft), + LayoutB, + typename layout::LayoutTranspose::type + >::type; + static ComplexTransform const TransformAMma2 = (kSideModeA == SideMode::kLeft) ? + ComplexTransform::kConjugate : ComplexTransform::kNone; + static ComplexTransform const TransformBMma2 = (kSideModeA == SideMode::kLeft) ? + ComplexTransform::kNone : ComplexTransform::kConjugate; + + using Mma2 = typename mctlass::gemm::threadblock::DefaultMultistageTrmmComplex< + ElementA, LayoutAMma2, + ElementB, LayoutBMma2, + kSideModeA, InvertFillMode::mode, kDiagTypeMma2, + ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, + Stages, TransformAMma2, TransformBMma2, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp< + ThreadblockShape, typename Mma1::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator>::Epilogue; + + /// Define the kernel-level Symm operator. + using SymmKernel = kernel::SymmUniversal; + +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture complex datatype (hermitian) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Side Mode for A (kLeft or kRight) + SideMode kSideModeA, + /// Fill Mode for A (kLower or kUpper) + FillMode kFillModeA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial> +struct DefaultSymmComplex< + ElementA, LayoutA, kSideModeA, kFillModeA, ElementB, LayoutB, ElementC, + layout::RowMajor, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm80, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, + Operator, SplitKSerial, BlasMode::kHermitian> { + + static BlasMode const kBlasMode = BlasMode::kHermitian; + + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + /// TRMM - with diagonal: alpha * A * B or alpha * B * A + static const DiagType kDiagTypeMma1 = DiagType::kNonUnit; + static ComplexTransform const TransformAMma1 = ComplexTransform::kNone; + static ComplexTransform const TransformBMma1 = ComplexTransform::kNone; + using Mma1 = typename mctlass::gemm::threadblock::DefaultMultistageTrmmComplex< + ElementA, LayoutA, + ElementB, LayoutB, + kSideModeA, kFillModeA, kDiagTypeMma1, + ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, + Stages, TransformAMma1, TransformBMma1, Operator, BlasMode::kHermitian>::ThreadblockMma; + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + /// TRMM - withOUT diagonal - with conjugate transpose: alpha * AT * B or alpha * B * AT + static const DiagType kDiagTypeMma2 = DiagType::kZero; + using LayoutAMma2 = typename platform::conditional< + (kSideModeA == SideMode::kLeft), + typename layout::LayoutTranspose::type, + LayoutA + >::type; + using LayoutBMma2 = typename platform::conditional< + (kSideModeA == SideMode::kLeft), + LayoutB, + typename layout::LayoutTranspose::type + >::type; + static ComplexTransform const TransformAMma2 = (kSideModeA == SideMode::kLeft) ? + ComplexTransform::kConjugate : ComplexTransform::kNone; + static ComplexTransform const TransformBMma2 = (kSideModeA == SideMode::kLeft) ? + ComplexTransform::kNone : ComplexTransform::kConjugate; + + using Mma2 = typename mctlass::gemm::threadblock::DefaultMultistageTrmmComplex< + ElementA, LayoutAMma2, + ElementB, LayoutBMma2, + kSideModeA, InvertFillMode::mode, kDiagTypeMma2, + ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, + Stages, TransformAMma2, TransformBMma2, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp< + ThreadblockShape, typename Mma1::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator>::Epilogue; + + /// Define the kernel-level Symm operator. + using SymmKernel = kernel::SymmUniversal; + +}; + +//////////////////////////////////////////////////////////////////////////////// + + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_symm_universal.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_symm_universal.h new file mode 100644 index 0000000..5887133 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_symm_universal.h @@ -0,0 +1,342 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level SYMM/HEMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" + +#include "mctlass/complex.h" +#include "mctlass/layout/matrix.h" + +#include "mctlass/gemm/kernel/symm_universal.h" +#include "mctlass/gemm/kernel/default_symm.h" +#include "mctlass/gemm/kernel/default_symm_complex.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Side Mode for A (kLeft or kRight) + SideMode SideModeA, + /// Fill Mode for A (kLower or kUpper) + FillMode FillModeA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by SYRK + typename Operator, + /// Blas3 computation mode (symmetric/hermitian) + BlasMode BlasMode_ = BlasMode::kSymmetric, + /// + typename Enable = void + > +struct DefaultSymmUniversal; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Real-valued SYMM/HEMM update kernels +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Side Mode for A (kLeft or kRight) + SideMode SideModeA, + /// Fill Mode for A (kLower or kUpper) + FillMode FillModeA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by SYMM/HEMM + typename Operator> +struct DefaultSymmUniversal< + ElementA, + LayoutA, + SideModeA, + FillModeA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + BlasMode::kSymmetric, + typename std::enable_if< ! mctlass::is_complex::value>::type +> { + + using DefaultSymmkernel = typename kernel::DefaultSymm< + ElementA, + LayoutA, + SideModeA, + FillModeA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + BlasMode::kSymmetric + >::SymmKernel; + + /// Define the kernel in terms of the default kernel + using SymmKernel = kernel::SymmUniversal< + typename DefaultSymmkernel::Mma1, + typename DefaultSymmkernel::Mma2, + typename DefaultSymmkernel::Epilogue, + ThreadblockSwizzle, + SideModeA, + FillModeA + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Complex-valued SYMM/HEMM update kernels +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Side Mode for A (kLeft or kRight) + SideMode SideModeA, + /// Fill Mode for A (kLower or kUpper) + FillMode FillModeA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by SYRK + typename Operator, + // BlasMode + BlasMode kBlasMode + > + +struct DefaultSymmUniversal< + ElementA, + LayoutA, + SideModeA, + FillModeA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + kBlasMode, + typename std::enable_if::value>::type +> { + + using DefaultSymmkernel = typename kernel::DefaultSymmComplex< + ElementA, + LayoutA, + SideModeA, + FillModeA, + ElementB, + LayoutB, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + Operator, + SplitKSerial, + kBlasMode + >::SymmKernel; + + /// Define the kernel in terms of the default kernel + using SymmKernel = kernel::SymmUniversal< + typename DefaultSymmkernel::Mma1, + typename DefaultSymmkernel::Mma2, + typename DefaultSymmkernel::Epilogue, + ThreadblockSwizzle, + SideModeA, + FillModeA + >; +}; + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_trmm.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_trmm.h new file mode 100644 index 0000000..1099477 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_trmm.h @@ -0,0 +1,269 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +// +/*! \file + \brief + Default kernel-level TRMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. +*/ + +#pragma once + +#include "mctlass/blas3.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/trmm_universal.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/default_trmm.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" + +#include "mctlass/epilogue/threadblock/default_epilogue_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_simt.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Side Mode for the kernel + SideMode SideMode_, + /// Fill Mode for the triangular matrix + FillMode FillMode_, + /// Diag Type for the triangular matrix + DiagType DiagType_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator> +struct DefaultTrmm; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Hopper Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Diag Type for the triangular matrix + DiagType kDiagType, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator> +struct DefaultTrmm { + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultTrmm< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + kSideMode, kFillMode, kDiagType, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm90, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount>::Epilogue; + + /// Define the kernel-level TRMM operator. + using TrmmKernel = kernel::TrmmUniversal; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of A matrix in units of elements + int kAlignmentB, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Diag Type for the triangular matrix + DiagType kDiagType, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by GEMM + typename Operator> +struct DefaultTrmm { + + /// Define the threadblock-scoped triagular matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultTrmm< + ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, + kSideMode, kFillMode, kDiagType, + ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, + ThreadblockShape, WarpShape, InstructionShape, Stages, + Operator>::ThreadblockMma; + + static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueTensorOp< + ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp, + EpilogueOutputOp::kCount>::Epilogue; + + /// Define the kernel-level TRMM operator. + using TrmmKernel = kernel::TrmmUniversal; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_trmm_complex.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_trmm_complex.h new file mode 100644 index 0000000..5680ad8 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_trmm_complex.h @@ -0,0 +1,265 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level TRMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" + +#include "mctlass/layout/matrix.h" + +#include "mctlass/epilogue/threadblock/epilogue.h" +#include "mctlass/epilogue/thread/linear_combination.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/trmm_universal.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/default_multistage_trmm_complex.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle.h" +#include "mctlass/epilogue/threadblock/default_epilogue_complex_tensor_op.h" +#include "mctlass/epilogue/threadblock/default_epilogue_simt.h" + +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Side Mode for the kernel + SideMode SideMode_, + /// Fill Mode for the triangular matrix + FillMode FillMode_, + /// Diag Type for the triangular matrix + DiagType DiagType_, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Multiply-add operator + // (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial +> +struct DefaultTrmmComplex; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Hopper Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Diag Type for the triangular matrix + DiagType kDiagType, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Multiply-add operator + // (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial + > +struct DefaultTrmmComplex< + ElementA, LayoutA, ElementB, LayoutB, + kSideMode, kFillMode, kDiagType, + ElementC, layout::RowMajor, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm90, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, TransformA, TransformB, Operator, SplitKSerial> { + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMultistageTrmmComplex< + ElementA, LayoutA, ElementB, LayoutB, + kSideMode, kFillMode, kDiagType, + ElementAccumulator,layout::RowMajor, arch::OpClassTensorOp, arch::Sm90, ThreadblockShape, + WarpShape, InstructionShape, Stages, TransformA, TransformB, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp< + ThreadblockShape, typename Mma::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator>::Epilogue; + + /// Define the kernel-level TRMM operator. + using TrmmKernel = kernel::TrmmUniversal; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Ampere Architecture +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Diag Type for the triangular matrix + DiagType kDiagType, + /// Element type for C and D matrix operands + typename ElementC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Multiply-add operator + // (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator, + /// If true, kernel is configured to support serial reduction in the epilogue + bool SplitKSerial + > +struct DefaultTrmmComplex< + ElementA, LayoutA, ElementB, LayoutB, + kSideMode, kFillMode, kDiagType, + ElementC, layout::RowMajor, ElementAccumulator, arch::OpClassTensorOp, + arch::Sm80, ThreadblockShape, WarpShape, InstructionShape, + EpilogueOutputOp, ThreadblockSwizzle, Stages, TransformA, TransformB, Operator, SplitKSerial> { + + /// Define the threadblock-scoped matrix multiply-accumulate + using Mma = typename mctlass::gemm::threadblock::DefaultMultistageTrmmComplex< + ElementA, LayoutA, ElementB, LayoutB, + kSideMode, kFillMode, kDiagType, + ElementAccumulator,layout::RowMajor, arch::OpClassTensorOp, arch::Sm80, ThreadblockShape, + WarpShape, InstructionShape, Stages, TransformA, TransformB, Operator>::ThreadblockMma; + + /// Define the epilogue + using Epilogue = + typename mctlass::epilogue::threadblock::DefaultEpilogueComplexTensorOp< + ThreadblockShape, typename Mma::Operator, 1, EpilogueOutputOp, + EpilogueOutputOp::kCount, Operator>::Epilogue; + + /// Define the kernel-level TRMM operator. + using TrmmKernel = kernel::TrmmUniversal; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/default_trmm_universal.h b/csrc/mctlass/include/mctlass/gemm/kernel/default_trmm_universal.h new file mode 100644 index 0000000..6f0e581 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/default_trmm_universal.h @@ -0,0 +1,359 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + Default kernel-level TRMM definitions combine threadblock-scoped matrix multiply-add with + the appropriate threadblock-scoped epilogue. + + Note, MCTLASS epilogues universally target row-major outputs. Column-major outputs are + accommodated by exchanging A and B operands and assuming transposed layouts. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" + +#include "mctlass/complex.h" +#include "mctlass/layout/matrix.h" + +#include "mctlass/gemm/kernel/trmm_universal.h" +#include "mctlass/gemm/kernel/default_trmm.h" +#include "mctlass/gemm/kernel/default_trmm_complex.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Diag Type for the triangular matrix + DiagType kDiagType, + /// Element type for C and D matrix operands + typename ElementC_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by TRMM + typename Operator, + /// + typename Enable = void + > +struct DefaultTrmmUniversal; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Real-valued TRMM kernels +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Diag Type for the triangular matrix + DiagType kDiagType, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by TRMM + typename Operator> +struct DefaultTrmmUniversal< + ElementA, + LayoutA, + ComplexTransform::kNone, // transform A + kAlignmentA, + ElementB, + LayoutB, + ComplexTransform::kNone, // transform B + kAlignmentB, + kSideMode, + kFillMode, + kDiagType, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + typename std::enable_if< ! mctlass::is_complex::value>::type +> { + + using DefaultTrmmKernel = typename kernel::DefaultTrmm< + ElementA, + LayoutA, + kAlignmentA, + ElementB, + LayoutB, + kAlignmentB, + kSideMode, + kFillMode, + kDiagType, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator + >::TrmmKernel; + + /// Define the kernel in terms of the default kernel + using TrmmKernel = kernel::TrmmUniversal< + typename DefaultTrmmKernel::Mma, + typename DefaultTrmmKernel::Epilogue, + ThreadblockSwizzle, + kSideMode, + kFillMode, + kDiagType + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Complex-valued TRMM kernels +// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Complex elementwise transformation on A operand + ComplexTransform TransformA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Complex elementwise transformation on B operand + ComplexTransform TransformB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Diag Type for the triangular matrix + DiagType kDiagType, + /// Element type for C and D matrix operands + typename ElementC, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Warp-level tile size (concept: GemmShape) + typename InstructionShape, + /// Epilogue output operator + typename EpilogueOutputOp, + /// Threadblock-level swizzling operator + typename ThreadblockSwizzle, + /// Number of stages used in the pipelined mainloop + int Stages, + /// If true, kernel is configured to support serial reduction in the + /// epilogue + bool SplitKSerial, + /// Operation performed by TRMM + typename Operator + > +struct DefaultTrmmUniversal< + ElementA, + LayoutA, + TransformA, + kAlignmentA, + ElementB, + LayoutB, + TransformB, + kAlignmentB, + kSideMode, + kFillMode, + kDiagType, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + SplitKSerial, + Operator, + typename std::enable_if::value>::type +> { + + using DefaultTrmmKernel = typename kernel::DefaultTrmmComplex< + ElementA, + LayoutA, + ElementB, + LayoutB, + kSideMode, + kFillMode, + kDiagType, + ElementC, + LayoutC, + ElementAccumulator, + OperatorClass, + ArchTag, + ThreadblockShape, + WarpShape, + InstructionShape, + EpilogueOutputOp, + ThreadblockSwizzle, + Stages, + TransformA, + TransformB, + Operator, + SplitKSerial + >::TrmmKernel; + + /// Define the kernel in terms of the default kernel + using TrmmKernel = kernel::TrmmUniversal< + typename DefaultTrmmKernel::Mma, + typename DefaultTrmmKernel::Epilogue, + ThreadblockSwizzle, + kSideMode, + kFillMode, + kDiagType + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/ell_gemm.h b/csrc/mctlass/include/mctlass/gemm/kernel/ell_gemm.h new file mode 100644 index 0000000..8041e8a --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/ell_gemm.h @@ -0,0 +1,830 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Template for a Block-Ell sparse gemm kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/semaphore.h" +#include "mctlass/arch/arch.h" + +#include "mctlass/transform/threadblock/ell_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + bool SplitKSerial, ///! If true, code supporting split-K via serial reduction is enabled. + bool IsASparse ///! If true, A is sparse matrix +> +struct EllGemm { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using OutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static bool const kSplitKSerial = SplitKSerial; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Parameters structure + struct Params { + mctlass::gemm::GemmCoord problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorA::TensorRef ref_A; + typename Mma::IteratorB::Params params_B; + typename Mma::IteratorB::TensorRef ref_B; + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::TensorRef ref_C; + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::OutputTileIterator::TensorRef ref_D; + typename OutputOp::Params output_op; + int *semaphore; + int gemm_k_iterations; + int gemm_k_size; + const int* ell_idx; + int ell_ncol; + int ell_blocksize; + int ell_base_idx; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): swizzle_log_tile(0), semaphore(0), gemm_k_iterations(0), gemm_k_size(0) { } + + MCTLASS_HOST_DEVICE + Params( + mctlass::gemm::GemmCoord const & problem_size, + mctlass::gemm::GemmCoord const & grid_tiled_shape, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::OutputTileIterator::TensorRef ref_C, + typename Epilogue::OutputTileIterator::TensorRef ref_D, + const int* ell_idx, + int ell_ncol, + int ell_blocksize, + int ell_base_idx, + typename OutputOp::Params output_op = typename OutputOp::Params(), + int *workspace = nullptr + ): + problem_size(problem_size), + grid_tiled_shape(grid_tiled_shape), + swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)), + params_A(ref_A.layout()), + ref_A(ref_A), + params_B(ref_B.layout()), + ref_B(ref_B), + params_C(ref_C.layout()), + ref_C(ref_C), + params_D(ref_D.layout()), + ref_D(ref_D), + output_op(output_op), + ell_idx(ell_idx), + ell_ncol(ell_ncol), + ell_blocksize(ell_blocksize), + ell_base_idx(ell_base_idx) + { + + int total_gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + int gemm_k_iterations = (total_gemm_k_iterations + grid_tiled_shape.k() - 1) / grid_tiled_shape.k(); + + gemm_k_size = gemm_k_iterations * Mma::Shape::kK; + + semaphore = workspace; + } + }; + + /// Shared memory storage structure + struct SharedStorage { + union{ + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + typename mctlass::transform::threadblock::ell::SharedStorage ell; + }; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + EllGemm() { } + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::OutputTileIterator::TensorRef ref_C, + typename Epilogue::OutputTileIterator::TensorRef ref_D) { + + static int const kAlignmentA = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + if (!TensorRef_aligned(ref_A, kAlignmentA)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_B, kAlignmentB)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_C, kAlignmentC)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_D, kAlignmentC)) { + return Status::kErrorMisalignedOperand; + } + + if ((problem_size.m() % kAlignmentA) || (problem_size.k() % kAlignmentA) || + (problem_size.n() % kAlignmentB) || (problem_size.k() % kAlignmentB) || + (problem_size.m() % kAlignmentC) || (problem_size.n() % kAlignmentC)) { + + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + int tile_in_ell_block = (params.ell_blocksize + Mma::Shape::kM - 1 ) / Mma::Shape::kM; + int ell_block_offset_m = threadblock_tile_offset.m() / tile_in_ell_block; + int tile_offset_m = threadblock_tile_offset.m() % tile_in_ell_block; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + int lane_idx = threadIdx.x % 64; + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // skip computation if matrix is 0 + if (params.ell_ncol > 0) { + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + ell_block_offset_m * params.ell_blocksize + + tile_offset_m * Mma::Shape::kM, + threadblock_tile_offset.k() * params.gemm_k_size + }; + + mctlass::MatrixCoord tb_offset_B{ + threadblock_tile_offset.k() * params.gemm_k_size, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + int ell_idx_start = + (threadblock_tile_offset.m() / tile_in_ell_block) * + (params.ell_ncol / params.ell_blocksize); + const int* ell_idx_ptr = &(params.ell_idx[ell_idx_start]); + + // Problem size is a function of threadblock index in the K dimension + int problem_size_k = min( + params.problem_size.k(), + (threadblock_tile_offset.k() + 1) * params.gemm_k_size); + problem_size_k = min(problem_size_k, params.ell_ncol); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = + (problem_size_k - tb_offset_A.column() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, + params.ref_A.data(), + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + params.params_B, + params.ref_B.data(), + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B); + + // Define coef for ELL index depending on LayoutB + int ell_stride = iterator_B.get_stride(); + + typename mctlass::transform::threadblock::ell::Iterator ell_iterator( + shared_storage.ell, + ell_idx_ptr, + params.ell_blocksize, + params.ell_base_idx, + Mma::Shape::kK, + problem_size_k, + ell_stride, + thread_idx + ); + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + if (!kSplitKSerial || gemm_k_iterations > 0) { + // check if index computations can be skipped + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + constexpr bool is_double = (sizeof(Mma::IteratorA::Element) == 8); + constexpr bool is_multiple_alignment = + (kAlignmentA > 1) && (kAlignmentB > 1) && (kAlignmentC > 1); + const bool is_specialized_blocksize = + ((params.ell_blocksize) & (params.ell_blocksize-1)) == 0 + && params.ell_blocksize >= Mma::Shape::kK; + // Compute threadblock-scoped matrix multiply-add + if ((is_double || is_multiple_alignment) && is_specialized_blocksize) { + mma.operator()( + gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators, ell_iterator); + } + else { + mma.operator()( + gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators, ell_iterator); + } + } + } // if (params.ell_ncols > 0) + + // + // Epilogue + // + + OutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + ell_block_offset_m = threadblock_tile_offset.m() / tile_in_ell_block; + tile_offset_m = threadblock_tile_offset.m() % tile_in_ell_block; + + //assume identity swizzle + MatrixCoord threadblock_offset( + ell_block_offset_m * params.ell_blocksize + + tile_offset_m * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + //avoid out of bounds + MatrixCoord threadblock_extent( + min(params.problem_size.m(), + ell_block_offset_m * params.ell_blocksize + + min((tile_offset_m + 1) * Mma::Shape::kM, params.ell_blocksize)), + min(params.problem_size.n(), + (threadblock_tile_offset.n()+1) * Mma::Shape::kN) + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + // If performing a reduction via split-K, fetch the initial synchronization + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + params.ref_C.data(), + threadblock_extent, + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + params.ref_D.data(), + threadblock_extent, + thread_idx, + threadblock_offset + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + } + + // Execute the epilogue operator to update the destination tensor. + epilogue(output_op, iterator_D, accumulators, iterator_C); + + // + // Release the semaphore + // + + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +// B is Sparse +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + bool SplitKSerial ///! If true, code supporting split-K via serial reduction is enabled. +> +struct EllGemm { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using OutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static bool const kSplitKSerial = SplitKSerial; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Parameters structure + struct Params { + mctlass::gemm::GemmCoord problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorA::TensorRef ref_A; + typename Mma::IteratorB::Params params_B; + typename Mma::IteratorB::TensorRef ref_B; + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::TensorRef ref_C; + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::OutputTileIterator::TensorRef ref_D; + typename OutputOp::Params output_op; + int *semaphore; + int gemm_k_iterations; + int gemm_k_size; + const int* ell_idx; + int ell_ncol; + int ell_blocksize; + int ell_base_idx; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): swizzle_log_tile(0), semaphore(0), gemm_k_iterations(0), gemm_k_size(0) { } + + MCTLASS_HOST_DEVICE + Params( + mctlass::gemm::GemmCoord const & problem_size, + mctlass::gemm::GemmCoord const & grid_tiled_shape, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::OutputTileIterator::TensorRef ref_C, + typename Epilogue::OutputTileIterator::TensorRef ref_D, + const int* ell_idx, + int ell_ncol, + int ell_blocksize, + int ell_base_idx, + typename OutputOp::Params output_op = typename OutputOp::Params(), + int *workspace = nullptr + ): + problem_size(problem_size), + grid_tiled_shape(grid_tiled_shape), + swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)), + params_A(ref_A.layout()), + ref_A(ref_A), + params_B(ref_B.layout()), + ref_B(ref_B), + params_C(ref_C.layout()), + ref_C(ref_C), + params_D(ref_D.layout()), + ref_D(ref_D), + output_op(output_op), + ell_idx(ell_idx), + ell_ncol(ell_ncol), + ell_blocksize(ell_blocksize), + ell_base_idx(ell_base_idx) + { + + int total_gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + int gemm_k_iterations = (total_gemm_k_iterations + grid_tiled_shape.k() - 1) / grid_tiled_shape.k(); + + gemm_k_size = gemm_k_iterations * Mma::Shape::kK; + + semaphore = workspace; + } + }; + + /// Shared memory storage structure + struct SharedStorage { + union{ + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + typename mctlass::transform::threadblock::ell::SharedStorage ell; + }; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + EllGemm() { } + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::OutputTileIterator::TensorRef ref_C, + typename Epilogue::OutputTileIterator::TensorRef ref_D) { + + static int const kAlignmentA = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + if (!TensorRef_aligned(ref_A, kAlignmentA)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_B, kAlignmentB)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_C, kAlignmentC)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_D, kAlignmentC)) { + return Status::kErrorMisalignedOperand; + } + + if ((problem_size.m() % kAlignmentA) || (problem_size.k() % kAlignmentA) || + (problem_size.n() % kAlignmentB) || (problem_size.k() % kAlignmentB) || + (problem_size.m() % kAlignmentC) || (problem_size.n() % kAlignmentC)) { + + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + int tile_in_ell_block = (params.ell_blocksize + Mma::Shape::kN - 1 ) / Mma::Shape::kN; + int ell_block_offset_n = threadblock_tile_offset.n() / tile_in_ell_block; + int tile_offset_n = threadblock_tile_offset.n() % tile_in_ell_block; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + int lane_idx = threadIdx.x % 64; + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // skip computation if matrix is 0 + if (params.ell_ncol > 0) { + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.k() * params.gemm_k_size, + }; + + mctlass::MatrixCoord tb_offset_B{ + threadblock_tile_offset.k() * params.gemm_k_size, + ell_block_offset_n * params.ell_blocksize + + tile_offset_n * Mma::Shape::kN, + }; + + int ell_idx_start = + (threadblock_tile_offset.n() / tile_in_ell_block) * + (params.ell_ncol / params.ell_blocksize); + const int* ell_idx_ptr = &(params.ell_idx[ell_idx_start]); + + // Problem size is a function of threadblock index in the K dimension + int problem_size_k = min( + params.problem_size.k(), + (threadblock_tile_offset.k() + 1) * params.gemm_k_size); + problem_size_k = min(problem_size_k, params.ell_ncol); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = + (problem_size_k - tb_offset_A.column() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, + params.ref_A.data(), + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + params.params_B, + params.ref_B.data(), + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B); + + // Define coef for ELL index depending on LayoutA + int ell_stride = iterator_A.get_stride(); + + typename mctlass::transform::threadblock::ell::Iterator ell_iterator( + shared_storage.ell, + ell_idx_ptr, + params.ell_blocksize, + params.ell_base_idx, + Mma::Shape::kK, + problem_size_k, + ell_stride, + thread_idx + ); + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + if (!kSplitKSerial || gemm_k_iterations > 0) { + // check if index computations can be skipped + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + constexpr bool is_double = (sizeof(Mma::IteratorA::Element) == 8); + constexpr bool is_multiple_alignment = + (kAlignmentA > 1) && (kAlignmentB > 1) && (kAlignmentC > 1); + const bool is_specialized_blocksize = + ((params.ell_blocksize) & (params.ell_blocksize-1)) == 0 + && params.ell_blocksize >= Mma::Shape::kK; + // Compute threadblock-scoped matrix multiply-add + if ((is_double || is_multiple_alignment) && is_specialized_blocksize) { + mma.operator()( + gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators, ell_iterator); + } + else { + mma.operator()( + gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators, ell_iterator); + } + } + } // if (params.ell_ncols > 0) + + // + // Epilogue + // + + OutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + ell_block_offset_n = threadblock_tile_offset.n() / tile_in_ell_block; + tile_offset_n = threadblock_tile_offset.n() % tile_in_ell_block; + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + ell_block_offset_n * params.ell_blocksize + + tile_offset_n * Mma::Shape::kN + ); + + //avoid out of bounds + MatrixCoord threadblock_extent( + min(params.problem_size.m(), + (threadblock_tile_offset.m()+1) * Mma::Shape::kM), + min(params.problem_size.n(), + ell_block_offset_n * params.ell_blocksize + + min((tile_offset_n + 1) * Mma::Shape::kN, params.ell_blocksize)) + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + // If performing a reduction via split-K, fetch the initial synchronization + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + params.ref_C.data(), + threadblock_extent, + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + params.ref_D.data(), + threadblock_extent, + thread_idx, + threadblock_offset + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + } + + // Execute the epilogue operator to update the destination tensor. + epilogue(output_op, iterator_D, accumulators, iterator_C); + + // + // Release the semaphore + // + + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm.h new file mode 100644 index 0000000..e9f98c5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm.h @@ -0,0 +1,381 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/semaphore.h" +#include "mctlass/arch/arch.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + bool SplitKSerial ///! If true, code supporting split-K via serial reduction is enabled. +> +struct Gemm { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using OutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static bool const kSplitKSerial = SplitKSerial; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + // static int const kThreadCount = 32 * WarpCount::kCount; + static int const kThreadCount = 64 * WarpCount::kCount; + + /// Parameters structure + struct Params { + mctlass::gemm::GemmCoord problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorA::TensorRef ref_A; + typename Mma::IteratorB::Params params_B; + typename Mma::IteratorB::TensorRef ref_B; + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::TensorRef ref_C; + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::OutputTileIterator::TensorRef ref_D; + typename OutputOp::Params output_op; + int *semaphore; + int gemm_k_size; + // For gather+scatter operations + int const *gather_A_indices; + int const *gather_B_indices; + int const *scatter_D_indices; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): swizzle_log_tile(0), semaphore(0), gemm_k_size(0) { } + + MCTLASS_HOST_DEVICE + Params( + mctlass::gemm::GemmCoord const & problem_size, + mctlass::gemm::GemmCoord const & grid_tiled_shape, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::OutputTileIterator::TensorRef ref_C, + typename Epilogue::OutputTileIterator::TensorRef ref_D, + typename OutputOp::Params output_op = typename OutputOp::Params(), + int *workspace = nullptr, + int const *gather_A_indices = nullptr, + int const *gather_B_indices = nullptr, + int const *scatter_D_indices = nullptr + ): + problem_size(problem_size), + grid_tiled_shape(grid_tiled_shape), + swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)), + params_A(ref_A.layout()), + ref_A(ref_A), + params_B(ref_B.layout()), + ref_B(ref_B), + params_C(ref_C.layout()), + ref_C(ref_C), + params_D(ref_D.layout()), + ref_D(ref_D), + output_op(output_op), + gather_A_indices(gather_A_indices), + gather_B_indices(gather_B_indices), + scatter_D_indices(scatter_D_indices) { + + int total_gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + int gemm_k_iterations = (total_gemm_k_iterations + grid_tiled_shape.k() - 1) / grid_tiled_shape.k(); + + gemm_k_size = gemm_k_iterations * Mma::Shape::kK; + + semaphore = workspace; + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Gemm() { } + + /// Determines whether kernel satisfies alignment + MCTLASS_HOST_DEVICE + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::OutputTileIterator::TensorRef ref_C, + typename Epilogue::OutputTileIterator::TensorRef ref_D) { + + static int const kAlignmentA = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Epilogue::OutputTileIterator::kElementsPerAccess; + + if (!TensorRef_aligned(ref_A, kAlignmentA)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_B, kAlignmentB)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_C, kAlignmentC)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_D, kAlignmentC)) { + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.k() * params.gemm_k_size, + }; + + mctlass::MatrixCoord tb_offset_B{ + threadblock_tile_offset.k() * params.gemm_k_size, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + // Problem size is a function of threadblock index in the K dimension + int problem_size_k = min( + params.problem_size.k(), + (threadblock_tile_offset.k() + 1) * params.gemm_k_size); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - tb_offset_A.column() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, + params.ref_A.data(), + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_A, + params.gather_A_indices); + + typename Mma::IteratorB iterator_B( + params.params_B, + params.ref_B.data(), + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B, + params.gather_B_indices); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + if (!kSplitKSerial || gemm_k_iterations > 0) { + // Compute threadblock-scoped matrix multiply-add + mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators); + } + + // + // Epilogue + // + + OutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + // If performing a reduction via split-K, fetch the initial synchronization + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + params.ref_C.data(), + params.problem_size.mn(), + thread_idx, + threadblock_offset, + params.scatter_D_indices + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + params.ref_D.data(), + params.problem_size.mn(), + thread_idx, + threadblock_offset, + params.scatter_D_indices + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + + } + + // Execute the epilogue operator to update the destination tensor. + epilogue(output_op, iterator_D, accumulators, iterator_C); + + // + // Release the semaphore + // + + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_array.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_array.h new file mode 100644 index 0000000..759784e --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_array.h @@ -0,0 +1,264 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_ ///! Threadblock swizzling function +> +struct GemmArray { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using OutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Parameters structure + struct Params { + mctlass::gemm::GemmCoord problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorA::Element const * const * ptr_A; + typename Mma::IteratorB::Params params_B; + typename Mma::IteratorB::Element const * const * ptr_B; + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::Element const * const * ptr_C; + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::OutputTileIterator::Element * const * ptr_D; + int64_t stride_D; + typename OutputOp::Params epilogue; + int batch_count; + int gemm_k_iterations; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params() : + swizzle_log_tile(0) { } + + MCTLASS_HOST_DEVICE + Params( + mctlass::gemm::GemmCoord const & problem_size_, + mctlass::gemm::GemmCoord const & grid_tiled_shape_, + typename Mma::IteratorA::Element const * const * ptr_A_, + typename Mma::IteratorA::Layout layout_A, + typename Mma::IteratorB::Element const * const * ptr_B_, + typename Mma::IteratorB::Layout layout_B, + typename Epilogue::OutputTileIterator::Element const * const * ptr_C_, + typename Epilogue::OutputTileIterator::Layout layout_C, + typename Epilogue::OutputTileIterator::Element * const * ptr_D_, + typename Epilogue::OutputTileIterator::Layout layout_D, + typename OutputOp::Params epilogue_, + int batch_count_ + ): + problem_size(problem_size_), + grid_tiled_shape(grid_tiled_shape_), + swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)), + params_A(layout_A), + ptr_A(ptr_A_), + params_B(layout_B), + ptr_B(ptr_B_), + params_C(layout_C), + ptr_C(ptr_C_), + params_D(layout_D), + ptr_D(ptr_D_), + epilogue(epilogue_), + batch_count(batch_count_), + gemm_k_iterations((problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK) { + + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + GemmArray() { } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + + // Each CTA handles multiple batch indices to accommodate limited range of CUDA grid's Z dimension + for (int batch_idx = threadblock_swizzle.get_batch_idx(); + batch_idx < params.batch_count; + batch_idx += gridDim.z) { + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + 0 + }; + + mctlass::MatrixCoord tb_offset_B{ + 0, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, + const_cast(params.ptr_A[batch_idx]), + params.problem_size.mk(), + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + params.params_B, + const_cast(params.ptr_B[batch_idx]), + params.problem_size.kn(), + thread_idx, + tb_offset_B); + + // + // Main loop + // + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + + // Compute threadblock-scoped matrix multiply-add + mma(params.gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators); + + // + // Epilogue + // + + OutputOp output_op(params.epilogue); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + // Tile iterator writing to output tile + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + const_cast(params.ptr_C[batch_idx]), + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to output tile + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + params.ptr_D[batch_idx], + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // run efficient epilogue + epilogue(output_op, iterator_D, accumulators, iterator_C); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_batched.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_batched.h new file mode 100644 index 0000000..9b0092e --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_batched.h @@ -0,0 +1,279 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_ ///! Threadblock swizzling function +> +struct GemmBatched { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using OutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Parameters structure + struct Params { + mctlass::gemm::GemmCoord problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorA::TensorRef ref_A; + int64_t stride_A; + typename Mma::IteratorB::Params params_B; + typename Mma::IteratorB::TensorRef ref_B; + int64_t stride_B; + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::TensorRef ref_C; + int64_t stride_C; + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::OutputTileIterator::TensorRef ref_D; + int64_t stride_D; + typename OutputOp::Params epilogue; + int batch_count; + int gemm_k_iterations; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params() : swizzle_log_tile(0) { } + + MCTLASS_HOST_DEVICE + Params( + mctlass::gemm::GemmCoord const & problem_size_, + mctlass::gemm::GemmCoord const & grid_tiled_shape_, + typename Mma::IteratorA::TensorRef ref_A_, + int64_t stride_A_, + typename Mma::IteratorB::TensorRef ref_B_, + int64_t stride_B_, + typename Epilogue::OutputTileIterator::TensorRef ref_C_, + int64_t stride_C_, + typename Epilogue::OutputTileIterator::TensorRef ref_D_, + int64_t stride_D_, + typename OutputOp::Params epilogue_, + int batch_count_ + ): + problem_size(problem_size_), + grid_tiled_shape(grid_tiled_shape_), + swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)), + params_A(ref_A_.layout()), + ref_A(ref_A_), + stride_A(stride_A_), + params_B(ref_B_.layout()), + ref_B(ref_B_), + stride_B(stride_B_), + params_C(ref_C_.layout()), + ref_C(ref_C_), + stride_C(stride_C_), + params_D(ref_D_.layout()), + ref_D(ref_D_), + stride_D(stride_D_), + epilogue(epilogue_), + batch_count(batch_count_), + gemm_k_iterations((problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK) { + + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + GemmBatched() { } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + + // Each CTA handles multiple batch indices to accommodate limited range of CUDA grid's Z dimension + for (int batch_idx = threadblock_swizzle.get_batch_idx(); + batch_idx < params.batch_count; + batch_idx += gridDim.z) { + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + 0 + }; + + mctlass::MatrixCoord tb_offset_B{ + 0, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, + params.ref_A.data(), + params.problem_size.mk(), + thread_idx, + tb_offset_A); + + iterator_A.add_pointer_offset(params.stride_A * batch_idx); + + typename Mma::IteratorB iterator_B( + params.params_B, + params.ref_B.data(), + params.problem_size.kn(), + thread_idx, + tb_offset_B); + + iterator_B.add_pointer_offset(params.stride_B * batch_idx); + + + // + // Main loop + // + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + + // Compute threadblock-scoped matrix multiply-add + mma(params.gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators); + + // + // Epilogue + // + + OutputOp output_op(params.epilogue); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + // Tile iterator writing to output tile + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + params.ref_C.data(), + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + iterator_C.add_pointer_offset(params.stride_C * batch_idx); + + // Tile iterator writing to output tile + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + params.ref_D.data(), + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + iterator_D.add_pointer_offset(params.stride_D * batch_idx); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // run efficient epilogue + epilogue(output_op, iterator_D, accumulators, iterator_C); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_grouped.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_grouped.h new file mode 100644 index 0000000..177cf96 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_grouped.h @@ -0,0 +1,481 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Problem visitor for grouped GEMMs +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/semaphore.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/trace.h" +#include "mctlass/gemm/kernel/gemm_transpose_operands.h" +#include "mctlass/gemm/kernel/gemm_grouped_problem_visitor.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + GroupScheduleMode GroupScheduleMode_, ///! Type of scheduling to perform + bool Transposed = false +> +struct GemmGrouped { +public: + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static GroupScheduleMode const kGroupScheduleMode = GroupScheduleMode_; + static bool const kTransposed = Transposed; + + // Optional transpose + using MapArguments = kernel::detail::MapArguments< + typename Mma::IteratorA::Element, + typename Mma::IteratorA::Layout, + Mma::kTransformA, + Mma::IteratorA::AccessType::kElements, + typename Mma::IteratorB::Element, + typename Mma::IteratorB::Layout, + Mma::kTransformB, + Mma::IteratorB::AccessType::kElements, + typename Mma::LayoutC, + kTransposed + >; + + // Public-facing type definitions related to operand element type, layout, and complex conjugate + // operation. Must interact with the 'kTransposed' notion. + using ElementA = typename MapArguments::ElementA; + using LayoutA = typename MapArguments::LayoutA; + using ElementB = typename MapArguments::ElementB; + using LayoutB = typename MapArguments::LayoutB; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename MapArguments::LayoutC; + + static ComplexTransform const kTransformA = MapArguments::kTransformA; + static ComplexTransform const kTransformB = MapArguments::kTransformB; + + // Type definitions about the mainloop. + using Operator = typename Mma::Operator; + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::InstructionShape; + using ArchTag = typename Mma::ArchTag; + + static int const kStages = Mma::kStages; + static int const kAlignmentA = MapArguments::kAlignmentA; + static int const kAlignmentB = MapArguments::kAlignmentB; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 64 * WarpCount::kCount; + + using ProblemVisitor = GemmGroupedProblemVisitor< + ThreadblockShape, + kGroupScheduleMode, + kThreadCount, + kThreadCount, + kTransposed>; + + // + // Structures + // + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord *problem_sizes; + int problem_count; + int threadblock_count; + + typename EpilogueOutputOp::Params output_op; + + ElementA ** ptr_A; + ElementB ** ptr_B; + ElementC ** ptr_C; + ElementC ** ptr_D; + + typename LayoutA::Stride::LongIndex *lda; + typename LayoutB::Stride::LongIndex *ldb; + typename LayoutC::Stride::LongIndex *ldc; + typename LayoutC::Stride::LongIndex *ldd; + + // Only used by device-level operator + GemmCoord *host_problem_sizes; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments(): + problem_count(0), + threadblock_count(0), + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + lda(nullptr), + ldb(nullptr), + ldc(nullptr), + ldd(nullptr), + host_problem_sizes(nullptr) + { + + } + + /// Ctor + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord *problem_sizes, + int problem_count, + int threadblock_count, + typename EpilogueOutputOp::Params output_op, + ElementA ** ptr_A, + ElementB ** ptr_B, + ElementC ** ptr_C, + ElementC ** ptr_D, + typename LayoutA::Stride::LongIndex *lda, + typename LayoutB::Stride::LongIndex *ldb, + typename LayoutC::Stride::LongIndex *ldc, + typename LayoutC::Stride::LongIndex *ldd, + GemmCoord *host_problem_sizes=nullptr + ): + problem_sizes(problem_sizes), + problem_count(problem_count), + threadblock_count(threadblock_count), + output_op(output_op), + ptr_A(ptr_A), + ptr_B(ptr_B), + ptr_C(ptr_C), + ptr_D(ptr_D), + lda(lda), + ldb(ldb), + ldc(ldc), + ldd(ldd), + host_problem_sizes(host_problem_sizes) + { + + } + }; + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params { + + typename ProblemVisitor::Params problem_visitor; + int threadblock_count; + + typename EpilogueOutputOp::Params output_op; + + ElementA ** ptr_A; + ElementB ** ptr_B; + ElementC ** ptr_C; + ElementC ** ptr_D; + + typename LayoutA::Stride::LongIndex *lda; + typename LayoutB::Stride::LongIndex *ldb; + typename LayoutC::Stride::LongIndex *ldc; + typename LayoutC::Stride::LongIndex *ldd; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + lda(nullptr), + ldb(nullptr), + ldc(nullptr), + ldd(nullptr) + { } + + MCTLASS_HOST_DEVICE + Params(Arguments const &args, + void *workspace = nullptr, + int tile_count = 0): + problem_visitor(args.problem_sizes, args.problem_count, workspace, tile_count), + threadblock_count(args.threadblock_count), + output_op(args.output_op), + ptr_A(args.ptr_A), + ptr_B(args.ptr_B), + ptr_C(args.ptr_C), + ptr_D(args.ptr_D), + lda(args.lda), + ldb(args.ldb), + ldc(args.ldc), + ldd(args.ldd) + { + + } + + MCTLASS_HOST_DEVICE + void update( + Arguments const &args, + void *workspace = nullptr, + int tile_count = 0) { + + problem_visitor = typename ProblemVisitor::Params(args.problem_sizes, args.problem_count, + workspace, tile_count); + threadblock_count = args.threadblock_count; + output_op = args.output_op; + ptr_A = args.ptr_A; + ptr_B = args.ptr_B; + ptr_C = args.ptr_C; + ptr_D = args.ptr_D; + lda = args.lda; + ldb = args.ldb; + ldc = args.ldc; + ldd = args.ldd; + } + }; + + /// Shared memory storage structure + struct SharedStorage { + union { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + } kernel; + + // ProblemVisitor shared storage can't be overlapped with others + typename ProblemVisitor::SharedStorage problem_visitor; + }; + +public: + + // + // Methods + // + + MCTLASS_DEVICE + GemmGrouped() { } + + /// Determines whether kernel satisfies alignment + static Status can_implement(mctlass::gemm::GemmCoord const & problem_size) { + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return Status::kSuccess; + } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // + // These types shadow the type-level definitions and support the ability to implement + // a 'transposed' GEMM that computes the transposed problems. + // + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + + // + // Problem visitor. + // + ProblemVisitor problem_visitor( + params.problem_visitor, + shared_storage.problem_visitor, + blockIdx.x); + + // Outer 'persistent' loop to iterate over tiles + while (problem_visitor.next_tile()) { + + GemmCoord problem_size = problem_visitor.problem_size(); + int32_t problem_idx = problem_visitor.problem_index(); + int32_t threadblock_idx = int32_t(problem_visitor.threadblock_idx()); + + GemmCoord grid_shape = problem_visitor.grid_shape(problem_size); + + mctlass::gemm::GemmCoord threadblock_offset( + int(threadblock_idx / grid_shape.n()) * Mma::Shape::kM, + int(threadblock_idx % grid_shape.n()) * Mma::Shape::kN, + 0); + + // Load element pointers. Exchange pointers and strides if working on the transpose + ElementA *ptr_A = reinterpret_cast((kTransposed ? params.ptr_B[problem_idx] : params.ptr_A[problem_idx])); + typename LayoutA::LongIndex ldm_A = (kTransposed ? params.ldb[problem_idx] : params.lda[problem_idx]); + + ElementB *ptr_B = reinterpret_cast((kTransposed ? params.ptr_A[problem_idx] : params.ptr_B[problem_idx])); + typename LayoutB::LongIndex ldm_B = (kTransposed ? params.lda[problem_idx] : params.ldb[problem_idx]); + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_offset.m(), + 0, + }; + + mctlass::MatrixCoord tb_offset_B{ + 0, + threadblock_offset.n() + }; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + LayoutA(ldm_A), + ptr_A, + {problem_size.m(), problem_size.k()}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + LayoutB(ldm_B), + ptr_B, + {problem_size.k(), problem_size.n()}, + thread_idx, + tb_offset_B); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + // + // Matrix multiply phase + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.kernel.main_loop, thread_idx, warp_idx, lane_idx); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Wait for all threads to finish their epilogue phases from the previous tile. + __syncthreads(); + + // Compute threadblock-scoped matrix multiply-add + mma( + gemm_k_iterations, + accumulators, + iterator_A, + iterator_B, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + ElementC *ptr_C = params.ptr_C[problem_idx]; + ElementC *ptr_D = params.ptr_D[problem_idx]; + + LayoutC layout_C(params.ldc[problem_idx]); + LayoutC layout_D(params.ldd[problem_idx]); + + typename Epilogue::OutputTileIterator::Params params_C(layout_C); + typename Epilogue::OutputTileIterator::Params params_D(layout_D); + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params_C, + ptr_C, + problem_size.mn(), + thread_idx, + threadblock_offset.mn() + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params_D, + ptr_D, + problem_size.mn(), + thread_idx, + threadblock_offset.mn() + ); + + Epilogue epilogue( + shared_storage.kernel.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Execute the epilogue operator to update the destination tensor. + epilogue( + output_op, + iterator_D, + accumulators, + iterator_C); + + // Next tile + problem_visitor.advance(gridDim.x); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_grouped_problem_visitor.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_grouped_problem_visitor.h new file mode 100644 index 0000000..c982c0f --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_grouped_problem_visitor.h @@ -0,0 +1,122 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Scheduler for grouped GEMM +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/gemm/kernel/grouped_problem_visitor.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { +// Helper for correctly representing problem sizes in grouped kernels +template < + typename ThreadblockShape, + bool Transposed +> +struct GemmGroupedProblemSizeHelper { + + static bool const kTransposed = Transposed; + + MCTLASS_HOST_DEVICE + static mctlass::gemm::GemmCoord grid_shape(const mctlass::gemm::GemmCoord& problem) { + return mctlass::gemm::GemmCoord( + ((problem.m() - 1 + ThreadblockShape::kM) / ThreadblockShape::kM), + ((problem.n() - 1 + ThreadblockShape::kN) / ThreadblockShape::kN), + 1); + } + + MCTLASS_HOST_DEVICE + static void possibly_transpose_problem(mctlass::gemm::GemmCoord& problem) { + if (kTransposed) { + swap(problem.m(), problem.n()); + } + } + + MCTLASS_HOST_DEVICE + static int32_t tile_count(const mctlass::gemm::GemmCoord& grid) { + return grid.m() * grid.n(); + } +}; + +} // namespace detail + +/// Visitor class to abstract away the algorithm for iterating over tiles +template +struct GemmGroupedProblemVisitor : public GroupedProblemVisitor< + detail::GemmGroupedProblemSizeHelper, + ThreadblockShape, + GroupScheduleMode_, + PrefetchTileCount, + ThreadCount> { + + static bool const kTransposed = Transposed; + + using ProblemSizeHelper = detail::GemmGroupedProblemSizeHelper; + using Base = GroupedProblemVisitor; + using Params = typename Base::Params; + using SharedStorage = typename Base::SharedStorage; + + // + // Methods + // + MCTLASS_DEVICE + GemmGroupedProblemVisitor( + Params const ¶ms_, + SharedStorage &shared_storage_, + int32_t block_idx + ): Base (params_, shared_storage_, block_idx) + {} +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_grouped_softmax_mainloop_fusion.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_grouped_softmax_mainloop_fusion.h new file mode 100644 index 0000000..3cf062e --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_grouped_softmax_mainloop_fusion.h @@ -0,0 +1,510 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Problem visitor for grouped GEMMs with a softmax fused beforehand +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/semaphore.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/trace.h" +#include "mctlass/gemm/kernel/gemm_transpose_operands.h" +#include "mctlass/gemm/kernel/gemm_grouped_problem_visitor.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + GroupScheduleMode GroupScheduleMode_, ///! Type of scheduling to perform + bool Transposed = false +> +struct GemmGroupedSoftmaxMainloopFusion { +public: + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static GroupScheduleMode const kGroupScheduleMode = GroupScheduleMode_; + static bool const kTransposed = Transposed; + + // Optional transpose + using MapArguments = kernel::detail::MapArguments< + typename Mma::IteratorA::Element, + typename Mma::IteratorA::Layout, + Mma::kTransformA, + Mma::IteratorA::AccessType::kElements, + typename Mma::IteratorB::Element, + typename Mma::IteratorB::Layout, + Mma::kTransformB, + Mma::IteratorB::AccessType::kElements, + typename Mma::LayoutC, + kTransposed + >; + + // Public-facing type definitions related to operand element type, layout, and complex conjugate + // operation. Must interact with the 'kTransposed' notion. + using ElementA = typename MapArguments::ElementA; + using LayoutA = typename MapArguments::LayoutA; + using ElementB = typename MapArguments::ElementB; + using LayoutB = typename MapArguments::LayoutB; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename MapArguments::LayoutC; + + using ElementScaleBias = typename Mma::IteratorNormSum::Element; + + static ComplexTransform const kTransformA = MapArguments::kTransformA; + static ComplexTransform const kTransformB = MapArguments::kTransformB; + + // Type definitions about the mainloop. + using Operator = typename Mma::Operator; + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::InstructionShape; + using ArchTag = typename Mma::ArchTag; + + static int const kStages = Mma::kStages; + static int const kAlignmentA = MapArguments::kAlignmentA; + static int const kAlignmentB = MapArguments::kAlignmentB; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + using ProblemVisitor = GemmGroupedProblemVisitor< + ThreadblockShape, + kGroupScheduleMode, + kThreadCount, + kThreadCount, + kTransposed>; + + // + // Structures + // + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmCoord *problem_sizes; + int problem_count; + int threadblock_count; + + typename EpilogueOutputOp::Params output_op; + + ElementA ** ptr_A; + ElementB ** ptr_B; + ElementC ** ptr_C; + ElementC ** ptr_D; + void ** ptr_norm; + void ** ptr_sum; + + typename LayoutA::Stride::LongIndex *lda; + typename LayoutB::Stride::LongIndex *ldb; + typename LayoutC::Stride::LongIndex *ldc; + typename LayoutC::Stride::LongIndex *ldd; + + // Only used by device-level operator + GemmCoord *host_problem_sizes; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments(): + problem_count(0), + threadblock_count(0), + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + ptr_norm(nullptr), + ptr_sum(nullptr), + lda(nullptr), + ldb(nullptr), + ldc(nullptr), + ldd(nullptr), + host_problem_sizes(nullptr) + { + + } + + /// Ctor + MCTLASS_HOST_DEVICE + Arguments( + GemmCoord *problem_sizes, + int problem_count, + int threadblock_count, + typename EpilogueOutputOp::Params output_op, + ElementA ** ptr_A, + ElementB ** ptr_B, + ElementC ** ptr_C, + ElementC ** ptr_D, + void ** ptr_norm, + void ** ptr_sum, + typename LayoutA::Stride::LongIndex *lda, + typename LayoutB::Stride::LongIndex *ldb, + typename LayoutC::Stride::LongIndex *ldc, + typename LayoutC::Stride::LongIndex *ldd, + GemmCoord *host_problem_sizes=nullptr + ): + problem_sizes(problem_sizes), + problem_count(problem_count), + threadblock_count(threadblock_count), + output_op(output_op), + ptr_A(ptr_A), + ptr_B(ptr_B), + ptr_C(ptr_C), + ptr_D(ptr_D), + ptr_norm(ptr_norm), + ptr_sum(ptr_sum), + lda(lda), + ldb(ldb), + ldc(ldc), + ldd(ldd), + host_problem_sizes(host_problem_sizes) + { + + } + }; + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params { + + typename ProblemVisitor::Params problem_visitor; + int threadblock_count; + + typename EpilogueOutputOp::Params output_op; + + ElementA ** ptr_A; + ElementB ** ptr_B; + ElementC ** ptr_C; + ElementC ** ptr_D; + + void ** ptr_norm; + void ** ptr_sum; + + typename LayoutA::Stride::LongIndex *lda; + typename LayoutB::Stride::LongIndex *ldb; + typename LayoutC::Stride::LongIndex *ldc; + typename LayoutC::Stride::LongIndex *ldd; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + ptr_norm(nullptr), + ptr_sum(nullptr), + lda(nullptr), + ldb(nullptr), + ldc(nullptr), + ldd(nullptr) + { } + + MCTLASS_HOST_DEVICE + Params(Arguments const &args, + void *workspace = nullptr, + int tile_count = 0): + problem_visitor(args.problem_sizes, args.problem_count, workspace, tile_count), + threadblock_count(args.threadblock_count), + output_op(args.output_op), + ptr_A(args.ptr_A), + ptr_B(args.ptr_B), + ptr_C(args.ptr_C), + ptr_D(args.ptr_D), + ptr_norm(args.ptr_norm), + ptr_sum(args.ptr_sum), + lda(args.lda), + ldb(args.ldb), + ldc(args.ldc), + ldd(args.ldd) + { + + } + + MCTLASS_HOST_DEVICE + void update( + Arguments const &args, + void *workspace = nullptr, + int tile_count = 0) { + + problem_visitor = typename ProblemVisitor::Params(args.problem_sizes, args.problem_count, + workspace, tile_count); + threadblock_count = args.threadblock_count; + output_op = args.output_op; + ptr_A = args.ptr_A; + ptr_B = args.ptr_B; + ptr_C = args.ptr_C; + ptr_D = args.ptr_D; + ptr_norm = args.ptr_norm; + ptr_sum = args.ptr_sum; + lda = args.lda; + ldb = args.ldb; + ldc = args.ldc; + ldd = args.ldd; + } + }; + + /// Shared memory storage structure + struct SharedStorage { + union { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + } kernel; + + // ProblemVisitor shared storage can't be overlapped with others + typename ProblemVisitor::SharedStorage problem_visitor; + }; + +public: + + // + // Methods + // + + MCTLASS_DEVICE + GemmGroupedSoftmaxMainloopFusion() { } + + /// Determines whether kernel satisfies alignment + static Status can_implement(mctlass::gemm::GemmCoord const & problem_size) { + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return Status::kSuccess; + } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // + // These types shadow the type-level definitions and support the ability to implement + // a 'transposed' GEMM that computes the transposed problems. + // + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + + // + // Problem visitor. + // + ProblemVisitor problem_visitor( + params.problem_visitor, + shared_storage.problem_visitor, + blockIdx.x); + + // Outer 'persistent' loop to iterate over tiles + while (problem_visitor.next_tile()) { + + GemmCoord problem_size = problem_visitor.problem_size(); + int32_t problem_idx = problem_visitor.problem_index(); + int32_t threadblock_idx = int32_t(problem_visitor.threadblock_idx()); + + GemmCoord grid_shape = problem_visitor.grid_shape(problem_size); + + mctlass::gemm::GemmCoord threadblock_offset( + int(threadblock_idx / grid_shape.n()) * Mma::Shape::kM, + int(threadblock_idx % grid_shape.n()) * Mma::Shape::kN, + 0); + + // Load element pointers. Exchange pointers and strides if working on the transpose + ElementA *ptr_A = reinterpret_cast((kTransposed ? params.ptr_B[problem_idx] : params.ptr_A[problem_idx])); + typename LayoutA::LongIndex ldm_A = (kTransposed ? params.ldb[problem_idx] : params.lda[problem_idx]); + + ElementB *ptr_B = reinterpret_cast((kTransposed ? params.ptr_A[problem_idx] : params.ptr_B[problem_idx])); + typename LayoutB::LongIndex ldm_B = (kTransposed ? params.lda[problem_idx] : params.ldb[problem_idx]); + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_offset.m(), + 0, + }; + + mctlass::MatrixCoord tb_offset_B{ + 0, + threadblock_offset.n() + }; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + LayoutA(ldm_A), + ptr_A, + {problem_size.m(), problem_size.k()}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + LayoutB(ldm_B), + ptr_B, + {problem_size.k(), problem_size.n()}, + thread_idx, + tb_offset_B); + + // Construct iterator to the softmax norm/sum vector + typename Mma::IteratorNormSum iterator_norm_sum( + problem_size.m(), + static_cast(params.ptr_norm[problem_idx]), + static_cast(params.ptr_sum[problem_idx]), + thread_idx, + MatrixCoord(0, threadblock_offset.m()) + ); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + // + // Matrix multiply phase + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.kernel.main_loop, thread_idx, warp_idx, lane_idx); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Wait for all threads to finish their epilogue phases from the previous tile. + __syncthreads(); + + // Compute threadblock-scoped matrix multiply-add + mma( + gemm_k_iterations, + accumulators, + iterator_A, + iterator_B, + iterator_norm_sum, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + ElementC *ptr_C = params.ptr_C[problem_idx]; + ElementC *ptr_D = params.ptr_D[problem_idx]; + + LayoutC layout_C(params.ldc[problem_idx]); + LayoutC layout_D(params.ldd[problem_idx]); + + typename Epilogue::OutputTileIterator::Params params_C(layout_C); + typename Epilogue::OutputTileIterator::Params params_D(layout_D); + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params_C, + ptr_C, + problem_size.mn(), + thread_idx, + threadblock_offset.mn() + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params_D, + ptr_D, + problem_size.mn(), + thread_idx, + threadblock_offset.mn() + ); + + Epilogue epilogue( + shared_storage.kernel.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Execute the epilogue operator to update the destination tensor. + epilogue( + output_op, + iterator_D, + accumulators, + iterator_C); + + // Next tile + problem_visitor.advance(gridDim.x); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_layernorm_mainloop_fusion.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_layernorm_mainloop_fusion.h new file mode 100644 index 0000000..d383d92 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_layernorm_mainloop_fusion.h @@ -0,0 +1,785 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Template for a multistage GEMM kernel with layernorm operations fused in mainloop. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/semaphore.h" +#include "mctlass/gemm/kernel/params_universal_base.h" + +#include "mctlass/layout/matrix.h" + +#include "mctlass/trace.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_ ///! Threadblock swizzling function +> +struct GemmLayernormMainloopFusion { +public: + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + + using ElementScaleBias = typename Mma::IteratorVarMean::Element; + using LayoutScaleBias = typename Mma::IteratorVarMean::Layout; + + static ComplexTransform const kTransformA = Mma::kTransformA; + static ComplexTransform const kTransformB = Mma::kTransformB; + using Operator = typename Mma::Operator; + + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::InstructionShape; + using ArchTag = typename Mma::ArchTag; + + static int const kStages = Mma::kStages; + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Split-K preserves splits that are 128b aligned + static int const kSplitKAlignment = const_max(128 / sizeof_bits::value, 128 / sizeof_bits::value); + + // + // Structures + // + + /// Argument structure + struct Arguments : UniversalArgumentsBase + { + // + // Data members + // + + typename EpilogueOutputOp::Params epilogue; + + void const * ptr_A; + void const * ptr_B; + void const * ptr_var; + void const * ptr_mean; + void const * ptr_gamma; + void const * ptr_beta; + void const * ptr_C; + void * ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_var; + int64_t batch_stride_mean; + int64_t batch_stride_gamma; + int64_t batch_stride_beta; + int64_t batch_stride_C; + + typename LayoutA::Stride stride_a; + typename LayoutB::Stride stride_b; + typename LayoutScaleBias::Stride stride_var; + typename LayoutScaleBias::Stride stride_mean; + typename LayoutScaleBias::Stride stride_gamma; + typename LayoutScaleBias::Stride stride_beta; + typename LayoutC::Stride stride_c; + typename LayoutC::Stride stride_d; + + typename LayoutA::Stride::LongIndex lda; + typename LayoutB::Stride::LongIndex ldb; + typename LayoutScaleBias::Stride::LongIndex ld_var; + typename LayoutScaleBias::Stride::LongIndex ld_mean; + typename LayoutScaleBias::Stride::LongIndex ld_gamma; + typename LayoutScaleBias::Stride::LongIndex ld_beta; + typename LayoutC::Stride::LongIndex ldc; + typename LayoutC::Stride::LongIndex ldd; + + int const * ptr_gather_A_indices; + int const * ptr_gather_B_indices; + int const * ptr_scatter_D_indices; + + // + // Methods + // + + Arguments(): + ptr_A(nullptr), ptr_B(nullptr), ptr_C(nullptr), ptr_D(nullptr), + ptr_var(nullptr), ptr_mean(nullptr), + ptr_gamma(nullptr), ptr_beta(nullptr), + ptr_gather_A_indices(nullptr), + ptr_gather_B_indices(nullptr), + ptr_scatter_D_indices(nullptr) + {} + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void const * ptr_var, + void const * ptr_mean, + void const * ptr_gamma, + void const * ptr_beta, + void const * ptr_C, + void * ptr_D, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_var, + int64_t batch_stride_mean, + int64_t batch_stride_gamma, + int64_t batch_stride_beta, + int64_t batch_stride_C, + int64_t batch_stride_D, + typename LayoutA::Stride stride_a, + typename LayoutB::Stride stride_b, + typename LayoutScaleBias::Stride stride_var, + typename LayoutScaleBias::Stride stride_mean, + typename LayoutScaleBias::Stride stride_gamma, + typename LayoutScaleBias::Stride stride_beta, + typename LayoutC::Stride stride_c, + typename LayoutC::Stride stride_d, + int const *ptr_gather_A_indices = nullptr, + int const *ptr_gather_B_indices = nullptr, + int const *ptr_scatter_D_indices = nullptr) + : + UniversalArgumentsBase(mode, problem_size, batch_count, batch_stride_D), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), + ptr_var(ptr_var), ptr_mean(ptr_mean), + ptr_gamma(ptr_gamma), ptr_beta(ptr_beta), + batch_stride_A(batch_stride_A), batch_stride_B(batch_stride_B), batch_stride_C(batch_stride_C), + batch_stride_var(batch_stride_var), batch_stride_mean(batch_stride_mean), + batch_stride_gamma(batch_stride_gamma), batch_stride_beta(batch_stride_beta), + lda(0), ldb(0), ldc(0), ldd(0), + ld_var(0), ld_mean(0), + ld_gamma(0), ld_beta(0), + stride_a(stride_a), stride_b(stride_b), stride_c(stride_c), stride_d(stride_d), + stride_var(stride_var), stride_mean(stride_mean), + stride_gamma(stride_gamma), stride_beta(stride_beta), + ptr_gather_A_indices(ptr_gather_A_indices), ptr_gather_B_indices(ptr_gather_B_indices), + ptr_scatter_D_indices(ptr_scatter_D_indices) + { + MCTLASS_TRACE_HOST("GemmUniversal::Arguments::Arguments() - problem_size: " << problem_size); + } + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void const * ptr_var, + void const * ptr_mean, + void const * ptr_gamma, + void const * ptr_beta, + void const * ptr_C, + void * ptr_D, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_var, + int64_t batch_stride_mean, + int64_t batch_stride_gamma, + int64_t batch_stride_beta, + int64_t batch_stride_C, + int64_t batch_stride_D, + typename LayoutA::Stride::LongIndex lda, + typename LayoutB::Stride::LongIndex ldb, + typename LayoutScaleBias::Stride::LongIndex ld_var, + typename LayoutScaleBias::Stride::LongIndex ld_mean, + typename LayoutScaleBias::Stride::LongIndex ld_gamma, + typename LayoutScaleBias::Stride::LongIndex ld_beta, + typename LayoutC::Stride::LongIndex ldc, + typename LayoutC::Stride::LongIndex ldd, + int const *ptr_gather_A_indices = nullptr, + int const *ptr_gather_B_indices = nullptr, + int const *ptr_scatter_D_indices = nullptr) + : + UniversalArgumentsBase(mode, problem_size, batch_count, batch_stride_D), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), + ptr_var(ptr_var), ptr_mean(ptr_mean), + ptr_gamma(ptr_gamma), ptr_beta(ptr_beta), + batch_stride_A(batch_stride_A), batch_stride_B(batch_stride_B), batch_stride_C(batch_stride_C), + batch_stride_var(batch_stride_var), batch_stride_mean(batch_stride_mean), + batch_stride_gamma(batch_stride_gamma), batch_stride_beta(batch_stride_beta), + lda(lda), ldb(ldb), ldc(ldc), ldd(ldd), + ld_var(ld_var), ld_mean(ld_mean), + ld_gamma(ld_gamma), ld_beta(ld_beta), + ptr_gather_A_indices(ptr_gather_A_indices), ptr_gather_B_indices(ptr_gather_B_indices), + ptr_scatter_D_indices(ptr_scatter_D_indices) + { + stride_a = make_Coord(lda); + stride_b = make_Coord(ldb); + stride_c = make_Coord(ldc); + stride_d = make_Coord(ldd); + stride_var = make_Coord(ld_var); + stride_mean = make_Coord(ld_mean); + stride_gamma = make_Coord(ld_gamma); + stride_beta = make_Coord(ld_beta); + MCTLASS_TRACE_HOST("GemmUniversal::Arguments::Arguments() - problem_size: " << problem_size); + } + + /// Returns arguments for the transposed problem + Arguments transposed_problem() const { + Arguments args(*this); + + std::swap(args.problem_size.m(), args.problem_size.n()); + std::swap(args.ptr_A, args.ptr_B); + std::swap(args.lda, args.ldb); + std::swap(args.stride_a, args.stride_b); + std::swap(args.batch_stride_A, args.batch_stride_B); + std::swap(args.ptr_gather_A_indices, args.ptr_gather_B_indices); + + return args; + } + }; + + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params : UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC> + { + using ParamsBase = UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC>; + + // + // Data members + // + + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorB::Params params_B; + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::Params params_D; + + typename EpilogueOutputOp::Params output_op; + + void * ptr_A; + void * ptr_B; + void * ptr_var; + void * ptr_mean; + void * ptr_gamma; + void * ptr_beta; + void * ptr_C; + void * ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_var; + int64_t batch_stride_mean; + int64_t batch_stride_gamma; + int64_t batch_stride_beta; + int64_t batch_stride_C; + + int * ptr_gather_A_indices; + int * ptr_gather_B_indices; + int * ptr_scatter_D_indices; + + // + // Host dispatch API + // + + /// Default constructor + Params() = default; + + /// Constructor + Params( + Arguments const &args, /// GEMM application arguments + int device_sms, /// Number of SMs on the device + int sm_occupancy) /// Kernel SM occupancy (in thread blocks) + : + ParamsBase(args, device_sms, sm_occupancy), + params_A(args.lda ? make_Coord_with_padding(args.lda) : args.stride_a), + params_B(args.ldb ? make_Coord_with_padding(args.ldb) : args.stride_b), + params_C(args.ldc ? make_Coord_with_padding(args.ldc) : args.stride_c), + params_D(args.ldd ? make_Coord_with_padding(args.ldd) : args.stride_d), + output_op(args.epilogue), + ptr_A(const_cast(args.ptr_A)), + ptr_B(const_cast(args.ptr_B)), + ptr_var(const_cast(args.ptr_var)), + ptr_mean(const_cast(args.ptr_mean)), + ptr_gamma(const_cast(args.ptr_gamma)), + ptr_beta(const_cast(args.ptr_beta)), + ptr_C(const_cast(args.ptr_C)), + ptr_D(args.ptr_D), + batch_stride_A(args.batch_stride_A), + batch_stride_B(args.batch_stride_B), + batch_stride_var(args.batch_stride_var), + batch_stride_mean(args.batch_stride_mean), + batch_stride_gamma(args.batch_stride_gamma), + batch_stride_beta(args.batch_stride_beta), + batch_stride_C(args.batch_stride_C), + ptr_gather_A_indices(const_cast(args.ptr_gather_A_indices)), + ptr_gather_B_indices(const_cast(args.ptr_gather_B_indices)), + ptr_scatter_D_indices(const_cast(args.ptr_scatter_D_indices)) + {} + + /// Lightweight update given a subset of arguments. + void update(Arguments const &args) + { + ptr_A = const_cast(args.ptr_A); + ptr_B = const_cast(args.ptr_B); + ptr_var = const_cast(args.ptr_var); + ptr_mean = const_cast(args.ptr_mean); + ptr_gamma = const_cast(args.ptr_gamma); + ptr_beta = const_cast(args.ptr_beta); + ptr_C = const_cast(args.ptr_C); + ptr_D = args.ptr_D; + + batch_stride_A = args.batch_stride_A; + batch_stride_B = args.batch_stride_B; + batch_stride_C = args.batch_stride_C; + batch_stride_var = args.batch_stride_var; + batch_stride_mean = args.batch_stride_mean; + batch_stride_gamma = args.batch_stride_gamma; + batch_stride_beta = args.batch_stride_beta; + this->batch_stride_D = args.batch_stride_D; + + ptr_gather_A_indices = const_cast(args.ptr_gather_A_indices); + ptr_gather_B_indices = const_cast(args.ptr_gather_B_indices); + ptr_scatter_D_indices = const_cast(args.ptr_scatter_D_indices); + + output_op = args.epilogue; + + MCTLASS_TRACE_HOST("GemmUniversal::Params::update()"); + } + }; + + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + +public: + + // + // Host dispatch API + // + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size) { + + MCTLASS_TRACE_HOST("GemmUniversal::can_implement()"); + + static int const kAlignmentA = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Epilogue::OutputTileIterator::kElementsPerAccess; + + bool isAMisaligned = false; + bool isBMisaligned = false; + bool isCMisaligned = false; + + if (platform::is_same::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } else if (platform::is_same::value) { + isAMisaligned = problem_size.m() % kAlignmentA; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } + + if (platform::is_same::value) { + isBMisaligned = problem_size.n() % kAlignmentB; + } else if (platform::is_same::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } + + if (platform::is_same::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } else if (platform::is_same::value) { + isCMisaligned = problem_size.m() % kAlignmentC; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } + + if (isAMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for A operand"); + return Status::kErrorMisalignedOperand; + } + + if (isBMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for B operand"); + return Status::kErrorMisalignedOperand; + } + + if (isCMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for C operand"); + return Status::kErrorMisalignedOperand; + } + + MCTLASS_TRACE_HOST(" returning kSuccess"); + + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + +public: + + // + // Device-only API + // + + // Factory invocation + MCTLASS_DEVICE + static void invoke( + Params const ¶ms, + SharedStorage &shared_storage) + { + GemmLayernormMainloopFusion op; + op(params, shared_storage); + } + + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + int offset_k = 0; + int problem_size_k = params.problem_size.k(); + + ElementA *ptr_A = static_cast(params.ptr_A); + ElementB *ptr_B = static_cast(params.ptr_B); + + // + // Fetch pointers based on mode. + // + if (params.mode == GemmUniversalMode::kGemm || + params.mode == GemmUniversalMode::kGemmSplitKParallel) { + + if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) { + + problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size; + } + + offset_k = threadblock_tile_offset.k() * params.gemm_k_size; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_A += threadblock_tile_offset.k() * params.batch_stride_A; + ptr_B += threadblock_tile_offset.k() * params.batch_stride_B; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_A = static_cast(params.ptr_A)[threadblock_tile_offset.k()]; + ptr_B = static_cast(params.ptr_B)[threadblock_tile_offset.k()]; + } + + __syncthreads(); + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + offset_k, + }; + + mctlass::MatrixCoord tb_offset_B{ + offset_k, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, + ptr_A, + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_A, + params.ptr_gather_A_indices); + + typename Mma::IteratorB iterator_B( + params.params_B, + ptr_B, + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B, + params.ptr_gather_B_indices); + + // Construct iterators to A var/mean vector + typename Mma::IteratorVarMean iterator_var_mean( + params.problem_size.m(), + static_cast(params.ptr_var), + static_cast(params.ptr_mean), + thread_idx, + MatrixCoord(0, (threadblock_tile_offset.m() * Mma::Shape::kM)) + ); + + // Construct iterators to A scale/bias vector + typename Mma::IteratorGammaBeta iterator_gamma_beta( + problem_size_k, + static_cast(params.ptr_gamma), + static_cast(params.ptr_beta), + thread_idx, + MatrixCoord( + 0, (threadblock_tile_offset.k() * Mma::Shape::kK) + ) + ); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma( + gemm_k_iterations, + accumulators, + iterator_A, + iterator_B, + iterator_var_mean, + iterator_gamma_beta, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + ElementC *ptr_C = static_cast(params.ptr_C); + ElementC *ptr_D = static_cast(params.ptr_D); + + // + // Fetch pointers based on mode. + // + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + if (params.mode == GemmUniversalMode::kGemm) { + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + } + else if (params.mode == GemmUniversalMode::kGemmSplitKParallel) { + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_C += threadblock_tile_offset.k() * params.batch_stride_C; + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_C = static_cast(params.ptr_C)[threadblock_tile_offset.k()]; + ptr_D = static_cast(params.ptr_D)[threadblock_tile_offset.k()]; + } + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + ptr_C, + params.problem_size.mn(), + thread_idx, + threadblock_offset, + params.ptr_scatter_D_indices + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.problem_size.mn(), + thread_idx, + threadblock_offset, + params.ptr_scatter_D_indices + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + } + + // Execute the epilogue operator to update the destination tensor. + epilogue( + output_op, + iterator_D, + accumulators, + iterator_C); + + // + // Release the semaphore + // + + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_params.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_params.h new file mode 100644 index 0000000..e8ec5d7 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_params.h @@ -0,0 +1,199 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/semaphore.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" +#include "mctlass/epilogue/threadblock/predicated_tile_iterator_params.h" +#include "mctlass/transform/threadblock/predicated_tile_access_iterator_params.h" + +#include "mctlass/trace.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +struct GemmParams { + + // + // Type definitions + // + using Index = int32_t; + using LongIndex = int64_t; + + using MmaIteratorParams = typename mctlass::transform::threadblock::PredicatedTileAccessIteratorParams; + using EpilogueIteratorParams = typename mctlass::epilogue::threadblock::PredicatedTileIteratorParams; + + // + // Data members + // + + mctlass::gemm::GemmCoord problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + + // Data members for Mma::Iterator::Params + MmaIteratorParams params_itr_a; + MmaIteratorParams params_itr_b; + + // Data member for Epilogue::OutputTileIterator::Params + EpilogueIteratorParams params_itr_c; + EpilogueIteratorParams params_itr_d; + + + GemmUniversalMode mode; + int batch_count; + int gemm_k_size; + + void * ptr_A; + void * ptr_B; + void * ptr_C; + void * ptr_D; + + LongIndex lda; + LongIndex ldb; + LongIndex ldc; + LongIndex ldd; + + LongIndex batch_stride_A; + LongIndex batch_stride_B; + LongIndex batch_stride_C; + LongIndex batch_stride_D; + + int *semaphore; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + GemmParams() {} + + MCTLASS_HOST_DEVICE + GemmParams( + mctlass::gemm::GemmCoord problem_size_, + mctlass::gemm::GemmCoord grid_tiled_shape_, + int swizzle_log_tile_, + GemmUniversalMode mode_, + int batch_count_, + int gemm_k_size_, + void const * ptr_A_, + void const * ptr_B_, + void const * ptr_C_, + void * ptr_D_, + LongIndex lda_, + LongIndex ldb_, + LongIndex ldc_, + LongIndex ldd_, + int64_t batch_stride_A_, + int64_t batch_stride_B_, + int64_t batch_stride_C_, + int64_t batch_stride_D_, + MmaIteratorParams const & params_itr_a_, + MmaIteratorParams const & params_itr_b_, + EpilogueIteratorParams const & params_itr_c_, + EpilogueIteratorParams const & params_itr_d_, + void *workspace_ = nullptr) : + problem_size(problem_size_), + grid_tiled_shape(grid_tiled_shape_), + swizzle_log_tile(swizzle_log_tile_), + mode(mode_), + batch_count(batch_count_), + gemm_k_size(gemm_k_size_), + ptr_A(const_cast(ptr_A_)), + ptr_B(const_cast(ptr_B_)), + ptr_C(const_cast(ptr_C_)), + ptr_D(ptr_D_), + lda(lda_), + ldb(ldb_), + ldc(ldc_), + ldd(ldd_), + batch_stride_A(batch_stride_A_), + batch_stride_B(batch_stride_B_), + batch_stride_C(batch_stride_C_), + batch_stride_D(batch_stride_D_), + params_itr_a(params_itr_a_), + params_itr_b(params_itr_b_), + params_itr_c(params_itr_c_), + params_itr_d(params_itr_d_), + semaphore(static_cast(workspace_) + ) { } + + + MCTLASS_HOST_DEVICE + void update( + void const * ptr_A_, + void const * ptr_B_, + void const * ptr_C_, + void * ptr_D_, + int64_t batch_stride_A_, + int64_t batch_stride_B_, + int64_t batch_stride_C_, + int64_t batch_stride_D_, + void *workspace_ = nullptr) { + + ptr_A = const_cast(ptr_A_); + ptr_B = const_cast(ptr_B_); + ptr_C = const_cast(ptr_C_); + ptr_D = ptr_D_; + + batch_stride_A = batch_stride_A_; + batch_stride_B = batch_stride_B_; + batch_stride_C = batch_stride_C_; + batch_stride_D = batch_stride_D_; + + + semaphore = static_cast(workspace_); + MCTLASS_TRACE_HOST("GemmParams::update()"); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_pipelined.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_pipelined.h new file mode 100644 index 0000000..5e407b4 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_pipelined.h @@ -0,0 +1,158 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/aligned_buffer.h" +#include "mctlass/array.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/gemm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +__global__ void GemmPipelined( + mctlass::gemm::GemmCoord problem_size, + mctlass::gemm::GemmCoord grid_tiled_shape, + typename Mma::IteratorA::Params params_A, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::Params params_B, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::Params params_epilogue + ) { + + // Shared storage needed by threadblock-scoped matrix multiply-accumulate + __shared__ union { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + } shared_storage; + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + int swizzle_log_tile = ThreadblockSwizzle().get_log_tile(grid_tiled_shape); + + mctlass::gemm::GemmCoord tb_tile_offset = threadblock_swizzle.get_tile_offset(swizzle_log_tile); + + if (grid_tiled_shape.m() <= tb_tile_offset.m() || + grid_tiled_shape.n() <= tb_tile_offset.n()) { + + return; + } + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + tb_tile_offset.m() * Mma::Shape::kM, + tb_tile_offset.k() + }; + + mctlass::MatrixCoord tb_offset_B{ + tb_tile_offset.k(), + tb_tile_offset.n() * Mma::Shape::kN + }; + + // Compute position within threadblock + int tb_thread_id = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params_A, + ref_A.data(), + {problem_size.m(), problem_size.k()}, + tb_thread_id, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + params_B, + ref_B.data(), + {problem_size.k(), problem_size.n()}, + tb_thread_id, + tb_offset_B); + + int warp_id = canonical_warp_idx(); + int lane_id = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, tb_thread_id, warp_id, lane_id); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + mma(problem_size, accumulators, iterator_A, iterator_B, accumulators); + + // + // Epilogue + // + + Epilogue epilogue( + params_epilogue, + shared_storage.epilogue, + tb_thread_id, + warp_id, + lane_id); + + tb_tile_offset = threadblock_swizzle.get_tile_offset(swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + tb_tile_offset.m() * Mma::Shape::kM, + tb_tile_offset.n() * Mma::Shape::kN + ); + + // run efficient epilogue + epilogue({problem_size.m(), problem_size.n()}, accumulators, threadblock_offset); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_planar_complex.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_planar_complex.h new file mode 100644 index 0000000..eb6e4d1 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_planar_complex.h @@ -0,0 +1,723 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/semaphore.h" +#include "mctlass/gemm/kernel/params_universal_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_ ///! Threadblock swizzling function +> +struct GemmPlanarComplex { +public: + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + using Operator = typename Mma::Operator; + using ArchTag = typename Mma::ArchTag; + + static ComplexTransform const kTransformA = Mma::kTransformA; + static ComplexTransform const kTransformB = Mma::kTransformB; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Split-K preserves splits that are 128b aligned + static int const kSplitKAlignment = const_max( + 128 / sizeof_bits::value, + 128 / sizeof_bits::value); + + // + // Additional types needed for reflection + // + + using ElementAccumulator = typename Mma::Policy::Operator::ElementC; + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::Shape; + + static int const kStages = Mma::kStages; + + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + // + // Arguments structure + // + + /// Argument structure + struct Arguments : UniversalArgumentsBase + { + // + // Data members + // + + typename EpilogueOutputOp::Params epilogue; + + void const * ptr_A_real; + void const * ptr_A_imag; + + void const * ptr_B_real; + void const * ptr_B_imag; + + void const * ptr_C_real; + void const * ptr_C_imag; + + void * ptr_D_real; + void * ptr_D_imag; + + typename LayoutA::Stride::Index lda_real; + typename LayoutA::Stride::Index lda_imag; + typename LayoutB::Stride::Index ldb_real; + typename LayoutB::Stride::Index ldb_imag; + typename LayoutC::Stride::Index ldc_real; + typename LayoutC::Stride::Index ldc_imag; + typename LayoutC::Stride::Index ldd_real; + typename LayoutC::Stride::Index ldd_imag; + + int64_t batch_stride_A; + int64_t batch_stride_A_imag; + int64_t batch_stride_B; + int64_t batch_stride_B_imag; + int64_t batch_stride_C; + int64_t batch_stride_C_imag; + int64_t batch_stride_D_imag; + + // + // Methods + // + + Arguments() : + ptr_A_real(nullptr), + ptr_A_imag(nullptr), + ptr_B_real(nullptr), + ptr_B_imag(nullptr), + ptr_C_real(nullptr), + ptr_C_imag(nullptr), + ptr_D_real(nullptr), + ptr_D_imag(nullptr) + {} + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A_real, + void const * ptr_A_imag, + void const * ptr_B_real, + void const * ptr_B_imag, + void const * ptr_C_real, + void const * ptr_C_imag, + void * ptr_D_real, + void * ptr_D_imag, + typename LayoutA::Stride::Index lda_real, + typename LayoutA::Stride::Index lda_imag, + typename LayoutB::Stride::Index ldb_real, + typename LayoutB::Stride::Index ldb_imag, + typename LayoutC::Stride::Index ldc_real, + typename LayoutC::Stride::Index ldc_imag, + typename LayoutC::Stride::Index ldd_real, + typename LayoutC::Stride::Index ldd_imag, + int64_t batch_stride_A = 0, + int64_t batch_stride_A_imag = 0, + int64_t batch_stride_B = 0, + int64_t batch_stride_B_imag = 0, + int64_t batch_stride_C = 0, + int64_t batch_stride_C_imag = 0, + int64_t batch_stride_D = 0, + int64_t batch_stride_D_imag = 0) + : + UniversalArgumentsBase(mode, problem_size, batch_count, batch_stride_D), + epilogue(epilogue), + ptr_A_real(ptr_A_real), + ptr_A_imag(ptr_A_imag), + ptr_B_real(ptr_B_real), + ptr_B_imag(ptr_B_imag), + ptr_C_real(ptr_C_real), + ptr_C_imag(ptr_C_imag), + ptr_D_real(ptr_D_real), + ptr_D_imag(ptr_D_imag), + lda_real(lda_real), + lda_imag(lda_imag), + ldb_real(ldb_real), + ldb_imag(ldb_imag), + ldc_real(ldc_real), + ldc_imag(ldc_imag), + ldd_real(ldd_real), + ldd_imag(ldd_imag), + batch_stride_A(batch_stride_A), + batch_stride_A_imag(batch_stride_A_imag), + batch_stride_B(batch_stride_B), + batch_stride_B_imag(batch_stride_B_imag), + batch_stride_C(batch_stride_C), + batch_stride_C_imag(batch_stride_C_imag), + batch_stride_D_imag(batch_stride_D_imag) + {} + + /// Returns arguments for the transposed problem + Arguments transposed_problem() const { + Arguments args(*this); + + std::swap(args.problem_size.m(), args.problem_size.n()); + std::swap(args.ptr_A_real, args.ptr_B_real); + std::swap(args.ptr_A_imag, args.ptr_B_imag); + std::swap(args.lda_real, args.ldb_real); + std::swap(args.lda_imag, args.ldb_imag); + std::swap(args.batch_stride_A, args.batch_stride_B); + std::swap(args.batch_stride_A_imag, args.batch_stride_B_imag); + + return args; + } + }; + + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params : UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC> + { + using ParamsBase = UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC>; + + // + // Data members + // + + typename Mma::IteratorA::Params params_A_real; + typename Mma::IteratorA::Params params_A_imag; + typename Mma::IteratorB::Params params_B_real; + typename Mma::IteratorB::Params params_B_imag; + typename Epilogue::OutputTileIterator::Params params_C_real; + typename Epilogue::OutputTileIterator::Params params_C_imag; + typename Epilogue::OutputTileIterator::Params params_D_real; + typename Epilogue::OutputTileIterator::Params params_D_imag; + + typename EpilogueOutputOp::Params output_op; + + void * ptr_A_real; + void * ptr_A_imag; + void * ptr_B_real; + void * ptr_B_imag; + void * ptr_C_real; + void * ptr_C_imag; + void * ptr_D_real; + void * ptr_D_imag; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + + int64_t batch_stride_A_imag; + int64_t batch_stride_B_imag; + int64_t batch_stride_C_imag; + int64_t batch_stride_D_imag; + + // + // Host dispatch API + // + + /// Default constructor + Params() = default; + + /// Constructor + Params( + Arguments const &args, /// GEMM application arguments + int device_sms, /// Number of SMs on the device + int sm_occupancy) /// Kernel SM occupancy (in thread blocks) + : + ParamsBase(args, device_sms, sm_occupancy), + params_A_real(args.lda_real), + params_A_imag(args.lda_imag), + params_B_real(args.ldb_real), + params_B_imag(args.ldb_imag), + params_C_real(args.ldc_real), + params_C_imag(args.ldc_imag), + params_D_real(args.ldd_real), + params_D_imag(args.ldd_imag), + output_op(args.epilogue), + ptr_A_real(const_cast(args.ptr_A_real)), + ptr_A_imag(const_cast(args.ptr_A_imag)), + ptr_B_real(const_cast(args.ptr_B_real)), + ptr_B_imag(const_cast(args.ptr_B_imag)), + ptr_C_real(const_cast(args.ptr_C_real)), + ptr_C_imag(const_cast(args.ptr_C_imag)), + ptr_D_real(args.ptr_D_real), + ptr_D_imag(args.ptr_D_imag), + batch_stride_A(args.batch_stride_A), + batch_stride_B(args.batch_stride_B), + batch_stride_C(args.batch_stride_C), + batch_stride_A_imag(args.batch_stride_A_imag), + batch_stride_B_imag(args.batch_stride_B_imag), + batch_stride_C_imag(args.batch_stride_C_imag), + batch_stride_D_imag(args.batch_stride_D_imag) + {} + + /// Returns the workspace size (in bytes) needed for this problem geometry + size_t get_workspace_size() const + { + size_t workspace_bytes = ParamsBase::get_workspace_size(); + if (this->mode == GemmUniversalMode::kGemmSplitKParallel) + { + // Double the size returned by the base class because we need to + // accumulate two ElementC components + workspace_bytes *= 2; + } + + return workspace_bytes; + } + + /// Lightweight update given a subset of arguments. + void update(Arguments const &args) + { + ptr_A_real = const_cast(args.ptr_A_real); + ptr_A_imag = const_cast(args.ptr_A_imag); + + ptr_B_real = const_cast(args.ptr_B_real); + ptr_B_imag = const_cast(args.ptr_B_imag); + + ptr_C_real = const_cast(args.ptr_C_real); + ptr_C_imag = const_cast(args.ptr_C_imag); + + ptr_D_real = const_cast(args.ptr_D_real); + ptr_D_imag = const_cast(args.ptr_D_imag); + + batch_stride_A = args.batch_stride_A; + batch_stride_B = args.batch_stride_B; + batch_stride_C = args.batch_stride_C; + this->batch_stride_D = args.batch_stride_D; + batch_stride_A_imag = args.batch_stride_A_imag; + batch_stride_B_imag = args.batch_stride_B_imag; + batch_stride_C_imag = args.batch_stride_C_imag; + batch_stride_D_imag = args.batch_stride_D_imag; + + output_op = args.epilogue; + } + }; + + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + +public: + + // + // Host dispatch API + // + + /// Determines whether kernel satisfies alignment + static Status can_implement(Arguments const &args) + { + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + bool isAMisaligned = false; + bool isBMisaligned = false; + bool isCMisaligned = false; + + if (platform::is_same::value) { + isAMisaligned = args.problem_size.k() % kAlignmentA; + } else if (platform::is_same::value) { + isAMisaligned = args.problem_size.m() % kAlignmentA; + } + + if (platform::is_same::value) { + isBMisaligned = args.problem_size.n() % kAlignmentB; + } else if (platform::is_same::value) { + isBMisaligned = args.problem_size.k() % kAlignmentB; + } + + if (platform::is_same::value) { + isCMisaligned = args.problem_size.n() % kAlignmentC; + } else if (platform::is_same::value) { + isCMisaligned = args.problem_size.m() % kAlignmentC; + } + + if (isAMisaligned || isBMisaligned || isCMisaligned) { + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + +public: + + // + // Device-only API + // + + // Factory invocation + MCTLASS_DEVICE + static void invoke( + Params const ¶ms, + SharedStorage &shared_storage) + { + GemmPlanarComplex op; + op(params, shared_storage); + } + + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + int offset_k = 0; + int problem_size_k = params.problem_size.k(); + + ElementA *ptr_A_real = static_cast(params.ptr_A_real); + ElementA *ptr_A_imag = static_cast(params.ptr_A_imag); + + ElementB *ptr_B_real = static_cast(params.ptr_B_real); + ElementB *ptr_B_imag = static_cast(params.ptr_B_imag); + + // + // Fetch pointers based on mode. + // + if (params.mode == GemmUniversalMode::kGemm || + params.mode == GemmUniversalMode::kGemmSplitKParallel) { + + if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) { + + problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size; + } + + offset_k = threadblock_tile_offset.k() * params.gemm_k_size; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_A_real += int64_t(threadblock_tile_offset.k()) * params.batch_stride_A; + ptr_A_imag += int64_t(threadblock_tile_offset.k()) * params.batch_stride_A_imag; + ptr_B_real += int64_t(threadblock_tile_offset.k()) * params.batch_stride_B; + ptr_B_imag += int64_t(threadblock_tile_offset.k()) * params.batch_stride_B_imag; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_A_real = static_cast(params.ptr_A_real)[threadblock_tile_offset.k()]; + ptr_A_imag = static_cast(params.ptr_A_imag)[threadblock_tile_offset.k()]; + ptr_B_real = static_cast(params.ptr_B_real)[threadblock_tile_offset.k()]; + ptr_B_imag = static_cast(params.ptr_B_imag)[threadblock_tile_offset.k()]; + } + + __syncthreads(); + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + offset_k, + }; + + mctlass::MatrixCoord tb_offset_B{ + offset_k, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A_real( + params.params_A_real, + ptr_A_real, + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorA iterator_A_imag( + params.params_A_imag, + ptr_A_imag, + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B_real( + params.params_B_real, + ptr_B_real, + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B); + + typename Mma::IteratorB iterator_B_imag( + params.params_B_imag, + ptr_B_imag, + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma( + gemm_k_iterations, + accumulators, + iterator_A_real, + iterator_A_imag, + iterator_B_real, + iterator_B_imag, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + ElementC *ptr_C_real = static_cast(params.ptr_C_real); + ElementC *ptr_C_imag = static_cast(params.ptr_C_imag); + ElementC *ptr_D_real = static_cast(params.ptr_D_real); + ElementC *ptr_D_imag = static_cast(params.ptr_D_imag); + + // + // Fetch pointers based on mode. + // + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + if (params.mode == GemmUniversalMode::kGemm) { + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + } + else if (params.mode == GemmUniversalMode::kGemmSplitKParallel) { + ptr_D_real += threadblock_tile_offset.k() * params.batch_stride_D; + ptr_D_imag += threadblock_tile_offset.k() * params.batch_stride_D_imag; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_C_real += int64_t(threadblock_tile_offset.k()) * params.batch_stride_C; + ptr_C_imag += int64_t(threadblock_tile_offset.k()) * params.batch_stride_C_imag; + ptr_D_real += int64_t(threadblock_tile_offset.k()) * params.batch_stride_D; + ptr_D_imag += int64_t(threadblock_tile_offset.k()) * params.batch_stride_D_imag; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_C_real = static_cast(params.ptr_C_real)[threadblock_tile_offset.k()]; + ptr_C_imag = static_cast(params.ptr_C_imag)[threadblock_tile_offset.k()]; + ptr_D_real = static_cast(params.ptr_D_real)[threadblock_tile_offset.k()]; + ptr_D_imag = static_cast(params.ptr_D_imag)[threadblock_tile_offset.k()]; + } + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C_real( + params.params_C_real, + ptr_C_real, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + typename Epilogue::OutputTileIterator iterator_C_imag( + params.params_C_imag, + ptr_C_imag, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D_real( + params.params_D_real, + ptr_D_real, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + typename Epilogue::OutputTileIterator iterator_D_imag( + params.params_D_imag, + ptr_D_imag, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // + // Construct epilogue + // + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C_real = iterator_D_real; + iterator_C_imag = iterator_D_imag; + } + + semaphore.wait(threadblock_tile_offset.k()); + + __threadfence(); + } + + + // Execute the epilogue operator to update the destination tensor. + epilogue( + output_op, + iterator_D_real, + iterator_D_imag, + accumulators, + iterator_C_real, + iterator_C_imag); + + // + // Release the semaphore + // + + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_planar_complex_array.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_planar_complex_array.h new file mode 100644 index 0000000..2192416 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_planar_complex_array.h @@ -0,0 +1,617 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/semaphore.h" +#include "mctlass/gemm/kernel/params_universal_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_ ///! Threadblock swizzling function +> +struct GemmPlanarComplexArray { +public: + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + using Operator = typename Mma::Operator; + using ArchTag = typename Mma::ArchTag; + + static ComplexTransform const kTransformA = Mma::kTransformA; + static ComplexTransform const kTransformB = Mma::kTransformB; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Split-K preserves splits that are 128b aligned + static int const kSplitKAlignment = const_max( + 128 / sizeof_bits::value, + 128 / sizeof_bits::value); + + // + // Additional types needed for reflection + // + + using ElementAccumulator = typename Mma::Policy::Operator::ElementC; + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::Shape; + + static int const kStages = Mma::kStages; + + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + // + // Arguments structure + // + + /// Argument structure + struct Arguments : UniversalArgumentsBase + { + // + // Data members + // + + typename EpilogueOutputOp::Params epilogue; + + int const *ptr_M; + int const *ptr_N; + int const *ptr_K; + + void const * const * ptr_A_real; + void const * const * ptr_A_imag; + + void const * const * ptr_B_real; + void const * const * ptr_B_imag; + + void const * const * ptr_C_real; + void const * const * ptr_C_imag; + + void * const * ptr_D_real; + void * const * ptr_D_imag; + + typename LayoutA::Stride::Index lda_real; + typename LayoutA::Stride::Index lda_imag; + typename LayoutB::Stride::Index ldb_real; + typename LayoutB::Stride::Index ldb_imag; + typename LayoutC::Stride::Index ldc_real; + typename LayoutC::Stride::Index ldc_imag; + typename LayoutC::Stride::Index ldd_real; + typename LayoutC::Stride::Index ldd_imag; + + // + // Methods + // + + Arguments(): + ptr_M(nullptr), + ptr_N(nullptr), + ptr_K(nullptr), + ptr_A_real(nullptr), + ptr_A_imag(nullptr), + ptr_B_real(nullptr), + ptr_B_imag(nullptr), + ptr_C_real(nullptr), + ptr_C_imag(nullptr), + ptr_D_real(nullptr), + ptr_D_imag(nullptr) + {} + + /// constructs an arguments structure + Arguments( + GemmCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params epilogue, + int const *ptr_M, + int const *ptr_N, + int const *ptr_K, + void const * const * ptr_A_real, + void const * const * ptr_A_imag, + void const * const * ptr_B_real, + void const * const * ptr_B_imag, + void const * const * ptr_C_real, + void const * const * ptr_C_imag, + void * const * ptr_D_real, + void * const * ptr_D_imag, + typename LayoutA::Stride::Index lda_real, + typename LayoutA::Stride::Index lda_imag, + typename LayoutB::Stride::Index ldb_real, + typename LayoutB::Stride::Index ldb_imag, + typename LayoutC::Stride::Index ldc_real, + typename LayoutC::Stride::Index ldc_imag, + typename LayoutC::Stride::Index ldd_real, + typename LayoutC::Stride::Index ldd_imag) + : + UniversalArgumentsBase(mode, problem_size, batch_count, batch_stride_D), + epilogue(epilogue), + ptr_M(ptr_M), + ptr_N(ptr_N), + ptr_K(ptr_K), + ptr_A_real(ptr_A_real), + ptr_A_imag(ptr_A_imag), + ptr_B_real(ptr_B_real), + ptr_B_imag(ptr_B_imag), + ptr_C_real(ptr_C_real), + ptr_C_imag(ptr_C_imag), + ptr_D_real(ptr_D_real), + ptr_D_imag(ptr_D_imag), + lda_real(lda_real), + lda_imag(lda_imag), + ldb_real(ldb_real), + ldb_imag(ldb_imag), + ldc_real(ldc_real), + ldc_imag(ldc_imag), + ldd_real(ldd_real), + ldd_imag(ldd_imag) + {} + + /// Returns arguments for the transposed problem + Arguments transposed_problem() const { + Arguments args(*this); + + std::swap(args.problem_size.m(), args.problem_size.n()); + std::swap(args.ptr_M, args.ptr_N); + std::swap(args.ptr_A_real, args.ptr_B_real); + std::swap(args.ptr_A_imag, args.ptr_B_imag); + std::swap(args.lda_real, args.ldb_real); + std::swap(args.lda_imag, args.ldb_imag); + + return args; + } + }; + + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params : UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC> + { + using ParamsBase = UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC>; + + // + // Data members + // + + typename Mma::IteratorA::Params params_A_real; + typename Mma::IteratorA::Params params_A_imag; + typename Mma::IteratorB::Params params_B_real; + typename Mma::IteratorB::Params params_B_imag; + typename Epilogue::OutputTileIterator::Params params_C_real; + typename Epilogue::OutputTileIterator::Params params_C_imag; + typename Epilogue::OutputTileIterator::Params params_D_real; + typename Epilogue::OutputTileIterator::Params params_D_imag; + + typename EpilogueOutputOp::Params output_op; + + int const *ptr_M; + int const *ptr_N; + int const *ptr_K; + + void const * const * ptr_A_real; + void const * const * ptr_A_imag; + void const * const * ptr_B_real; + void const * const * ptr_B_imag; + void const * const * ptr_C_real; + void const * const * ptr_C_imag; + void * const * ptr_D_real; + void * const * ptr_D_imag; + + // + // Host dispatch API + // + + /// Default constructor + Params() = default; + + /// Constructor + Params( + Arguments const &args, /// GEMM application arguments + int device_sms, /// Number of SMs on the device + int sm_occupancy) /// Kernel SM occupancy (in thread blocks) + : + ParamsBase(args, device_sms, sm_occupancy), + ptr_M(args.ptr_M), + ptr_N(args.ptr_N), + ptr_K(args.ptr_K), + params_A_real(args.lda_real), + params_A_imag(args.lda_imag), + params_B_real(args.ldb_real), + params_B_imag(args.ldb_imag), + params_C_real(args.ldc_real), + params_C_imag(args.ldc_imag), + params_D_real(args.ldd_real), + params_D_imag(args.ldd_imag), + output_op(args.epilogue), + ptr_A_real(args.ptr_A_real), + ptr_A_imag(args.ptr_A_imag), + ptr_B_real(args.ptr_B_real), + ptr_B_imag(args.ptr_B_imag), + ptr_C_real(args.ptr_C_real), + ptr_C_imag(args.ptr_C_imag), + ptr_D_real(args.ptr_D_real), + ptr_D_imag(args.ptr_D_imag) + {} + + /// Lightweight update given a subset of arguments. + void update(Arguments const &args) + { + ptr_M = args.ptr_M; + ptr_N = args.ptr_N; + ptr_K = args.ptr_K; + + ptr_A_real = args.ptr_A_real; + ptr_A_imag = args.ptr_A_imag; + + ptr_B_real = args.ptr_B_real; + ptr_B_imag = args.ptr_B_imag; + + ptr_C_real = args.ptr_C_real; + ptr_C_imag = args.ptr_C_imag; + + ptr_D_real = args.ptr_D_real; + ptr_D_imag = args.ptr_D_imag; + + output_op = args.epilogue; + } + }; + + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + +public: + + // + // Host dispatch API + // + + /// Determines whether kernel satisfies alignment + static Status can_implement(Arguments const &args) { + + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + bool isAMisaligned = false; + bool isBMisaligned = false; + bool isCMisaligned = false; + + if (platform::is_same::value) { + isAMisaligned = args.problem_size.k() % kAlignmentA; + } else if (platform::is_same::value) { + isAMisaligned = args.problem_size.m() % kAlignmentA; + } + + if (platform::is_same::value) { + isBMisaligned = args.problem_size.n() % kAlignmentB; + } else if (platform::is_same::value) { + isBMisaligned = args.problem_size.k() % kAlignmentB; + } + + if (platform::is_same::value) { + isCMisaligned = args.problem_size.n() % kAlignmentC; + } else if (platform::is_same::value) { + isCMisaligned = args.problem_size.m() % kAlignmentC; + } + + if (isAMisaligned || isBMisaligned || isCMisaligned) { + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + + +public: + + // + // Device-only API + // + + // Factory invocation + MCTLASS_DEVICE + static void invoke( + Params const ¶ms, + SharedStorage &shared_storage) + { + GemmPlanarComplexArray op; + op(params, shared_storage); + } + + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + int batch_idx = threadblock_tile_offset.k(); + + int problem_size_m = params.problem_size.m(); + int problem_size_n = params.problem_size.n(); + int problem_size_k = params.problem_size.k(); + + ElementA *ptr_A_real = static_cast(const_cast(params.ptr_A_real[batch_idx])); + ElementA *ptr_A_imag = static_cast(const_cast(params.ptr_A_imag[batch_idx])); + + ElementB *ptr_B_real = static_cast(const_cast(params.ptr_B_real[batch_idx])); + ElementB *ptr_B_imag = static_cast(const_cast(params.ptr_B_imag[batch_idx])); + + // + // If pointers for problem sizes are specified, these are loaded from global memory + // + + if (params.ptr_M) { + problem_size_m = params.ptr_M[batch_idx]; + } + + if (params.ptr_N) { + problem_size_n = params.ptr_N[batch_idx]; + } + + if (params.ptr_K) { + problem_size_k = params.ptr_K[batch_idx]; + } + + int const kBlockCountM = (problem_size_m + Mma::Shape::kM - 1) / Mma::Shape::kM; + int const kBlockCountN = (problem_size_n + Mma::Shape::kN - 1) / Mma::Shape::kN; + + int const kGemmKIterations = (problem_size_k + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // + // Each threadblock loops over the logical problem size which the kernel may have discovered + // after the grid is launched. + // + + MCTLASS_PRAGMA_NO_UNROLL + for (int block_m = threadblock_tile_offset.m(); + block_m < kBlockCountM; + block_m += params.grid_tiled_shape.m()) { + + MCTLASS_PRAGMA_NO_UNROLL + for (int block_n = threadblock_tile_offset.n(); + block_n < kBlockCountN; + block_n += params.grid_tiled_shape.n()) { + + // + // Compute indices within threadblock and warp. + // + int thread_idx = threadIdx.x; + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + int lane_idx = threadIdx.x % 64; + + // + // Proceed with regular GEMM logic. + // + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ block_m * Mma::Shape::kM, 0}; + mctlass::MatrixCoord tb_offset_B{ 0, block_n * Mma::Shape::kN }; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A_real( + params.params_A_real, + ptr_A_real, + {problem_size_m, problem_size_k}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorA iterator_A_imag( + params.params_A_imag, + ptr_A_imag, + {problem_size_m, problem_size_k}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B_real( + params.params_B_real, + ptr_B_real, + {problem_size_k, problem_size_n}, + thread_idx, + tb_offset_B); + + typename Mma::IteratorB iterator_B_imag( + params.params_B_imag, + ptr_B_imag, + {problem_size_k, problem_size_n}, + thread_idx, + tb_offset_B); + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + mma( + kGemmKIterations, + accumulators, + iterator_A_real, + iterator_A_imag, + iterator_B_real, + iterator_B_imag, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + //assume identity swizzle + MatrixCoord threadblock_offset( + block_m * Mma::Shape::kM, + block_n * Mma::Shape::kN + ); + + ElementC *ptr_C_real = static_cast(const_cast(params.ptr_C_real[batch_idx])); + ElementC *ptr_C_imag = static_cast(const_cast(params.ptr_C_imag[batch_idx])); + ElementC *ptr_D_real = static_cast(params.ptr_D_real[batch_idx]); + ElementC *ptr_D_imag = static_cast(params.ptr_D_imag[batch_idx]); + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C_real( + params.params_C_real, + ptr_C_real, + {problem_size_m, problem_size_n}, + thread_idx, + threadblock_offset + ); + + typename Epilogue::OutputTileIterator iterator_C_imag( + params.params_C_imag, + ptr_C_imag, + {problem_size_m, problem_size_n}, + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D_real( + params.params_D_real, + ptr_D_real, + {problem_size_m, problem_size_n}, + thread_idx, + threadblock_offset + ); + + typename Epilogue::OutputTileIterator iterator_D_imag( + params.params_D_imag, + ptr_D_imag, + {problem_size_m, problem_size_n}, + thread_idx, + threadblock_offset + ); + + // + // Construct epilogue + // + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Execute the epilogue operator to update the destination tensor. + epilogue( + output_op, + iterator_D_real, + iterator_D_imag, + accumulators, + iterator_C_real, + iterator_C_imag); + + + } // for block_n + } // for block_m + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_splitk_parallel.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_splitk_parallel.h new file mode 100644 index 0000000..47f993c --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_splitk_parallel.h @@ -0,0 +1,253 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for GEMM performing a reduction over K partitions in parallel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_ ///! Threadblock swizzling function +> +struct GemmSplitKParallel { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using OutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + static int const kAlignmentK = Mma::Operator::Shape::kK; + + /// Parameters structure + struct Params { + mctlass::gemm::GemmCoord problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorA::TensorRef ref_A; + typename Mma::IteratorB::Params params_B; + typename Mma::IteratorB::TensorRef ref_B; + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::OutputTileIterator::TensorRef ref_D; + typename OutputOp::Params output_op; + int64_t splitk_slice_stride; + int gemm_k_size; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): swizzle_log_tile(0) { } + + MCTLASS_HOST_DEVICE + Params( + mctlass::gemm::GemmCoord const & problem_size, + mctlass::gemm::GemmCoord const & grid_tiled_shape, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::OutputTileIterator::TensorRef ref_D, + typename OutputOp::Params output_op, + int64_t splitk_slice_stride + ): + problem_size(problem_size), + grid_tiled_shape(grid_tiled_shape), + swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)), + params_A(ref_A.layout()), + ref_A(ref_A), + params_B(ref_B.layout()), + ref_B(ref_B), + params_D(ref_D.layout()), + ref_D(ref_D), + output_op(output_op), + splitk_slice_stride(splitk_slice_stride) { + + int full_gemm_k_iterations = problem_size.k() / Mma::Shape::kK; + int gemm_k_iterations = full_gemm_k_iterations / grid_tiled_shape.k(); + + gemm_k_size = gemm_k_iterations * Mma::Shape::kK; + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + GemmSplitKParallel() { } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.k() * params.gemm_k_size, + }; + + mctlass::MatrixCoord tb_offset_B{ + threadblock_tile_offset.k() * params.gemm_k_size, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + // Problem size is a function of threadblock index in the K dimension + int problem_size_k; + if (threadblock_tile_offset.k() + 1 == params.grid_tiled_shape.k()) { + problem_size_k = params.problem_size.k(); + } + else { + problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size; + } + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - tb_offset_A.column() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, + params.ref_A.data(), + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + params.params_B, + params.ref_B.data(), + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B); + + int warp_idx = threadIdx.x / 64; + int lane_idx = threadIdx.x % 64; + + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators); + + // + // Epilogue + // + + OutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + // Tile iterator writing to output tile + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + params.ref_D.data(), + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + iterator_D.add_pointer_offset(params.splitk_slice_stride * threadblock_tile_offset.k()); + + // Execute the epilogue + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Run efficient epilogue + epilogue(output_op, iterator_D, accumulators, iterator_D); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_streamk_with_fused_epilogue.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_streamk_with_fused_epilogue.h new file mode 100644 index 0000000..9126509 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_streamk_with_fused_epilogue.h @@ -0,0 +1,2405 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Stream-K Gemm kernel compatible with fused epilogues + that broadcast a bias vector over the MMA output. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/layout/layout.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/barrier.h" +#include "mctlass/block_striped.h" +#include "mctlass/semaphore.h" + +#include "mctlass/trace.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + bool IsSingleSource = Epilogue_::kIsSingleSource +> +struct GemmStreamkWithFusedEpilogue; + +// GemmStreamkWithFusedEpilogue with two sources +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_ ///! Threadblock swizzling function +> +struct GemmStreamkWithFusedEpilogue { + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + + /// The per-thread tile of raw accumulators + using AccumulatorTile = typename Mma::FragmentC; + + static ComplexTransform const kTransformA = Mma::kTransformA; + static ComplexTransform const kTransformB = Mma::kTransformB; + using Operator = typename Mma::Operator; + + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::InstructionShape; + using ArchTag = typename Mma::ArchTag; + + static int const kStages = Mma::kStages; + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Workspace bytes per thread block + static size_t const kWorkspaceBytesPerBlock = + __NV_STD_MAX( + kThreadCount * sizeof(AccumulatorTile), + Epilogue::kWorkspaceBytesPerBlock); + + /// Block-striped reduction utility + using BlockStripedReduceT = BlockStripedReduce; + + + + // + // Structures + // + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmUniversalMode mode; + GemmCoord problem_size; + int batch_count; // Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor + + typename EpilogueOutputOp::Params epilogue; + + void const * ptr_A; + void const * ptr_B; + void const * ptr_C1; + void const * ptr_C2; + void * ptr_D; + + void * ptr_Vector; + void * ptr_Tensor; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C1; + int64_t batch_stride_C2; + int64_t batch_stride_D; + int64_t batch_stride_Vector; + int64_t batch_stride_Tensor; + + typename LayoutA::Stride::Index lda; + typename LayoutB::Stride::Index ldb; + typename LayoutC::Stride::Index ldc1; + typename LayoutC::Stride::Index ldc2; + typename LayoutC::Stride::Index ldd; + typename LayoutC::Stride::Index ldr; + typename LayoutC::Stride::Index ldt; + + int avail_sms; /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling) + + + // + // Methods + // + + /// Default Constructor + Arguments(): + mode(GemmUniversalMode::kGemm), + batch_count(1), + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C1(nullptr), + ptr_C2(nullptr), + ptr_D(nullptr), + avail_sms(-1) + {} + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_split, /// Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor (1 defaults to StreamK, >1 emulates Split-K) + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void const * ptr_C1, + void const * ptr_C2, + void * ptr_D, + void * ptr_Vector, + void * ptr_Tensor, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C1, + int64_t batch_stride_C2, + int64_t batch_stride_D, + int64_t batch_stride_Vector, + int64_t batch_stride_Tensor, + typename LayoutA::Stride::Index lda, + typename LayoutB::Stride::Index ldb, + typename LayoutC::Stride::Index ldc1, + typename LayoutC::Stride::Index ldc2, + typename LayoutC::Stride::Index ldd, + typename LayoutC::Stride::Index ldr, + typename LayoutC::Stride::Index ldt, + int avail_sms = -1) /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling) + : + mode(mode), + problem_size(problem_size), + batch_count(batch_split), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_C1(ptr_C1), ptr_C2(ptr_C2), ptr_D(ptr_D), + ptr_Vector(ptr_Vector), + ptr_Tensor(ptr_Tensor), + batch_stride_A(batch_stride_A), + batch_stride_B(batch_stride_B), + batch_stride_C1(batch_stride_C1), + batch_stride_C2(batch_stride_C2), + batch_stride_Vector(batch_stride_Vector), + batch_stride_Tensor(batch_stride_Tensor), + lda(lda), ldb(ldb), ldc1(ldc1), ldc2(ldc2), ldd(ldd), ldr(ldr), ldt(ldt), avail_sms(avail_sms) + { + MCTLASS_TRACE_HOST("GemmStreamkWithFusedEpilogue::Arguments::Arguments() - problem_size: " << problem_size); + MCTLASS_TRACE_HOST(" ptr_Vector: " << (void *)this->ptr_Vector); + MCTLASS_TRACE_HOST(" ptr_Tensor: " << (void *)this->ptr_Tensor); + MCTLASS_TRACE_HOST(" ldr: " << this->ldr); + MCTLASS_TRACE_HOST(" ldt: " << this->ldt); + MCTLASS_TRACE_HOST(" avail_sms: " << this->avail_sms); + } + + /// Returns arguments for the transposed problem + Arguments transposed_problem() const { + Arguments args(*this); + + std::swap(args.problem_size.m(), args.problem_size.n()); + std::swap(args.ptr_A, args.ptr_B); + std::swap(args.lda, args.ldb); + std::swap(args.batch_stride_A, args.batch_stride_B); + + return args; + } + }; + + + /// Parameters structure + struct Params + { + public: + + // + // Data members + // + + void * ptr_A; + void * ptr_B; + + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorB::Params params_B; + + int64_t batch_stride_A; + int64_t batch_stride_B; + + GemmUniversalMode mode; + + ThreadblockSwizzle block_mapping; + + void *barrier_workspace; + void *partials_workspace; + + typename EpilogueOutputOp::Params output_op; + + void * ptr_C1; + void * ptr_C2; + void * ptr_D; + void * ptr_Tensor; + void * ptr_Vector; + + typename Epilogue::OutputTileIterator::Params params_C1; + typename Epilogue::OutputTileIterator::Params params_C2; + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::TensorTileIterator::Params params_Tensor; + + int64_t batch_stride_C1; + int64_t batch_stride_C2; + int64_t batch_stride_D; + int64_t batch_stride_Vector; + int64_t batch_stride_Tensor; + + typename LayoutC::Stride::Index ldr; + + protected: + + // + // Host-only dispatch-utilities + // + + /// Pad the given allocation size up to the nearest cache line + static size_t cacheline_align_up(size_t size) + { + static const int CACHELINE_SIZE = 128; + return (size + CACHELINE_SIZE - 1) / CACHELINE_SIZE * CACHELINE_SIZE; + } + + /// Get the workspace size needed for barrier + size_t get_barrier_workspace_size() const + { + // For atomic reduction, each SK-block needs a synchronization flag. For parallel reduction, + // each reduction block needs its own synchronization flag. + int sk_blocks = block_mapping.sk_regions() * block_mapping.sk_blocks_per_region(); + int num_flags = fast_max(sk_blocks, block_mapping.reduction_blocks); + + return cacheline_align_up(sizeof(typename Barrier::T) * num_flags); + } + + /// Get the workspace size needed for intermediate partial sums + size_t get_partials_workspace_size() const + { + int sk_blocks = block_mapping.sk_regions() * block_mapping.sk_blocks_per_region(); + return cacheline_align_up(kWorkspaceBytesPerBlock * sk_blocks); + } + + + public: + + // + // Host dispatch API + // + + /// Default constructor + Params() = default; + + /// Constructor + Params( + Arguments const &args, /// GEMM application arguments + int device_sms, /// Number of SMs on the device + int sm_occupancy) /// Kernel SM occupancy (in thread blocks) + : + params_A(args.lda), + params_B(args.ldb), + params_C1(args.ldc1), + params_C2(args.ldc2), + params_D(args.ldd), + params_Tensor(args.ldt), + output_op(args.epilogue), + mode(args.mode), + ptr_A(const_cast(args.ptr_A)), + ptr_B(const_cast(args.ptr_B)), + ptr_C1(const_cast(args.ptr_C1)), + ptr_C2(const_cast(args.ptr_C2)), + ptr_D(args.ptr_D), + ptr_Vector(args.ptr_Vector), + ldr(args.ldr), + ptr_Tensor(args.ptr_Tensor), + batch_stride_A(args.batch_stride_A), + batch_stride_B(args.batch_stride_B), + batch_stride_C1(args.batch_stride_C1), + batch_stride_C2(args.batch_stride_C2), + batch_stride_D(args.batch_stride_D), + batch_stride_Vector(args.batch_stride_Vector), + batch_stride_Tensor(args.batch_stride_Tensor), + barrier_workspace(nullptr), + partials_workspace(nullptr) + { + MCTLASS_TRACE_HOST("GemmStreamkWithFusedEpilogue::Params::Params() - problem_size: " << problem_size); + MCTLASS_TRACE_HOST(" ptr_Vector: " << (void *)this->ptr_Vector); + MCTLASS_TRACE_HOST(" ptr_Tensor: " << (void *)this->ptr_Tensor); + MCTLASS_TRACE_HOST(" ldr: " << this->ldr); + MCTLASS_TRACE_HOST(" ldt: " << args.ldt); + MCTLASS_TRACE_HOST(" avail_sms: " << avail_sms); + + // Number of SMs to make available for StreamK decomposition + int avail_sms = (args.avail_sms == -1) ? + device_sms : + fast_min(args.avail_sms, device_sms); + + // Initialize the block mapping structure + block_mapping = ThreadblockSwizzle( + typename ThreadblockSwizzle::template KernelTraits(), + args.mode, + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count, + sm_occupancy, + device_sms, + avail_sms); + } + + /// Returns the workspace size (in bytes) needed for these parameters + size_t get_workspace_size() const + { + return + get_barrier_workspace_size() + + get_partials_workspace_size(); + } + + /// Assign and initialize the specified workspace buffer. Assumes + /// the memory allocated to workspace is at least as large as get_workspace_size(). + Status init_workspace( + void *workspace, + mcStream_t stream = nullptr) + { + uint8_t *ptr = static_cast(workspace); + + + // Establish partials workspace + partials_workspace = nullptr; + size_t partials_workspace_bytes = get_partials_workspace_size(); + if (partials_workspace_bytes > 0) + { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + partials_workspace = ptr; + ptr += partials_workspace_bytes; + } + + // Establish barrier workspace + barrier_workspace = nullptr; + size_t barrier_workspace_bytes = get_barrier_workspace_size(); + if (barrier_workspace_bytes > 0) + { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + barrier_workspace = ptr; + ptr += barrier_workspace_bytes; + } + + // Zero-initialize barrier workspace + if (barrier_workspace) + { + size_t barrier_workspace_bytes = get_barrier_workspace_size(); + + MCTLASS_TRACE_HOST(" Initialize " << barrier_workspace_bytes << " barrier bytes"); + + mcError_t result = mcMemsetAsync( + barrier_workspace, + 0, + barrier_workspace_bytes, + stream); + + if (result != mcSuccess) { + MCTLASS_TRACE_HOST(" mcMemsetAsync() returned error " << mcGetErrorString(result)); + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + + /// Returns the GEMM volume in thread block tiles + mctlass::gemm::GemmCoord get_tiled_shape() const + { + return block_mapping.tiled_shape(); + } + + /// Returns the total number of thread blocks to launch + int get_grid_blocks() const + { + dim3 grid_dims = get_grid_dims(); + return grid_dims.x * grid_dims.y * grid_dims.z; + } + + /// Returns the grid extents in thread blocks to launch + dim3 get_grid_dims() const + { + return block_mapping.get_grid_dims(); + } + + /// Lightweight update given a subset of arguments. Problem geometry is assumed + /// to remain the same. + MCTLASS_HOST_DEVICE + void update(Arguments const &args) + { + ptr_A = const_cast(args.ptr_A); + ptr_B = const_cast(args.ptr_B); + ptr_C1 = const_cast(args.ptr_C1); + ptr_C2 = const_cast(args.ptr_C2); + ptr_D = args.ptr_D; + + ptr_Vector = args.ptr_Vector; + ldr = args.ldr; + ptr_Tensor = args.ptr_Tensor; + + batch_stride_A = args.batch_stride_A; + batch_stride_B = args.batch_stride_B; + batch_stride_C1 = args.batch_stride_C1; + batch_stride_C2 = args.batch_stride_C2; + batch_stride_D = args.batch_stride_D; + batch_stride_Vector = args.batch_stride_Vector; + batch_stride_Tensor = args.batch_stride_Tensor; + + output_op = args.epilogue; + + MCTLASS_TRACE_HOST("GemmStreamkWithFusedEpilogue::Params::update()"); + MCTLASS_TRACE_HOST(" ptr_Vector: " << (void *)this->ptr_Vector); + MCTLASS_TRACE_HOST(" ptr_Tensor: " << (void *)this->ptr_Tensor); + MCTLASS_TRACE_HOST(" ldr: " << this->ldr); + } + }; + + /// Tile work descriptor + struct TileWorkDesc + { + /// The linear tile index + int tile_idx; + + /// The location of this tile (in threadblock-tile coordinates) in the output matrix + mctlass::gemm::GemmCoord tiled_coord; + + // The first global-scoped MAC-iteration this threadblock will perform for this tile + int iter_begin; + + // The starting index in the k-domain for MAC-iterations this threadblock will perform for this tile + int k_begin; + + // The ending index (one-past) in the k-domain for MAC-iterations this threadblock will perform for this tile + int k_end; + + /// The number of remaining MAC-iterations this threadblock will perform for this tile + int k_iters_remaining; + + // Whether this block will perform the first iteration of this tile + MCTLASS_DEVICE + bool tile_started() + { + return (k_begin == 0); + } + + // Whether this block will perform the last iteration of this tile + MCTLASS_DEVICE + bool tile_finished(Params const ¶ms) + { + return (k_end == params.block_mapping.problem_size.k()); + } + }; + + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + +protected: + + // + // Data members + // + + /// GEMM problem parameters + Params const ¶ms; + + /// Shared storage reference + SharedStorage &shared_storage; + + /// ID within the threadblock + int thread_idx; + + /// ID of warp + int warp_idx; + + /// ID of each thread within a warp + int lane_idx; + + /// Threadblock scoped epilogue + Epilogue epilogue; + + +public: + + // + // Host dispatch API + // + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size) { + + MCTLASS_TRACE_HOST("GemmStreamkWithFusedEpilogue::can_implement()"); + + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + bool isAMisaligned = false; + bool isBMisaligned = false; + bool isCMisaligned = false; + + if (platform::is_same::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } else if (platform::is_same::value) { + isAMisaligned = problem_size.m() % kAlignmentA; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } + + if (platform::is_same::value) { + isBMisaligned = problem_size.n() % kAlignmentB; + } else if (platform::is_same::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } + + if (platform::is_same::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } else if (platform::is_same::value) { + isCMisaligned = problem_size.m() % kAlignmentC; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } + + if (isAMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for A operand"); + return Status::kErrorMisalignedOperand; + } + + if (isBMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for B operand"); + return Status::kErrorMisalignedOperand; + } + + if (isCMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for C operand"); + return Status::kErrorMisalignedOperand; + } + + MCTLASS_TRACE_HOST(" returning kSuccess"); + + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + +protected: + + // + // Device-only utility methods + // + + /// Iterator for fetching tile fragments from A + MCTLASS_DEVICE + typename Mma::IteratorA init_iterator_A( + TileWorkDesc &tile_work, + GemmUniversalMode mode) + { + // The input A matrix + ElementA *ptr_A = static_cast(params.ptr_A); + + // Update input pointers based on batched/array mode + if (mode == GemmUniversalMode::kBatched) { + ptr_A += tile_work.tiled_coord.k() * params.batch_stride_A; + } + if (mode == GemmUniversalMode::kArray) { + ptr_A = static_cast(params.ptr_A)[tile_work.tiled_coord.k()]; + } + + int m_begin = tile_work.tiled_coord.m() * Mma::Shape::kM; + int m_end = params.block_mapping.problem_size.m(); + return Mma::IteratorA( + params.params_A, + ptr_A, + { m_end, tile_work.k_end }, + threadIdx.x, + { m_begin, tile_work.k_begin }); + + } + + + /// Iterator for fetching tile fragments from B + MCTLASS_DEVICE + typename Mma::IteratorB init_iterator_B( + TileWorkDesc &tile_work, + GemmUniversalMode mode) + { + // The input B matrix + ElementB *ptr_B = static_cast(params.ptr_B); + + // Update input pointers based on batched/array mode + if (mode == GemmUniversalMode::kBatched) { + ptr_B += tile_work.tiled_coord.k() * params.batch_stride_B; + } + if (mode == GemmUniversalMode::kArray) { + ptr_B = static_cast(params.ptr_B)[tile_work.tiled_coord.k()]; + } + + int n_begin = tile_work.tiled_coord.n() * Mma::Shape::kN; + int n_end = params.block_mapping.problem_size.n(); + return Mma::IteratorB( + params.params_B, + ptr_B, + { tile_work.k_end, n_end }, + threadIdx.x, + { tile_work.k_begin, n_begin }); + } + + + MCTLASS_DEVICE + void init_dp_tile_work( + TileWorkDesc &tile_work, + int tile_idx) + { + // The linear tile index + tile_work.tile_idx = tile_idx; + + // The first global-scoped MAC-iteration this threadblock will perform for this tile + tile_work.iter_begin = tile_idx * params.block_mapping.iters_per_tile(); + + // The number of MAC-iterations this threadblock will perform for this tile + tile_work.k_iters_remaining = params.block_mapping.iters_per_tile(); + + // The starting index in the k-domain for MAC-iterations this threadblock will perform for this tile + tile_work.k_begin = 0; + + // The ending index (one-past) in the k-domain for MAC-iterations this threadblock will perform for this tile + tile_work.k_end = params.block_mapping.problem_size.k(); + + // The location of this tile (in threadblock-tile coordinates) in the output matrix + tile_work.tiled_coord = params.block_mapping.get_tile_offset(tile_work.tile_idx); + } + + + MCTLASS_DEVICE + void init_sk_tile_work( + TileWorkDesc &tile_work, + int tile_idx, + int block_iter_begin, + int block_iter_end) + { + // The linear tile index + tile_work.tile_idx = tile_idx; + + // The first global-scoped MAC-iteration for this tile + int tile_iter_begin = tile_idx * params.block_mapping.iters_per_tile(); + + // The first global-scoped MAC-iteration this threadblock will perform for this tile + tile_work.iter_begin = max(block_iter_begin, tile_iter_begin); + + // The first tile-scoped MAC-iteration this threadblock will perform for this tile + int k_iter_begin = tile_work.iter_begin - tile_iter_begin; + + // The last (one past) tile-scoped MAC-iteration this threadblock will perform for this tile + int k_iter_end = block_iter_end - tile_iter_begin; + + // The number of MAC-iterations this threadblock will perform for this tile + tile_work.k_iters_remaining = k_iter_end - k_iter_begin; + + // The starting index in the k-domain for MAC-iterations this threadblock will perform for this tile + tile_work.k_begin = k_iter_begin * Mma::Shape::kK; + + // The ending index (one-past) in the k-domain for MAC-iterations this threadblock will perform for this tile + tile_work.k_end = min( + params.block_mapping.problem_size.k(), // extent of k domain + (k_iter_end * Mma::Shape::kK)); // extent of the threadblock's global iteration assignment + + // The location of this tile (in threadblock-tile coordinates) in the output matrix + tile_work.tiled_coord = params.block_mapping.get_tile_offset(tile_work.tile_idx); + } + + + /// Share accumulators with peers + MCTLASS_DEVICE + void share_accumulators( + AccumulatorTile const &accumulator_tile, + int block_idx, + int first_block_idx) + { + AccumulatorTile *accum_tile_workspace = reinterpret_cast(params.partials_workspace); + + int accum_tile_offset = first_block_idx * kThreadCount; + + if (block_idx == first_block_idx) + { + // First peer initializes the workspace partials + BlockStripedReduceT::store(accum_tile_workspace + accum_tile_offset, accumulator_tile, thread_idx); + } + else + { + // Subsequent peers atomically accumulate into the workspace partials + if (ThreadblockSwizzle::kReductionStrategy == ThreadblockSwizzle::kAtomic) + { + // Non-deterministic reduction order: wait for the first peer to have initialized the partials before we add to them + Barrier::wait_lt(params.barrier_workspace, thread_idx, first_block_idx, 1); + } + else + { + // Turnstile reduction order: wait until the previous peer has written + int wait_count = block_idx - first_block_idx; + Barrier::wait_eq(params.barrier_workspace, thread_idx, first_block_idx, wait_count); + } + + // Perform reduction in workspace + BlockStripedReduceT::reduce(accum_tile_workspace + accum_tile_offset, accumulator_tile, thread_idx); + } + + // Signal our arrival + Barrier::arrive_inc(params.barrier_workspace, thread_idx, first_block_idx); + } + + + /// Acquire accumulators from peers + MCTLASS_DEVICE + void acquire_accumulators( + AccumulatorTile &accumulator_tile, + int block_idx, + int first_block_idx) + { + AccumulatorTile *accum_tile_workspace = reinterpret_cast(params.partials_workspace); + + // Wait for arrival + int num_carry_in = block_idx - first_block_idx; + Barrier::wait_eq_reset(params.barrier_workspace, thread_idx, first_block_idx, num_carry_in); + + // Load and add peer-partials accumulator tile to local accumulator tile + int accum_tile_offset = first_block_idx * kThreadCount; + BlockStripedReduceT::load_add(accumulator_tile, accum_tile_workspace + accum_tile_offset, thread_idx); + } + + + /// Perform epilogue computations and output + MCTLASS_DEVICE + void do_epilogue( + TileWorkDesc &tile_work, + AccumulatorTile &accumulator_tile) + { + ElementC *ptr_C1 = static_cast(params.ptr_C1); + ElementC *ptr_C2 = static_cast(params.ptr_C2); + ElementC *ptr_D = static_cast(params.ptr_D); + typename Epilogue::ElementTensor *ptr_Tensor = static_cast(params.ptr_Tensor); + + // Define the reduction output pointer and move to the appropriate place + typename Epilogue::ElementVector *ptr_Vector = + static_cast(params.ptr_Vector); + + // Update pointers for batched/array mode(s) + if (params.mode == GemmUniversalMode::kBatched) { + ptr_C1 += tile_work.tiled_coord.k() * params.batch_stride_C1; + if (ptr_C2) { + ptr_C2 += tile_work.tiled_coord.k() * params.batch_stride_C2; + } + ptr_D += tile_work.tiled_coord.k() * params.batch_stride_D; + if (ptr_Tensor) { + ptr_Tensor += tile_work.tiled_coord.k() * params.batch_stride_Tensor; + } + if (ptr_Vector) { + ptr_Vector += tile_work.tiled_coord.k() * params.batch_stride_Vector; + } + } + if (params.mode == GemmUniversalMode::kArray) { + ptr_C1 = static_cast(params.ptr_C1)[tile_work.tiled_coord.k()]; + if (ptr_C2) { + ptr_C2 = static_cast(params.ptr_C2)[tile_work.tiled_coord.k()]; + } + ptr_D = static_cast(params.ptr_D)[tile_work.tiled_coord.k()]; + if (ptr_Tensor) { + ptr_Tensor = static_cast(params.ptr_Tensor)[tile_work.tiled_coord.k()]; + } + if (ptr_Vector) { + ptr_Vector = static_cast(params.ptr_Vector)[tile_work.tiled_coord.k()]; + } + } + + // Location of this tile in item-coords + MatrixCoord threadblock_item_begin( + tile_work.tiled_coord.m() * Mma::Shape::kM, + tile_work.tiled_coord.n() * Mma::Shape::kN + ); + + // Tile iterator loading from residual1. + typename Epilogue::OutputTileIterator iterator_C1( + params.params_C1, + ptr_C1, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Tile iterator loading from residual2. + typename Epilogue::OutputTileIterator iterator_C2( + params.params_C2, + ptr_C2, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Additional tensor to load from + typename Epilogue::TensorTileIterator tensor_iterator( + params.params_Tensor, + ptr_Tensor, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Move to appropriate location for this output tile + if (ptr_Vector) { + ptr_Vector += threadblock_item_begin.column() + tile_work.tiled_coord.m() * params.ldr; + } + + // Execute the epilogue operator to update the destination tensor. + epilogue( + EpilogueOutputOp(params.output_op), + ptr_Vector, + iterator_D, + accumulator_tile, + iterator_C1, + iterator_C2, + tensor_iterator, + params.block_mapping.problem_size.mn(), + threadblock_item_begin); + } + + + MCTLASS_DEVICE + void separate_reduction(int reduce_idx) + { + int peer_idx_begin, peer_idx_last, reduce_tile_idx, reduce_fragment_idx; + + // Reduce by sk-tile (every tile contributed to by one or more blocks) + reduce_tile_idx = reduce_idx / Epilogue::kAccumulatorFragments; + reduce_fragment_idx = reduce_idx % Epilogue::kAccumulatorFragments; + + int iter_tile_first = reduce_tile_idx * params.block_mapping.iters_per_tile(); + int iter_tile_last = iter_tile_first + params.block_mapping.iters_per_tile() - 1; + + peer_idx_begin = params.block_mapping.get_sk_block_idx(iter_tile_first); + peer_idx_last = params.block_mapping.get_sk_block_idx(iter_tile_last); + + // Wait for peers to complete + int peer_idx_end = peer_idx_last + 1; + int num_peers = peer_idx_end - peer_idx_begin; + Barrier::wait_eq_reset( + params.barrier_workspace, + thread_idx, + (reduce_tile_idx * Epilogue::kAccumulatorFragments) + reduce_fragment_idx, + num_peers); + + /// The location of this tile (in threadblock-tile coordinates) in the output matrix + GemmCoord tiled_coord = params.block_mapping.get_tile_offset(reduce_tile_idx); + + // Location of this tile in item-coords + MatrixCoord threadblock_item_begin( + tiled_coord.m() * Mma::Shape::kM, + tiled_coord.n() * Mma::Shape::kN + ); + + ElementC *ptr_C1 = static_cast(params.ptr_C1); + ElementC *ptr_C2 = static_cast(params.ptr_C2); + ElementC *ptr_D = static_cast(params.ptr_D); + typename Epilogue::ElementTensor *ptr_Tensor = static_cast(params.ptr_Tensor); + + // Define the reduction output pointer and move to the appropriate place + typename Epilogue::ElementVector *ptr_Vector = + static_cast(params.ptr_Vector); + + // Tile iterator loading from residual1. + typename Epilogue::OutputTileIterator iterator_C1( + params.params_C1, + ptr_C1, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Tile iterator loading from residual2. + typename Epilogue::OutputTileIterator iterator_C2( + params.params_C2, + ptr_C2, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Additional tensor to load from + typename Epilogue::TensorTileIterator tensor_iterator( + params.params_Tensor, + ptr_Tensor, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Move to appropriate location for this output tile + if (ptr_Vector) { + ptr_Vector += threadblock_item_begin.column() + tiled_coord.m() * params.ldr; + } + + // Execute the epilogue operator to update the destination tensor. + epilogue.reduce( + peer_idx_begin, + peer_idx_end, + reduce_fragment_idx, + params.partials_workspace, + EpilogueOutputOp(params.output_op), + ptr_Vector, + iterator_D, + iterator_C1, + iterator_C2, + tensor_iterator, + params.block_mapping.problem_size.mn(), + threadblock_item_begin); + } + + + MCTLASS_DEVICE + void process_tile( + TileWorkDesc tile_work, + int block_idx, + int dp_start_block_idx, + int block_iter_begin) + { + // Initialize input iterators + typename Mma::IteratorA iterator_A = init_iterator_A(tile_work, params.mode); + typename Mma::IteratorB iterator_B = init_iterator_B(tile_work, params.mode); + + // Initialize accumulators + AccumulatorTile accumulator_tile; + accumulator_tile.clear(); + + // Initialize MMA abstraction + Mma mma( + shared_storage.main_loop, + thread_idx, + warp_idx, + lane_idx); + + // Perform this tile's range of multiply-accumulate (MAC) iterations + mma(tile_work.k_iters_remaining, accumulator_tile, iterator_A, iterator_B, accumulator_tile); + + if ((ThreadblockSwizzle::kReductionStrategy == ThreadblockSwizzle::kAtomic) || + (params.block_mapping.reduction_blocks == 0) || + (block_idx >= dp_start_block_idx)) + { + // + // Cooperative SK peer reduction or DP block + // + + int first_block_idx = params.block_mapping.get_first_block_idx(tile_work.tile_idx, block_idx); + + if (!tile_work.tile_finished(params)) { + // Non "finishing" SK blocks must share their partial accumulator sums through global scratch workspace + share_accumulators(accumulator_tile, block_idx, first_block_idx); + } + else + { + // DP blocks and "finishing" SK blocks must perform epilogue operations and write the output tile + if (!tile_work.tile_started()) + { + // A "finishing" SK block must first aggregate its accumulator partial sums with those shared by peer threadblocks + acquire_accumulators(accumulator_tile, block_idx, first_block_idx); + } + + do_epilogue(tile_work, accumulator_tile); + } + } + else + { + // + // Separate peer reduction + // + + // Share accumulator partial sums with peer threadblock(s) through scratch workspace + epilogue.share(block_idx, params.partials_workspace, accumulator_tile, tile_work.tile_started()); + + // Signal arrival + Barrier::arrive_range_inc( + params.barrier_workspace, + thread_idx, + tile_work.tile_idx * Epilogue::kAccumulatorFragments, + Epilogue::kAccumulatorFragments); + } + } + + + /// Executes one GEMM + MCTLASS_DEVICE + void gemm() + { + // Initialize block's iteration range + int tile_idx = 0; + int block_iter_begin = 0; + int block_iters_remaining = 0; + + int block_idx = params.block_mapping.get_block_idx(); + + int sk_padding_start_block_idx = params.block_mapping.sk_regions() * params.block_mapping.sk_blocks_per_region(); + int dp_start_block_idx = params.block_mapping.sk_waves * params.block_mapping.avail_sms; + int reduce_start_block_idx = dp_start_block_idx + params.block_mapping.dp_blocks; + int grid_padding_start_block_idx = reduce_start_block_idx + params.block_mapping.reduction_blocks; + + // Initialize tile work descriptor + TileWorkDesc tile_work; + + bool dp_block = (block_idx >= dp_start_block_idx) && (block_idx < reduce_start_block_idx); + bool sk_block = (block_idx < sk_padding_start_block_idx); + bool reduce_block = (block_idx >= reduce_start_block_idx) && + (block_idx < grid_padding_start_block_idx) && + (ThreadblockSwizzle::kReductionStrategy == ThreadblockSwizzle::kMixed); + + if (dp_block) + { + // This is a DP block + int dp_block_idx = block_idx - dp_start_block_idx; + int first_dp_tile = (params.block_mapping.cohort_raster) ? 0 : params.block_mapping.sk_tiles; + + // Blocks in first DP wave get configured number of tiles + tile_idx = first_dp_tile + dp_block_idx; + int tile_allottment = params.block_mapping.dp_first_wave_tiles; + + // Blocks in subsequent DP waves get 1 tile + if (dp_block_idx >= params.block_mapping.avail_sms) { + tile_allottment = 1; + tile_idx += (params.block_mapping.dp_first_wave_tiles - 1) * params.block_mapping.avail_sms; + } + + block_iters_remaining = params.block_mapping.iters_per_tile() * tile_allottment; + + init_dp_tile_work(tile_work, tile_idx); + + // DP blocks exit if out of bounds or overlap an SK tile (only possible during cohort rasterization, where dp_first_wave_tiles must be 1) + if ((tile_idx < params.block_mapping.sk_tiles) || + (tile_work.tiled_coord.m() >= params.block_mapping.tiled_shape().m()) || + (tile_work.tiled_coord.n() >= params.block_mapping.tiled_shape().n())) + { + return; + } + } + else if (sk_block) + { + // This is a SK block + int block_iter_end; + params.block_mapping.get_iter_extents(block_idx, block_iter_begin, block_iter_end); + block_iters_remaining = block_iter_end - block_iter_begin; + + tile_idx = params.block_mapping.get_sk_tile_idx(block_iter_end - 1); + init_sk_tile_work(tile_work, tile_idx, block_iter_begin, block_iter_begin + block_iters_remaining); + } + else + { + if (reduce_block) + { + // This is a reduction threadblock + int reduce_block_idx = block_idx - reduce_start_block_idx; + separate_reduction(reduce_block_idx); + } + + return; + } + + // Iteration-processing loop body + MCTLASS_PRAGMA_NO_UNROLL + while (true) + { + // Perform this block's share of work for this tile + process_tile( + tile_work, + block_idx, + dp_start_block_idx, + block_iter_begin); + + block_iters_remaining -= tile_work.k_iters_remaining; + + if (block_iters_remaining == 0) + { + break; + } + + // Continue to next tile + __syncthreads(); + + if (block_idx >= dp_start_block_idx) + { + // DP block consume their tiles at stride + tile_idx += params.block_mapping.avail_sms; + init_dp_tile_work(tile_work, tile_idx); + } + else + { + // SK blocks consume their tiles in backwards order + tile_idx--; + init_sk_tile_work(tile_work, tile_idx, block_iter_begin, block_iter_begin + block_iters_remaining); + } + } + + } + + +public: + + // + // Device-only API + // + + // Factory invocation + MCTLASS_DEVICE + static void invoke( + Params const ¶ms, + SharedStorage &shared_storage) + { + GemmStreamkWithFusedEpilogue op(params, shared_storage); + op(); + } + + + // Constructor + MCTLASS_DEVICE + GemmStreamkWithFusedEpilogue( + Params const ¶ms, + SharedStorage &shared_storage) + : + params(params), + shared_storage(shared_storage), + thread_idx(threadIdx.x), + warp_idx(__shfl_sync(0xffffffff, threadIdx.x / 64, 0)), // broadcast the warp_id computed by lane 0 to ensure dependent code + lane_idx(threadIdx.x % 64), + epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx) + {} + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()() { + // Generic SK code path + gemm(); + + } +}; + + +// GemmStreamkWithFusedEpilogue with one source +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_ ///! Threadblock swizzling function +> +struct GemmStreamkWithFusedEpilogue { + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + + /// The per-thread tile of raw accumulators + using AccumulatorTile = typename Mma::FragmentC; + + static ComplexTransform const kTransformA = Mma::kTransformA; + static ComplexTransform const kTransformB = Mma::kTransformB; + using Operator = typename Mma::Operator; + + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::InstructionShape; + using ArchTag = typename Mma::ArchTag; + + static int const kStages = Mma::kStages; + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Workspace bytes per thread block + static size_t const kWorkspaceBytesPerBlock = + __NV_STD_MAX( + kThreadCount * sizeof(AccumulatorTile), + Epilogue::kWorkspaceBytesPerBlock); + + /// Block-striped reduction utility + using BlockStripedReduceT = BlockStripedReduce; + + + + // + // Structures + // + + /// Argument structure + struct Arguments + { + + // + // Data members + // + + GemmUniversalMode mode; + GemmCoord problem_size; + int batch_count; // Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor + + typename EpilogueOutputOp::Params epilogue; + + void const * ptr_A; + void const * ptr_B; + void const * ptr_C; + void * ptr_D; + + void * ptr_Vector; + void * ptr_Tensor; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + int64_t batch_stride_D; + int64_t batch_stride_Vector; + int64_t batch_stride_Tensor; + + typename LayoutA::Stride::Index lda; + typename LayoutB::Stride::Index ldb; + typename LayoutC::Stride::Index ldc; + typename LayoutC::Stride::Index ldd; + typename LayoutC::Stride::Index ldr; + typename LayoutC::Stride::Index ldt; + + int avail_sms; /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling) + + + // + // Methods + // + + /// Default Constructor + Arguments(): + mode(GemmUniversalMode::kGemm), + batch_count(1), + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + avail_sms(-1) + {} + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_split, /// Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor (1 defaults to StreamK, >1 emulates Split-K) + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void const * ptr_C, + void * ptr_D, + void * ptr_Vector, + void * ptr_Tensor, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C, + int64_t batch_stride_D, + int64_t batch_stride_Vector, + int64_t batch_stride_Tensor, + typename LayoutA::Stride::Index lda, + typename LayoutB::Stride::Index ldb, + typename LayoutC::Stride::Index ldc, + typename LayoutC::Stride::Index ldd, + typename LayoutC::Stride::Index ldr, + typename LayoutC::Stride::Index ldt, + int avail_sms = -1) /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling) + : + mode(mode), + problem_size(problem_size), + batch_count(batch_split), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), + ptr_Vector(ptr_Vector), + ptr_Tensor(ptr_Tensor), + batch_stride_A(batch_stride_A), + batch_stride_B(batch_stride_B), + batch_stride_C(batch_stride_C), + batch_stride_Vector(batch_stride_Vector), + batch_stride_Tensor(batch_stride_Tensor), + lda(lda), ldb(ldb), ldc(ldc), ldd(ldd), ldr(ldr), ldt(ldt), avail_sms(avail_sms) + { + MCTLASS_TRACE_HOST("GemmStreamkWithFusedEpilogue::Arguments::Arguments() - problem_size: " << problem_size); + MCTLASS_TRACE_HOST(" ptr_Vector: " << (void *)this->ptr_Vector); + MCTLASS_TRACE_HOST(" ptr_Tensor: " << (void *)this->ptr_Tensor); + MCTLASS_TRACE_HOST(" ldr: " << this->ldr); + MCTLASS_TRACE_HOST(" ldt: " << this->ldt); + MCTLASS_TRACE_HOST(" avail_sms: " << this->avail_sms); + } + + /// Returns arguments for the transposed problem + Arguments transposed_problem() const { + Arguments args(*this); + + std::swap(args.problem_size.m(), args.problem_size.n()); + std::swap(args.ptr_A, args.ptr_B); + std::swap(args.lda, args.ldb); + std::swap(args.batch_stride_A, args.batch_stride_B); + + return args; + } + }; + + + /// Parameters structure + struct Params + { + + public: + + // + // Data members + // + + void * ptr_A; + void * ptr_B; + + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorB::Params params_B; + + int64_t batch_stride_A; + int64_t batch_stride_B; + + GemmUniversalMode mode; + + ThreadblockSwizzle block_mapping; + + void *barrier_workspace; + void *partials_workspace; + + typename EpilogueOutputOp::Params output_op; + + void * ptr_C; + void * ptr_D; + void * ptr_Tensor; + void * ptr_Vector; + + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::TensorTileIterator::Params params_Tensor; + + int64_t batch_stride_C; + int64_t batch_stride_D; + int64_t batch_stride_Vector; + int64_t batch_stride_Tensor; + + + typename LayoutC::Stride::Index ldr; + + protected: + + // + // Host-only dispatch-utilities + // + + /// Pad the given allocation size up to the nearest cache line + static size_t cacheline_align_up(size_t size) + { + static const int CACHELINE_SIZE = 128; + return (size + CACHELINE_SIZE - 1) / CACHELINE_SIZE * CACHELINE_SIZE; + } + + /// Get the workspace size needed for barrier + size_t get_barrier_workspace_size() const + { + // For atomic reduction, each SK-block needs a synchronization flag. For parallel reduction, + // each reduction block needs its own synchronization flag. + int sk_blocks = block_mapping.sk_regions() * block_mapping.sk_blocks_per_region(); + int num_flags = fast_max(sk_blocks, block_mapping.reduction_blocks); + + return cacheline_align_up(sizeof(typename Barrier::T) * num_flags); + } + + /// Get the workspace size needed for intermediate partial sums + size_t get_partials_workspace_size() const + { + int sk_blocks = block_mapping.sk_regions() * block_mapping.sk_blocks_per_region(); + return cacheline_align_up(kWorkspaceBytesPerBlock * sk_blocks); + } + + + public: + // + // Host dispatch API + // + + /// Default constructor + Params() = default; + + /// Constructor + Params( + Arguments const &args, /// GEMM application arguments + int device_sms, /// Number of SMs on the device + int sm_occupancy) /// Kernel SM occupancy (in thread blocks) + : + params_A(args.lda), + params_B(args.ldb), + params_C(args.ldc), + params_D(args.ldd), + params_Tensor(args.ldt), + output_op(args.epilogue), + mode(args.mode), + ptr_A(const_cast(args.ptr_A)), + ptr_B(const_cast(args.ptr_B)), + ptr_C(const_cast(args.ptr_C)), + ptr_D(args.ptr_D), + ptr_Vector(args.ptr_Vector), + ldr(args.ldr), + ptr_Tensor(args.ptr_Tensor), + batch_stride_A(args.batch_stride_A), + batch_stride_B(args.batch_stride_B), + batch_stride_C(args.batch_stride_C), + batch_stride_D(args.batch_stride_D), + batch_stride_Vector(args.batch_stride_Vector), + batch_stride_Tensor(args.batch_stride_Tensor), + barrier_workspace(nullptr), + partials_workspace(nullptr) + { + MCTLASS_TRACE_HOST("GemmStreamkWithFusedEpilogue::Params::Params() - problem_size: " << problem_size); + MCTLASS_TRACE_HOST(" ptr_Vector: " << (void *)this->ptr_Vector); + MCTLASS_TRACE_HOST(" ptr_Tensor: " << (void *)this->ptr_Tensor); + MCTLASS_TRACE_HOST(" ldr: " << this->ldr); + MCTLASS_TRACE_HOST(" ldt: " << args.ldt); + MCTLASS_TRACE_HOST(" avail_sms: " << avail_sms); + + // Number of SMs to make available for StreamK decomposition + int avail_sms = (args.avail_sms == -1) ? + device_sms : + fast_min(args.avail_sms, device_sms); + + // Initialize the block mapping structure + block_mapping = ThreadblockSwizzle( + typename ThreadblockSwizzle::template KernelTraits(), + args.mode, + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count, + sm_occupancy, + device_sms, + avail_sms); + } + + /// Returns the workspace size (in bytes) needed for these parameters + size_t get_workspace_size() const + { + return + get_barrier_workspace_size() + + get_partials_workspace_size(); + } + + + /// Assign and initialize the specified workspace buffer. Assumes + /// the memory allocated to workspace is at least as large as get_workspace_size(). + Status init_workspace( + void *workspace, + mcStream_t stream = nullptr) + { + uint8_t *ptr = static_cast(workspace); + + // Establish partials workspace + partials_workspace = nullptr; + size_t partials_workspace_bytes = get_partials_workspace_size(); + if (partials_workspace_bytes > 0) + { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + partials_workspace = ptr; + ptr += partials_workspace_bytes; + } + + // Establish barrier workspace + barrier_workspace = nullptr; + size_t barrier_workspace_bytes = get_barrier_workspace_size(); + if (barrier_workspace_bytes > 0) + { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + barrier_workspace = ptr; + ptr += barrier_workspace_bytes; + } + + // Zero-initialize barrier workspace + if (barrier_workspace) + { + size_t barrier_workspace_bytes = get_barrier_workspace_size(); + + MCTLASS_TRACE_HOST(" Initialize " << barrier_workspace_bytes << " barrier bytes"); + + mcError_t result = mcMemsetAsync( + barrier_workspace, + 0, + barrier_workspace_bytes, + stream); + + if (result != mcSuccess) { + MCTLASS_TRACE_HOST(" mcMemsetAsync() returned error " << mcGetErrorString(result)); + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + + /// Returns the GEMM volume in thread block tiles + mctlass::gemm::GemmCoord get_tiled_shape() const + { + return block_mapping.tiled_shape(); + } + + + /// Returns the total number of thread blocks to launch + int get_grid_blocks() const + { + dim3 grid_dims = get_grid_dims(); + return grid_dims.x * grid_dims.y * grid_dims.z; + } + + + /// Returns the grid extents in thread blocks to launch + dim3 get_grid_dims() const + { + return block_mapping.get_grid_dims(); + } + + /// Lightweight update given a subset of arguments. Problem geometry is assumed + /// to remain the same. + MCTLASS_HOST_DEVICE + void update(Arguments const &args) + { + ptr_A = const_cast(args.ptr_A); + ptr_B = const_cast(args.ptr_B); + ptr_C = const_cast(args.ptr_C); + ptr_D = args.ptr_D; + + ptr_Vector = args.ptr_Vector; + ldr = args.ldr; + ptr_Tensor = args.ptr_Tensor; + + batch_stride_A = args.batch_stride_A; + batch_stride_B = args.batch_stride_B; + batch_stride_C = args.batch_stride_C; + batch_stride_D = args.batch_stride_D; + batch_stride_Vector = args.batch_stride_Vector; + batch_stride_Tensor = args.batch_stride_Tensor; + + output_op = args.epilogue; + + MCTLASS_TRACE_HOST("GemmStreamkWithFusedEpilogue::Params::update()"); + MCTLASS_TRACE_HOST(" ptr_Vector: " << (void *)this->ptr_Vector); + MCTLASS_TRACE_HOST(" ptr_Tensor: " << (void *)this->ptr_Tensor); + MCTLASS_TRACE_HOST(" ldr: " << this->ldr); + } + }; + + /// Tile work descriptor + struct TileWorkDesc + { + /// The linear tile index + int tile_idx; + + /// The location of this tile (in threadblock-tile coordinates) in the output matrix + mctlass::gemm::GemmCoord tiled_coord; + + // The first global-scoped MAC-iteration this threadblock will perform for this tile + int iter_begin; + + // The starting index in the k-domain for MAC-iterations this threadblock will perform for this tile + int k_begin; + + // The ending index (one-past) in the k-domain for MAC-iterations this threadblock will perform for this tile + int k_end; + + /// The number of remaining MAC-iterations this threadblock will perform for this tile + int k_iters_remaining; + + // Whether this block will perform the first iteration of this tile + MCTLASS_DEVICE + bool tile_started() + { + return (k_begin == 0); + } + + // Whether this block will perform the last iteration of this tile + MCTLASS_DEVICE + bool tile_finished(Params const ¶ms) + { + return (k_end == params.block_mapping.problem_size.k()); + } + }; + + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + +protected: + + // + // Data members + // + + /// GEMM problem parameters + Params const ¶ms; + + /// Shared storage reference + SharedStorage &shared_storage; + + /// ID within the threadblock + int thread_idx; + + /// ID of warp + int warp_idx; + + /// ID of each thread within a warp + int lane_idx; + + /// Threadblock scoped epilogue + Epilogue epilogue; + + +public: + + // + // Host dispatch API + // + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size) { + + MCTLASS_TRACE_HOST("GemmStreamkWithFusedEpilogue::can_implement()"); + + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + bool isAMisaligned = false; + bool isBMisaligned = false; + bool isCMisaligned = false; + + if (platform::is_same::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } else if (platform::is_same::value) { + isAMisaligned = problem_size.m() % kAlignmentA; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } + + if (platform::is_same::value) { + isBMisaligned = problem_size.n() % kAlignmentB; + } else if (platform::is_same::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } + + if (platform::is_same::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } else if (platform::is_same::value) { + isCMisaligned = problem_size.m() % kAlignmentC; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } + + if (isAMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for A operand"); + return Status::kErrorMisalignedOperand; + } + + if (isBMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for B operand"); + return Status::kErrorMisalignedOperand; + } + + if (isCMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for C operand"); + return Status::kErrorMisalignedOperand; + } + + MCTLASS_TRACE_HOST(" returning kSuccess"); + + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + +protected: + + // + // Device-only utility methods + // + + /// Iterator for fetching tile fragments from A + MCTLASS_DEVICE + typename Mma::IteratorA init_iterator_A( + TileWorkDesc &tile_work, + GemmUniversalMode mode) + { + // The input A matrix + ElementA *ptr_A = static_cast(params.ptr_A); + + // Update input pointers based on batched/array mode + if (mode == GemmUniversalMode::kBatched) { + ptr_A += tile_work.tiled_coord.k() * params.batch_stride_A; + } + if (mode == GemmUniversalMode::kArray) { + ptr_A = static_cast(params.ptr_A)[tile_work.tiled_coord.k()]; + } + + int m_begin = tile_work.tiled_coord.m() * Mma::Shape::kM; + int m_end = params.block_mapping.problem_size.m(); + return Mma::IteratorA( + params.params_A, + ptr_A, + { m_end, tile_work.k_end }, + threadIdx.x, + { m_begin, tile_work.k_begin }); + + } + + + /// Iterator for fetching tile fragments from B + MCTLASS_DEVICE + typename Mma::IteratorB init_iterator_B( + TileWorkDesc &tile_work, + GemmUniversalMode mode) + { + // The input B matrix + ElementB *ptr_B = static_cast(params.ptr_B); + + // Update input pointers based on batched/array mode + if (mode == GemmUniversalMode::kBatched) { + ptr_B += tile_work.tiled_coord.k() * params.batch_stride_B; + } + if (mode == GemmUniversalMode::kArray) { + ptr_B = static_cast(params.ptr_B)[tile_work.tiled_coord.k()]; + } + + int n_begin = tile_work.tiled_coord.n() * Mma::Shape::kN; + int n_end = params.block_mapping.problem_size.n(); + return Mma::IteratorB( + params.params_B, + ptr_B, + { tile_work.k_end, n_end }, + threadIdx.x, + { tile_work.k_begin, n_begin }); + } + + + MCTLASS_DEVICE + void init_dp_tile_work( + TileWorkDesc &tile_work, + int tile_idx) + { + // The linear tile index + tile_work.tile_idx = tile_idx; + + // The first global-scoped MAC-iteration this threadblock will perform for this tile + tile_work.iter_begin = tile_idx * params.block_mapping.iters_per_tile(); + + // The number of MAC-iterations this threadblock will perform for this tile + tile_work.k_iters_remaining = params.block_mapping.iters_per_tile(); + + // The starting index in the k-domain for MAC-iterations this threadblock will perform for this tile + tile_work.k_begin = 0; + + // The ending index (one-past) in the k-domain for MAC-iterations this threadblock will perform for this tile + tile_work.k_end = params.block_mapping.problem_size.k(); + + // The location of this tile (in threadblock-tile coordinates) in the output matrix + tile_work.tiled_coord = params.block_mapping.get_tile_offset(tile_work.tile_idx); + } + + + MCTLASS_DEVICE + void init_sk_tile_work( + TileWorkDesc &tile_work, + int tile_idx, + int block_iter_begin, + int block_iter_end) + { + // The linear tile index + tile_work.tile_idx = tile_idx; + + // The first global-scoped MAC-iteration for this tile + int tile_iter_begin = tile_idx * params.block_mapping.iters_per_tile(); + + // The first global-scoped MAC-iteration this threadblock will perform for this tile + tile_work.iter_begin = max(block_iter_begin, tile_iter_begin); + + // The first tile-scoped MAC-iteration this threadblock will perform for this tile + int k_iter_begin = tile_work.iter_begin - tile_iter_begin; + + // The last (one past) tile-scoped MAC-iteration this threadblock will perform for this tile + int k_iter_end = block_iter_end - tile_iter_begin; + + // The number of MAC-iterations this threadblock will perform for this tile + tile_work.k_iters_remaining = k_iter_end - k_iter_begin; + + // The starting index in the k-domain for MAC-iterations this threadblock will perform for this tile + tile_work.k_begin = k_iter_begin * Mma::Shape::kK; + + // The ending index (one-past) in the k-domain for MAC-iterations this threadblock will perform for this tile + tile_work.k_end = min( + params.block_mapping.problem_size.k(), // extent of k domain + (k_iter_end * Mma::Shape::kK)); // extent of the threadblock's global iteration assignment + + // The location of this tile (in threadblock-tile coordinates) in the output matrix + tile_work.tiled_coord = params.block_mapping.get_tile_offset(tile_work.tile_idx); + } + + + /// Share accumulators with peers + MCTLASS_DEVICE + void share_accumulators( + AccumulatorTile const &accumulator_tile, + int block_idx, + int first_block_idx) + { + AccumulatorTile *accum_tile_workspace = reinterpret_cast(params.partials_workspace); + + int accum_tile_offset = first_block_idx * kThreadCount; + + if (block_idx == first_block_idx) + { + // First peer initializes the workspace partials + BlockStripedReduceT::store(accum_tile_workspace + accum_tile_offset, accumulator_tile, thread_idx); + } + else + { + // Subsequent peers atomically accumulate into the workspace partials + if (ThreadblockSwizzle::kReductionStrategy == ThreadblockSwizzle::kAtomic) + { + // Non-deterministic reduction order: wait for the first peer to have initialized the partials before we add to them + Barrier::wait_lt(params.barrier_workspace, thread_idx, first_block_idx, 1); + } + else + { + // Turnstile reduction order: wait until the previous peer has written + int wait_count = block_idx - first_block_idx; + Barrier::wait_eq(params.barrier_workspace, thread_idx, first_block_idx, wait_count); + } + + // Perform reduction in workspace + BlockStripedReduceT::reduce(accum_tile_workspace + accum_tile_offset, accumulator_tile, thread_idx); + } + + // Signal our arrival + Barrier::arrive_inc(params.barrier_workspace, thread_idx, first_block_idx); + } + + + /// Acquire accumulators from peers + MCTLASS_DEVICE + void acquire_accumulators( + AccumulatorTile &accumulator_tile, + int block_idx, + int first_block_idx) + { + AccumulatorTile *accum_tile_workspace = reinterpret_cast(params.partials_workspace); + + // Wait for arrival + int num_carry_in = block_idx - first_block_idx; + Barrier::wait_eq_reset(params.barrier_workspace, thread_idx, first_block_idx, num_carry_in); + + // Load and add peer-partials accumulator tile to local accumulator tile + int accum_tile_offset = first_block_idx * kThreadCount; + BlockStripedReduceT::load_add(accumulator_tile, accum_tile_workspace + accum_tile_offset, thread_idx); + } + + + /// Perform epilogue computations and output + MCTLASS_DEVICE + void do_epilogue( + TileWorkDesc &tile_work, + AccumulatorTile &accumulator_tile) + { + ElementC *ptr_C = static_cast(params.ptr_C); + ElementC *ptr_D = static_cast(params.ptr_D); + typename Epilogue::ElementTensor *ptr_Tensor = static_cast(params.ptr_Tensor); + + // Define the reduction output pointer and move to the appropriate place + typename Epilogue::ElementVector *ptr_Vector = + static_cast(params.ptr_Vector); + + // Update pointers for batched/array mode(s) + if (params.mode == GemmUniversalMode::kBatched) { + ptr_C += tile_work.tiled_coord.k() * params.batch_stride_C; + ptr_D += tile_work.tiled_coord.k() * params.batch_stride_D; + if (ptr_Tensor) { + ptr_Tensor += tile_work.tiled_coord.k() * params.batch_stride_Tensor; + } + if (ptr_Vector) { + ptr_Vector += tile_work.tiled_coord.k() * params.batch_stride_Vector; + } + } + if (params.mode == GemmUniversalMode::kArray) { + ptr_C = static_cast(params.ptr_C)[tile_work.tiled_coord.k()]; + ptr_D = static_cast(params.ptr_D)[tile_work.tiled_coord.k()]; + if (ptr_Tensor) { + ptr_Tensor = static_cast(params.ptr_Tensor)[tile_work.tiled_coord.k()]; + } + if (ptr_Vector) { + ptr_Vector = static_cast(params.ptr_Vector)[tile_work.tiled_coord.k()]; + } + } + + // Location of this tile in item-coords + MatrixCoord threadblock_item_begin( + tile_work.tiled_coord.m() * Mma::Shape::kM, + tile_work.tiled_coord.n() * Mma::Shape::kN + ); + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + ptr_C, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Additional tensor to load from + typename Epilogue::TensorTileIterator tensor_iterator( + params.params_Tensor, + ptr_Tensor, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Move to appropriate location for this output tile + if (ptr_Vector) { + ptr_Vector += threadblock_item_begin.column() + tile_work.tiled_coord.m() * params.ldr; + } + + // Execute the epilogue operator to update the destination tensor. + epilogue( + EpilogueOutputOp(params.output_op), + ptr_Vector, + iterator_D, + accumulator_tile, + iterator_C, + tensor_iterator, + params.block_mapping.problem_size.mn(), + threadblock_item_begin); + } + + + MCTLASS_DEVICE + void separate_reduction(int reduce_idx) + { + int peer_idx_begin, peer_idx_last, reduce_tile_idx, reduce_fragment_idx; + + // Reduce by sk-tile (every tile contributed to by one or more blocks) + reduce_tile_idx = reduce_idx / Epilogue::kAccumulatorFragments; + reduce_fragment_idx = reduce_idx % Epilogue::kAccumulatorFragments; + + int iter_tile_first = reduce_tile_idx * params.block_mapping.iters_per_tile(); + int iter_tile_last = iter_tile_first + params.block_mapping.iters_per_tile() - 1; + + peer_idx_begin = params.block_mapping.get_sk_block_idx(iter_tile_first); + peer_idx_last = params.block_mapping.get_sk_block_idx(iter_tile_last); + + // Wait for peers to complete + int peer_idx_end = peer_idx_last + 1; + int num_peers = peer_idx_end - peer_idx_begin; + Barrier::wait_eq_reset( + params.barrier_workspace, + thread_idx, + (reduce_tile_idx * Epilogue::kAccumulatorFragments) + reduce_fragment_idx, + num_peers); + + /// The location of this tile (in threadblock-tile coordinates) in the output matrix + GemmCoord tiled_coord = params.block_mapping.get_tile_offset(reduce_tile_idx); + + // Location of this tile in item-coords + MatrixCoord threadblock_item_begin( + tiled_coord.m() * Mma::Shape::kM, + tiled_coord.n() * Mma::Shape::kN + ); + + ElementC *ptr_C = static_cast(params.ptr_C); + ElementC *ptr_D = static_cast(params.ptr_D); + typename Epilogue::ElementTensor *ptr_Tensor = static_cast(params.ptr_Tensor); + + // Define the reduction output pointer and move to the appropriate place + typename Epilogue::ElementVector *ptr_Vector = + static_cast(params.ptr_Vector); + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + ptr_C, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Additional tensor to load from + typename Epilogue::TensorTileIterator tensor_iterator( + params.params_Tensor, + ptr_Tensor, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Move to appropriate location for this output tile + if (ptr_Vector) { + ptr_Vector += threadblock_item_begin.column() + tiled_coord.m() * params.ldr; + } + + // Execute the epilogue operator to update the destination tensor. + epilogue.reduce( + peer_idx_begin, + peer_idx_end, + reduce_fragment_idx, + params.partials_workspace, + EpilogueOutputOp(params.output_op), + ptr_Vector, + iterator_D, + iterator_C, + tensor_iterator, + params.block_mapping.problem_size.mn(), + threadblock_item_begin); + } + + + MCTLASS_DEVICE + void process_tile( + TileWorkDesc tile_work, + int block_idx, + int dp_start_block_idx, + int block_iter_begin) + { + // Initialize input iterators + typename Mma::IteratorA iterator_A = init_iterator_A(tile_work, params.mode); + typename Mma::IteratorB iterator_B = init_iterator_B(tile_work, params.mode); + + // Initialize accumulators + AccumulatorTile accumulator_tile; + accumulator_tile.clear(); + + // Initialize MMA abstraction + Mma mma( + shared_storage.main_loop, + thread_idx, + warp_idx, + lane_idx); + + // Perform this tile's range of multiply-accumulate (MAC) iterations + mma(tile_work.k_iters_remaining, accumulator_tile, iterator_A, iterator_B, accumulator_tile); + + if ((ThreadblockSwizzle::kReductionStrategy == ThreadblockSwizzle::kAtomic) || + (params.block_mapping.reduction_blocks == 0) || + (block_idx >= dp_start_block_idx)) + { + // + // Cooperative SK peer reduction or DP block + // + + int first_block_idx = params.block_mapping.get_first_block_idx(tile_work.tile_idx, block_idx); + + if (!tile_work.tile_finished(params)) { + // Non "finishing" SK blocks must share their partial accumulator sums through global scratch workspace + share_accumulators(accumulator_tile, block_idx, first_block_idx); + } + else + { + // DP blocks and "finishing" SK blocks must perform epilogue operations and write the output tile + if (!tile_work.tile_started()) + { + // A "finishing" SK block must first aggregate its accumulator partial sums with those shared by peer threadblocks + acquire_accumulators(accumulator_tile, block_idx, first_block_idx); + } + + do_epilogue(tile_work, accumulator_tile); + } + } + else + { + // + // Separate peer reduction + // + + // Share accumulator partial sums with peer threadblock(s) through scratch workspace + epilogue.share(block_idx, params.partials_workspace, accumulator_tile, tile_work.tile_started()); + + // Signal arrival + Barrier::arrive_range_inc( + params.barrier_workspace, + thread_idx, + tile_work.tile_idx * Epilogue::kAccumulatorFragments, + Epilogue::kAccumulatorFragments); + } + } + + + /// Executes one GEMM + MCTLASS_DEVICE + void gemm() + { + // Initialize block's iteration range + int tile_idx = 0; + int block_iter_begin = 0; + int block_iters_remaining = 0; + + int block_idx = params.block_mapping.get_block_idx(); + + int sk_padding_start_block_idx = params.block_mapping.sk_regions() * params.block_mapping.sk_blocks_per_region(); + int dp_start_block_idx = params.block_mapping.sk_waves * params.block_mapping.avail_sms; + int reduce_start_block_idx = dp_start_block_idx + params.block_mapping.dp_blocks; + int grid_padding_start_block_idx = reduce_start_block_idx + params.block_mapping.reduction_blocks; + + // Initialize tile work descriptor + TileWorkDesc tile_work; + + bool dp_block = (block_idx >= dp_start_block_idx) && (block_idx < reduce_start_block_idx); + bool sk_block = (block_idx < sk_padding_start_block_idx); + bool reduce_block = (block_idx >= reduce_start_block_idx) && + (block_idx < grid_padding_start_block_idx) && + (ThreadblockSwizzle::kReductionStrategy == ThreadblockSwizzle::kMixed); + + if (dp_block) + { + // This is a DP block + int dp_block_idx = block_idx - dp_start_block_idx; + int first_dp_tile = (params.block_mapping.cohort_raster) ? 0 : params.block_mapping.sk_tiles; + + // Blocks in first DP wave get configured number of tiles + tile_idx = first_dp_tile + dp_block_idx; + int tile_allottment = params.block_mapping.dp_first_wave_tiles; + + // Blocks in subsequent DP waves get 1 tile + if (dp_block_idx >= params.block_mapping.avail_sms) { + tile_allottment = 1; + tile_idx += (params.block_mapping.dp_first_wave_tiles - 1) * params.block_mapping.avail_sms; + } + + block_iters_remaining = params.block_mapping.iters_per_tile() * tile_allottment; + + init_dp_tile_work(tile_work, tile_idx); + + // DP blocks exit if out of bounds or overlap an SK tile (only possible during cohort rasterization, where dp_first_wave_tiles must be 1) + if ((tile_idx < params.block_mapping.sk_tiles) || + (tile_work.tiled_coord.m() >= params.block_mapping.tiled_shape().m()) || + (tile_work.tiled_coord.n() >= params.block_mapping.tiled_shape().n())) + { + return; + } + } + else if (sk_block) + { + // This is a SK block + int block_iter_end; + params.block_mapping.get_iter_extents(block_idx, block_iter_begin, block_iter_end); + block_iters_remaining = block_iter_end - block_iter_begin; + + tile_idx = params.block_mapping.get_sk_tile_idx(block_iter_end - 1); + init_sk_tile_work(tile_work, tile_idx, block_iter_begin, block_iter_begin + block_iters_remaining); + } + else + { + if (reduce_block) + { + // This is a reduction threadblock + int reduce_block_idx = block_idx - reduce_start_block_idx; + separate_reduction(reduce_block_idx); + } + + return; + } + + // Iteration-processing loop body + MCTLASS_PRAGMA_NO_UNROLL + while (true) + { + // Perform this block's share of work for this tile + process_tile( + tile_work, + block_idx, + dp_start_block_idx, + block_iter_begin); + + block_iters_remaining -= tile_work.k_iters_remaining; + + if (block_iters_remaining == 0) + { + break; + } + + // Continue to next tile + __syncthreads(); + + if (block_idx >= dp_start_block_idx) + { + // DP block consume their tiles at stride + tile_idx += params.block_mapping.avail_sms; + init_dp_tile_work(tile_work, tile_idx); + } + else + { + // SK blocks consume their tiles in backwards order + tile_idx--; + init_sk_tile_work(tile_work, tile_idx, block_iter_begin, block_iter_begin + block_iters_remaining); + } + } + + } + + +public: + + // + // Device-only API + // + + // Factory invocation + MCTLASS_DEVICE + static void invoke( + Params const ¶ms, + SharedStorage &shared_storage) + { + GemmStreamkWithFusedEpilogue op(params, shared_storage); + op(); + } + + + // Constructor + MCTLASS_DEVICE + GemmStreamkWithFusedEpilogue( + Params const ¶ms, + SharedStorage &shared_storage) + : + params(params), + shared_storage(shared_storage), + thread_idx(threadIdx.x), + warp_idx(__shfl_sync(0xffffffff, threadIdx.x / 64, 0)), // broadcast the warp_id computed by lane 0 to ensure dependent code + lane_idx(threadIdx.x % 64), + epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx) + {} + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()() { + // Generic SK code path + gemm(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_transpose_operands.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_transpose_operands.h new file mode 100644 index 0000000..19dabb8 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_transpose_operands.h @@ -0,0 +1,124 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and + batched array variants. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementA_, + typename LayoutA_, + ComplexTransform TransformA, + int AlignmentA, + typename ElementB_, + typename LayoutB_, + ComplexTransform TransformB, + int AlignmentB, + typename LayoutC_, + bool Transpose +> +struct MapArguments { + using ElementA = ElementA_; + using LayoutA = LayoutA_; + static ComplexTransform const kTransformA = TransformA; + static int const kAlignmentA = AlignmentA; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + static ComplexTransform const kTransformB = TransformB; + static int const kAlignmentB = AlignmentB; + using LayoutC = LayoutC_; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementA_, + typename LayoutA_, + ComplexTransform TransformA, + int AlignmentA, + typename ElementB_, + typename LayoutB_, + ComplexTransform TransformB, + int AlignmentB, + typename LayoutC_ +> +struct MapArguments< + ElementA_, + LayoutA_, + TransformA, + AlignmentA, + ElementB_, + LayoutB_, + TransformB, + AlignmentB, + LayoutC_, + true +> { + using ElementA = ElementB_; + using LayoutA = typename layout::LayoutTranspose::type; + static ComplexTransform const kTransformA = TransformB; + static int const kAlignmentA = AlignmentB; + using ElementB = ElementA_; + using LayoutB = typename layout::LayoutTranspose::type; + static ComplexTransform const kTransformB = TransformA; + static int const kAlignmentB = AlignmentA; + using LayoutC = typename layout::LayoutTranspose::type; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} +} +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_universal.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_universal.h new file mode 100644 index 0000000..58103f8 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_universal.h @@ -0,0 +1,698 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/arch/arch.h" +#include "mctlass/fast_math.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/semaphore.h" +#include "mctlass/gemm/kernel/gemm_universal.hpp" + +#include "mctlass/layout/matrix.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/kernel/params_universal_base.h" + +#include "mctlass/trace.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_ ///! Threadblock swizzling function +> +class GemmUniversal< + Mma_, + Epilogue_, + ThreadblockSwizzle_, + void, + // 3.x kernels use the first template argument to define the ProblemShape tuple + // We use this invariant to SFINAE dispatch against either the 2.x API or the 3.x API + cute::enable_if_t::value> +> { +public: + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + + static ComplexTransform const kTransformA = Mma::kTransformA; + static ComplexTransform const kTransformB = Mma::kTransformB; + using Operator = typename Mma::Operator; + + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::InstructionShape; + using ArchTag = typename Mma::ArchTag; + + static int const kStages = Mma::kStages; + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 64 * WarpCount::kCount; + + /// Split-K preserves splits that are 128b aligned + static int const kSplitKAlignment = const_max(128 / sizeof_bits::value, 128 / sizeof_bits::value); + + // + // Structures + // + + /// Argument structure + struct Arguments : UniversalArgumentsBase + { + // + // Data members + // + + typename EpilogueOutputOp::Params epilogue; + + void const * ptr_A; + void const * ptr_B; + void const * ptr_C; + void * ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + + typename LayoutA::Stride stride_a; + typename LayoutB::Stride stride_b; + typename LayoutC::Stride stride_c; + typename LayoutC::Stride stride_d; + + typename LayoutA::Stride::LongIndex lda; + typename LayoutB::Stride::LongIndex ldb; + typename LayoutC::Stride::LongIndex ldc; + typename LayoutC::Stride::LongIndex ldd; + + int const * ptr_gather_A_indices; + int const * ptr_gather_B_indices; + int const * ptr_scatter_D_indices; + + // + // Methods + // + + Arguments(): + ptr_A(nullptr), ptr_B(nullptr), ptr_C(nullptr), ptr_D(nullptr), + ptr_gather_A_indices(nullptr), + ptr_gather_B_indices(nullptr), + ptr_scatter_D_indices(nullptr) + {} + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void const * ptr_C, + void * ptr_D, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C, + int64_t batch_stride_D, + typename LayoutA::Stride stride_a, + typename LayoutB::Stride stride_b, + typename LayoutC::Stride stride_c, + typename LayoutC::Stride stride_d, + int const *ptr_gather_A_indices = nullptr, + int const *ptr_gather_B_indices = nullptr, + int const *ptr_scatter_D_indices = nullptr) + : + UniversalArgumentsBase(mode, problem_size, batch_count, batch_stride_D), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), + batch_stride_A(batch_stride_A), batch_stride_B(batch_stride_B), batch_stride_C(batch_stride_C), + stride_a(stride_a), stride_b(stride_b), stride_c(stride_c), stride_d(stride_d), + ptr_gather_A_indices(ptr_gather_A_indices), ptr_gather_B_indices(ptr_gather_B_indices), + ptr_scatter_D_indices(ptr_scatter_D_indices) + { + lda = 0; + ldb = 0; + ldc = 0; + ldd = 0; + MCTLASS_TRACE_HOST("GemmUniversal::Arguments::Arguments() - problem_size: " << problem_size); + } + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void const * ptr_C, + void * ptr_D, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C, + int64_t batch_stride_D, + typename LayoutA::Stride::LongIndex lda, + typename LayoutB::Stride::LongIndex ldb, + typename LayoutC::Stride::LongIndex ldc, + typename LayoutC::Stride::LongIndex ldd, + int const *ptr_gather_A_indices = nullptr, + int const *ptr_gather_B_indices = nullptr, + int const *ptr_scatter_D_indices = nullptr + ): + UniversalArgumentsBase(mode, problem_size, batch_count, batch_stride_D), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), + batch_stride_A(batch_stride_A), batch_stride_B(batch_stride_B), batch_stride_C(batch_stride_C), + lda(lda), ldb(ldb), ldc(ldc), ldd(ldd), + ptr_gather_A_indices(ptr_gather_A_indices), ptr_gather_B_indices(ptr_gather_B_indices), + ptr_scatter_D_indices(ptr_scatter_D_indices) + { + stride_a = make_Coord(lda); + stride_b = make_Coord(ldb); + stride_c = make_Coord(ldc); + stride_d = make_Coord(ldd); + MCTLASS_TRACE_HOST("GemmUniversal::Arguments::Arguments() - problem_size: " << problem_size); + } + + /// Returns arguments for the transposed problem + Arguments transposed_problem() const + { + Arguments args(*this); + + std::swap(args.problem_size.m(), args.problem_size.n()); + std::swap(args.ptr_A, args.ptr_B); + std::swap(args.lda, args.ldb); + std::swap(args.stride_a, args.stride_b); + std::swap(args.batch_stride_A, args.batch_stride_B); + std::swap(args.ptr_gather_A_indices, args.ptr_gather_B_indices); + + return args; + } + }; + + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params : UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC> + { + using ParamsBase = UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC>; + + // + // Data members + // + + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorB::Params params_B; + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::Params params_D; + + typename EpilogueOutputOp::Params output_op; + + void * ptr_A; + void * ptr_B; + void * ptr_C; + void * ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + + int * ptr_gather_A_indices; + int * ptr_gather_B_indices; + int * ptr_scatter_D_indices; + + // + // Host dispatch API + // + + /// Default constructor + Params() = default; + + /// Constructor + Params( + Arguments const &args, /// GEMM application arguments + int device_sms, /// Number of SMs on the device + int sm_occupancy) /// Kernel SM occupancy (in thread blocks) + : + ParamsBase(args, device_sms, sm_occupancy), + params_A(args.lda ? make_Coord_with_padding(args.lda) : args.stride_a), + params_B(args.ldb ? make_Coord_with_padding(args.ldb) : args.stride_b), + params_C(args.ldc ? make_Coord_with_padding(args.ldc) : args.stride_c), + params_D(args.ldd ? make_Coord_with_padding(args.ldd) : args.stride_d), + output_op(args.epilogue), + ptr_A(const_cast(args.ptr_A)), + ptr_B(const_cast(args.ptr_B)), + ptr_C(const_cast(args.ptr_C)), + ptr_D(args.ptr_D), + batch_stride_A(args.batch_stride_A), + batch_stride_B(args.batch_stride_B), + batch_stride_C(args.batch_stride_C), + ptr_gather_A_indices(const_cast(args.ptr_gather_A_indices)), + ptr_gather_B_indices(const_cast(args.ptr_gather_B_indices)), + ptr_scatter_D_indices(const_cast(args.ptr_scatter_D_indices)) + {} + + /// Lightweight update given a subset of arguments. + void update(Arguments const &args) + { + MCTLASS_TRACE_HOST("GemmUniversal::Params::update()"); + + // Update input/output pointers + ptr_A = const_cast(args.ptr_A); + ptr_B = const_cast(args.ptr_B); + ptr_C = const_cast(args.ptr_C); + ptr_D = args.ptr_D; + + batch_stride_A = args.batch_stride_A; + batch_stride_B = args.batch_stride_B; + batch_stride_C = args.batch_stride_C; + this->batch_stride_D = args.batch_stride_D; + + ptr_gather_A_indices = const_cast(args.ptr_gather_A_indices); + ptr_gather_B_indices = const_cast(args.ptr_gather_B_indices); + ptr_scatter_D_indices = const_cast(args.ptr_scatter_D_indices); + + output_op = args.epilogue; + } + }; + + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + +public: + + // + // Host dispatch API + // + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size) + { + MCTLASS_TRACE_HOST("GemmUniversal::can_implement()"); + + static int const kAlignmentA = (cute::is_same>::value) + ? 32 + : (cute::is_same>::value) + ? 64 + : Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = (cute::is_same>::value) + ? 32 + : (cute::is_same>::value) + ? 64 + : Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = (cute::is_same>::value) + ? 32 + : (cute::is_same>::value) + ? 64 + : Epilogue::OutputTileIterator::kElementsPerAccess; + + bool isAMisaligned = false; + bool isBMisaligned = false; + bool isCMisaligned = false; + + if (cute::is_same::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } else if (cute::is_same::value) { + isAMisaligned = problem_size.m() % kAlignmentA; + } else if (cute::is_same>::value + || cute::is_same>::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } + + if (cute::is_same::value) { + isBMisaligned = problem_size.n() % kAlignmentB; + } else if (cute::is_same::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } else if (cute::is_same>::value + || cute::is_same>::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } + + if (cute::is_same::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } else if (cute::is_same::value) { + isCMisaligned = problem_size.m() % kAlignmentC; + } else if (cute::is_same>::value + || cute::is_same>::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } + + if (isAMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for A operand"); + return Status::kErrorMisalignedOperand; + } + + if (isBMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for B operand"); + return Status::kErrorMisalignedOperand; + } + + if (isCMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for C operand"); + return Status::kErrorMisalignedOperand; + } + + MCTLASS_TRACE_HOST(" returning kSuccess"); + + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + + +public: + + // + // Device-only API + // + + // Factory invocation + MCTLASS_DEVICE + static void invoke( + Params const ¶ms, + SharedStorage &shared_storage) + { + GemmUniversal op; + op(params, shared_storage); + } + + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()( + Params const ¶ms, + SharedStorage &shared_storage) + { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + int offset_k = 0; + int problem_size_k = params.problem_size.k(); + + ElementA *ptr_A = static_cast(params.ptr_A); + ElementB *ptr_B = static_cast(params.ptr_B); + + // + // Fetch pointers based on mode. + // + if (params.mode == GemmUniversalMode::kGemm || + params.mode == GemmUniversalMode::kGemmSplitKParallel) { + + if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) { + + problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size; + } + + offset_k = threadblock_tile_offset.k() * params.gemm_k_size; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_A += threadblock_tile_offset.k() * params.batch_stride_A; + ptr_B += threadblock_tile_offset.k() * params.batch_stride_B; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_A = static_cast(params.ptr_A)[threadblock_tile_offset.k()]; + ptr_B = static_cast(params.ptr_B)[threadblock_tile_offset.k()]; + } + + __syncthreads(); + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + offset_k, + }; + + mctlass::MatrixCoord tb_offset_B{ + offset_k, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, + ptr_A, + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_A, + params.ptr_gather_A_indices); + + typename Mma::IteratorB iterator_B( + params.params_B, + ptr_B, + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B, + params.ptr_gather_B_indices); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma( + gemm_k_iterations, + accumulators, + iterator_A, + iterator_B, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + ElementC *ptr_C = static_cast(params.ptr_C); + ElementC *ptr_D = static_cast(params.ptr_D); + + // + // Fetch pointers based on mode. + // + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + if (params.mode == GemmUniversalMode::kGemm) { + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + } + else if (params.mode == GemmUniversalMode::kGemmSplitKParallel) { + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_C += threadblock_tile_offset.k() * params.batch_stride_C; + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_C = static_cast(params.ptr_C)[threadblock_tile_offset.k()]; + ptr_D = static_cast(params.ptr_D)[threadblock_tile_offset.k()]; + } + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + ptr_C, + params.problem_size.mn(), + thread_idx, + threadblock_offset, + params.ptr_scatter_D_indices + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.problem_size.mn(), + thread_idx, + threadblock_offset, + params.ptr_scatter_D_indices + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + } + + + // Execute the epilogue operator to update the destination tensor. + epilogue( + output_op, + iterator_D, + accumulators, + iterator_C); + + // + // Release the semaphore + // + + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_universal.hpp b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_universal.hpp new file mode 100644 index 0000000..8ce4d2f --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_universal.hpp @@ -0,0 +1,68 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::kernel { + +//////////////////////////////////////////////////////////////////////////////// + +/* + * Stateless universal device GEMM kernel type that treats GEMM as + * a composition of a collective mainloop and a collective epilogue. + * + * Supports both the 2.x and 3.x APIs based on whether the first type is + * a cute::tuple<> or not. + * 2.x API implementation: mctlass/gemm/kernel/gemm_universal.h + * 3.x API implementation: mctlass/gemm/kernel/gemm_*.hpp + * + * In the following declaration, the name preceding the 'Or' refers to + * 3.x API type argument order, and the name succeeding the 'Or' refers to + * 2.x API type argument order. Template arguments without two names + * belong to the 3.x API only. +**/ +template < + class ProblemShapeOrThreadblockMma_, // (m, n, k) or (m, n, k, l) + class CollectiveMainloopOrEpilogue_, + class CollectiveEpilogueOrThreadblockSwizzle_, + class GridSwizzle_ = void, + class Enable = void +> +class GemmUniversal; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::kernel + +//////////////////////////////////////////////////////////////////////////////// + +#include "mctlass/gemm/kernel/sm70_gemm.hpp" diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_universal_streamk.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_universal_streamk.h new file mode 100644 index 0000000..3c709d6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_universal_streamk.h @@ -0,0 +1,1175 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/barrier.h" +#include "mctlass/block_striped.h" + +#include "mctlass/trace.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_ ///! Threadblock mapping function +> +struct GemmUniversalStreamk { +public: + + + // + // Types and constants + // + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + + /// The per-thread tile of raw accumulators + using AccumulatorTile = typename Mma::FragmentC; + + static ComplexTransform const kTransformA = Mma::kTransformA; + static ComplexTransform const kTransformB = Mma::kTransformB; + using Operator = typename Mma::Operator; + + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::InstructionShape; + using ArchTag = typename Mma::ArchTag; + + static int const kStages = Mma::kStages; + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Workspace bytes per thread block + static size_t const kWorkspaceBytesPerBlock = + __NV_STD_MAX( + kThreadCount * sizeof(AccumulatorTile), + Epilogue::kWorkspaceBytesPerBlock); + + /// Block-striped reduction utility + using BlockStripedReduceT = BlockStripedReduce; + + + + // + // Structures + // + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmUniversalMode mode; + GemmCoord problem_size; + int batch_count; // Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor + + typename EpilogueOutputOp::Params epilogue; + + void const * ptr_A; + void const * ptr_B; + void const * ptr_C; + void * ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + int64_t batch_stride_D; + + typename LayoutA::Stride stride_a; + typename LayoutB::Stride stride_b; + typename LayoutC::Stride stride_c; + typename LayoutC::Stride stride_d; + + typename LayoutA::Stride::LongIndex lda; + typename LayoutB::Stride::LongIndex ldb; + typename LayoutC::Stride::LongIndex ldc; + typename LayoutC::Stride::LongIndex ldd; + + int avail_sms; /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling) + + + // + // Methods + // + + /// Default Constructor + Arguments(): + mode(GemmUniversalMode::kGemm), + batch_count(1), + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + avail_sms(-1) + {} + + /// Constructor + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_split, /// Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor (1 defaults to StreamK, >1 emulates Split-K) + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void const * ptr_C, + void * ptr_D, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C, + int64_t batch_stride_D, + typename LayoutA::Stride stride_a, + typename LayoutB::Stride stride_b, + typename LayoutC::Stride stride_c, + typename LayoutC::Stride stride_d, + int avail_sms = -1 /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling) + ): + mode(mode), + problem_size(problem_size), + batch_count(batch_split), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), + batch_stride_A(batch_stride_A), batch_stride_B(batch_stride_B), batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D), + stride_a(stride_a), stride_b(stride_b), stride_c(stride_c), stride_d(stride_d), avail_sms(avail_sms) + { + MCTLASS_TRACE_HOST("GemmUniversalStreamk::Arguments::Arguments() - problem_size: " << problem_size); + } + + /// Constructor + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_split, /// Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor (1 defaults to StreamK, >1 emulates Split-K) + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void const * ptr_C, + void * ptr_D, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C, + int64_t batch_stride_D, + typename LayoutA::Stride::LongIndex lda, + typename LayoutB::Stride::LongIndex ldb, + typename LayoutC::Stride::LongIndex ldc, + typename LayoutC::Stride::LongIndex ldd, + int avail_sms = -1 /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling) + ): + mode(mode), + problem_size(problem_size), + batch_count(batch_split), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), + batch_stride_A(batch_stride_A), batch_stride_B(batch_stride_B), batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D), + lda(lda), ldb(ldb), ldc(ldc), ldd(ldd), avail_sms(avail_sms) + { + stride_a = make_Coord(lda); + stride_b = make_Coord(ldb); + stride_c = make_Coord(ldc); + stride_d = make_Coord(ldd); + MCTLASS_TRACE_HOST("GemmUniversalStreamk::Arguments::Arguments() - problem_size: " << problem_size); + } + + /// Returns arguments for the transposed problem + Arguments transposed_problem() const + { + Arguments args(*this); + + std::swap(args.problem_size.m(), args.problem_size.n()); + std::swap(args.ptr_A, args.ptr_B); + std::swap(args.lda, args.ldb); + std::swap(args.stride_a, args.stride_b); + std::swap(args.batch_stride_A, args.batch_stride_B); + + return args; + } + }; + + + /// Parameters structure + struct Params + { + public: + + // + // Data members + // + + void * ptr_A; + void * ptr_B; + + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorB::Params params_B; + + int64_t batch_stride_A; + int64_t batch_stride_B; + + GemmUniversalMode mode; + + ThreadblockSwizzle block_mapping; + + void *barrier_workspace; + void *partials_workspace; + + typename EpilogueOutputOp::Params output_op; + + void * ptr_D; + void * ptr_C; + + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::OutputTileIterator::Params params_C; + + int64_t batch_stride_D; + int64_t batch_stride_C; + + + protected: + + // + // Host-only dispatch-utilities + // + + /// Pad the given allocation size up to the nearest cache line + static size_t cacheline_align_up(size_t size) + { + static const int CACHELINE_SIZE = 128; + return (size + CACHELINE_SIZE - 1) / CACHELINE_SIZE * CACHELINE_SIZE; + } + + /// Get the workspace size needed for barrier + size_t get_barrier_workspace_size() const + { + // For atomic reduction, each SK-block needs a synchronization flag. For parallel reduction, + // each reduction block needs its own synchronization flag. + int sk_blocks = block_mapping.sk_regions() * block_mapping.sk_blocks_per_region(); + int num_flags = fast_max(sk_blocks, block_mapping.reduction_blocks); + + return cacheline_align_up(sizeof(typename Barrier::T) * num_flags); + } + + /// Get the workspace size needed for intermediate partial sums + size_t get_partials_workspace_size() const + { + int sk_blocks = block_mapping.sk_regions() * block_mapping.sk_blocks_per_region(); + return cacheline_align_up(kWorkspaceBytesPerBlock * sk_blocks); + } + + + public: + + // + // Host dispatch API + // + + /// Default constructor + Params() = default; + + + /// Constructor + Params( + Arguments const &args, /// GEMM application arguments + int device_sms, /// Number of SMs on the device + int sm_occupancy) /// Kernel SM occupancy (in thread blocks) + : + params_A(args.lda ? make_Coord_with_padding(args.lda) : args.stride_a), + params_B(args.ldb ? make_Coord_with_padding(args.ldb) : args.stride_b), + params_C(args.ldc ? make_Coord_with_padding(args.ldc) : args.stride_c), + params_D(args.ldd ? make_Coord_with_padding(args.ldd) : args.stride_d), + output_op(args.epilogue), + mode(args.mode), + ptr_A(const_cast(args.ptr_A)), + ptr_B(const_cast(args.ptr_B)), + ptr_C(const_cast(args.ptr_C)), + ptr_D(args.ptr_D), + batch_stride_A(args.batch_stride_A), + batch_stride_B(args.batch_stride_B), + batch_stride_C(args.batch_stride_C), + batch_stride_D(args.batch_stride_D), + barrier_workspace(nullptr), + partials_workspace(nullptr) + { + // Number of SMs to make available for StreamK decomposition + int avail_sms = (args.avail_sms == -1) ? + device_sms : + fast_min(args.avail_sms, device_sms); + + // Initialize the block mapping structure + block_mapping = ThreadblockSwizzle( + typename ThreadblockSwizzle::template KernelTraits(), + args.mode, + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count, + sm_occupancy, + device_sms, + avail_sms); + } + + + /// Returns the workspace size (in bytes) needed for these parameters + size_t get_workspace_size() const + { + return + get_barrier_workspace_size() + + get_partials_workspace_size(); + } + + + /// Assign and initialize the specified workspace buffer. Assumes + /// the memory allocated to workspace is at least as large as get_workspace_size(). + Status init_workspace( + void *workspace, + mcStream_t stream = nullptr) + { + uint8_t *ptr = static_cast(workspace); + + // Establish partials workspace + partials_workspace = nullptr; + size_t partials_workspace_bytes = get_partials_workspace_size(); + if (partials_workspace_bytes > 0) + { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + partials_workspace = ptr; + ptr += partials_workspace_bytes; + } + + // Establish barrier workspace + barrier_workspace = nullptr; + size_t barrier_workspace_bytes = get_barrier_workspace_size(); + if (barrier_workspace_bytes > 0) + { + if (!workspace) { + return Status::kErrorWorkspaceNull; + } + barrier_workspace = ptr; + ptr += barrier_workspace_bytes; + } + + // Zero-initialize barrier workspace + if (barrier_workspace) + { + size_t barrier_workspace_bytes = get_barrier_workspace_size(); + + MCTLASS_TRACE_HOST(" Initialize " << barrier_workspace_bytes << " barrier bytes"); + + mcError_t result = mcMemsetAsync( + barrier_workspace, + 0, + barrier_workspace_bytes, + stream); + + if (result != mcSuccess) { + MCTLASS_TRACE_HOST(" mcMemsetAsync() returned error " << mcGetErrorString(result)); + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + + /// Returns the GEMM volume in thread block tiles + mctlass::gemm::GemmCoord get_tiled_shape() const + { + return block_mapping.tiled_shape(); + } + + + /// Returns the total number of thread blocks to launch + int get_grid_blocks() const + { + dim3 grid_dims = get_grid_dims(); + return grid_dims.x * grid_dims.y * grid_dims.z; + } + + + /// Returns the grid extents in thread blocks to launch + dim3 get_grid_dims() const + { + return block_mapping.get_grid_dims(); + } + + + /// Lightweight update given a subset of arguments. + void update(Arguments const &args) + { + MCTLASS_TRACE_HOST("GemmUniversalStreamK::Params::update()"); + + // Update input/output pointers + ptr_A = const_cast(args.ptr_A); + ptr_B = const_cast(args.ptr_B); + ptr_C = const_cast(args.ptr_C); + ptr_D = args.ptr_D; + + batch_stride_A = args.batch_stride_A; + batch_stride_B = args.batch_stride_B; + batch_stride_C = args.batch_stride_C; + batch_stride_D = args.batch_stride_D; + + output_op = args.epilogue; + } + + }; + + /// Tile work descriptor + struct TileWorkDesc + { + /// The linear tile index + int tile_idx; + + /// The location of this tile (in threadblock-tile coordinates) in the output matrix + mctlass::gemm::GemmCoord tiled_coord; + + // The first global-scoped MAC-iteration this threadblock will perform for this tile + int iter_begin; + + // The starting index in the k-domain for MAC-iterations this threadblock will perform for this tile + int k_begin; + + // The ending index (one-past) in the k-domain for MAC-iterations this threadblock will perform for this tile + int k_end; + + /// The number of remaining MAC-iterations this threadblock will perform for this tile + int k_iters_remaining; + + // Whether this block will perform the first iteration of this tile + MCTLASS_DEVICE + bool tile_started() + { + return (k_begin == 0); + } + + // Whether this block will perform the last iteration of this tile + MCTLASS_DEVICE + bool tile_finished(Params const ¶ms) + { + return (k_end == params.block_mapping.problem_size.k()); + } + }; + + + /// Shared memory storage structure + union SharedStorage + { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + +protected: + + // + // Data members + // + + /// GEMM problem parameters + Params params; + + /// Shared storage reference + SharedStorage &shared_storage; + + /// ID within the threadblock + int thread_idx; + + /// ID of warp + int warp_idx; + + /// ID of each thread within a warp + int lane_idx; + + /// Threadblock scoped epilogue + Epilogue epilogue; + + +public: + + // + // Host-only dispatch API + // + + /// Determines whether the GEMM problem size satisfies this kernel's + /// alignment requirements + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size) + { + MCTLASS_TRACE_HOST("GemmUniversalStreamk::can_implement()"); + + static int const kAlignmentA = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Epilogue::OutputTileIterator::kElementsPerAccess; + + bool isAMisaligned = false; + bool isBMisaligned = false; + bool isCMisaligned = false; + + if (platform::is_same::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } else if (platform::is_same::value) { + isAMisaligned = problem_size.m() % kAlignmentA; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } + + if (platform::is_same::value) { + isBMisaligned = problem_size.n() % kAlignmentB; + } else if (platform::is_same::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } + + if (platform::is_same::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } else if (platform::is_same::value) { + isCMisaligned = problem_size.m() % kAlignmentC; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } + + if (isAMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for A operand"); + return Status::kErrorMisalignedOperand; + } + + if (isBMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for B operand"); + return Status::kErrorMisalignedOperand; + } + + if (isCMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for C operand"); + return Status::kErrorMisalignedOperand; + } + + MCTLASS_TRACE_HOST(" returning kSuccess"); + + return Status::kSuccess; + } + + /// Determines whether the GEMM problem satisfies this kernel's + /// alignment requirements + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + +protected: + + // + // Device-only utility methods + // + + /// Iterator for fetching tile fragments from A + MCTLASS_DEVICE + typename Mma::IteratorA init_iterator_A( + TileWorkDesc &tile_work, + GemmUniversalMode mode) + { + // The input A matrix + ElementA *ptr_A = static_cast(params.ptr_A); + + // Update input pointers based on batched/array mode + if (mode == GemmUniversalMode::kBatched) { + ptr_A += tile_work.tiled_coord.k() * params.batch_stride_A; + } + if (mode == GemmUniversalMode::kArray) { + ptr_A = static_cast(params.ptr_A)[tile_work.tiled_coord.k()]; + } + + int m_begin = tile_work.tiled_coord.m() * Mma::Shape::kM; + int m_end = params.block_mapping.problem_size.m(); + + return typename Mma::IteratorA( + params.params_A, + ptr_A, + { m_end, tile_work.k_end }, + threadIdx.x, + { m_begin, tile_work.k_begin }); + + } + + + /// Iterator for fetching tile fragments from B + MCTLASS_DEVICE + typename Mma::IteratorB init_iterator_B( + TileWorkDesc &tile_work, + GemmUniversalMode mode) + { + // The input B matrix + ElementB *ptr_B = static_cast(params.ptr_B); + + // Update input pointers based on batched/array mode + if (mode == GemmUniversalMode::kBatched) { + ptr_B += tile_work.tiled_coord.k() * params.batch_stride_B; + } + if (mode == GemmUniversalMode::kArray) { + ptr_B = static_cast(params.ptr_B)[tile_work.tiled_coord.k()]; + } + + int n_begin = tile_work.tiled_coord.n() * Mma::Shape::kN; + int n_end = params.block_mapping.problem_size.n(); + return typename Mma::IteratorB( + params.params_B, + ptr_B, + { tile_work.k_end, n_end }, + threadIdx.x, + { tile_work.k_begin, n_begin }); + } + + + MCTLASS_DEVICE + void init_dp_tile_work( + TileWorkDesc &tile_work, + int tile_idx) + { + // The linear tile index + tile_work.tile_idx = tile_idx; + + // The first global-scoped MAC-iteration this threadblock will perform for this tile + tile_work.iter_begin = tile_idx * params.block_mapping.iters_per_tile(); + + // The number of MAC-iterations this threadblock will perform for this tile + tile_work.k_iters_remaining = params.block_mapping.iters_per_tile(); + + // The starting index in the k-domain for MAC-iterations this threadblock will perform for this tile + tile_work.k_begin = 0; + + // The ending index (one-past) in the k-domain for MAC-iterations this threadblock will perform for this tile + tile_work.k_end = params.block_mapping.problem_size.k(); + + // The location of this tile (in threadblock-tile coordinates) in the output matrix + tile_work.tiled_coord = params.block_mapping.get_tile_offset(tile_work.tile_idx); + } + + + MCTLASS_DEVICE + void init_sk_tile_work( + TileWorkDesc &tile_work, + int tile_idx, + int block_iter_begin, + int block_iter_end) + { + // The linear tile index + tile_work.tile_idx = tile_idx; + + // The first global-scoped MAC-iteration for this tile + int tile_iter_begin = tile_idx * params.block_mapping.iters_per_tile(); + + // The first global-scoped MAC-iteration this threadblock will perform for this tile + tile_work.iter_begin = max(block_iter_begin, tile_iter_begin); + + // The first tile-scoped MAC-iteration this threadblock will perform for this tile + int k_iter_begin = tile_work.iter_begin - tile_iter_begin; + + // The last (one past) tile-scoped MAC-iteration this threadblock will perform for this tile + int k_iter_end = block_iter_end - tile_iter_begin; + + // The number of MAC-iterations this threadblock will perform for this tile + tile_work.k_iters_remaining = k_iter_end - k_iter_begin; + + // The starting index in the k-domain for MAC-iterations this threadblock will perform for this tile + tile_work.k_begin = k_iter_begin * Mma::Shape::kK; + + // The ending index (one-past) in the k-domain for MAC-iterations this threadblock will perform for this tile + tile_work.k_end = min( + params.block_mapping.problem_size.k(), // extent of k domain + (k_iter_end * Mma::Shape::kK)); // extent of the threadblock's global iteration assignment + + // The location of this tile (in threadblock-tile coordinates) in the output matrix + tile_work.tiled_coord = params.block_mapping.get_tile_offset(tile_work.tile_idx); + } + + + /// Share accumulators with peers + MCTLASS_DEVICE + void share_accumulators( + AccumulatorTile const &accumulator_tile, + int block_idx, + int first_block_idx) + { + AccumulatorTile *accum_tile_workspace = reinterpret_cast(params.partials_workspace); + + int accum_tile_offset = first_block_idx * kThreadCount; + + if (block_idx == first_block_idx) + { + // First peer initializes the workspace partials + BlockStripedReduceT::store(accum_tile_workspace + accum_tile_offset, accumulator_tile, thread_idx); + } + else + { + // Subsequent peers atomically accumulate into the workspace partials + if (ThreadblockSwizzle::kReductionStrategy == ThreadblockSwizzle::kAtomic) + { + // Non-deterministic reduction order: wait for the first peer to have initialized the partials before we add to them + Barrier::wait_lt(params.barrier_workspace, thread_idx, first_block_idx, 1); + } + else + { + // Turnstile reduction order: wait until the previous peer has written + int wait_count = block_idx - first_block_idx; + Barrier::wait_eq(params.barrier_workspace, thread_idx, first_block_idx, wait_count); + } + + // Perform reduction in workspace + BlockStripedReduceT::reduce(accum_tile_workspace + accum_tile_offset, accumulator_tile, thread_idx); + } + + // Signal our arrival + Barrier::arrive_inc(params.barrier_workspace, thread_idx, first_block_idx); + } + + + /// Acquire accumulators from peers + MCTLASS_DEVICE + void acquire_accumulators( + AccumulatorTile &accumulator_tile, + int block_idx, + int first_block_idx) + { + AccumulatorTile *accum_tile_workspace = reinterpret_cast(params.partials_workspace); + + // Wait for arrival + int num_carry_in = block_idx - first_block_idx; + Barrier::wait_eq_reset(params.barrier_workspace, thread_idx, first_block_idx, num_carry_in); + + // Load and add peer-partials accumulator tile to local accumulator tile + int accum_tile_offset = first_block_idx * kThreadCount; + BlockStripedReduceT::load_add(accumulator_tile, accum_tile_workspace + accum_tile_offset, thread_idx); + } + + + /// Perform epilogue computations and output + MCTLASS_DEVICE + void do_epilogue( + TileWorkDesc &tile_work, + AccumulatorTile &accumulator_tile) + { + ElementC *ptr_C = static_cast(params.ptr_C); + ElementC *ptr_D = static_cast(params.ptr_D); + + // Update pointers for batched/array mode(s) + if (params.mode == GemmUniversalMode::kBatched) { + ptr_C += tile_work.tiled_coord.k() * params.batch_stride_C; + ptr_D += tile_work.tiled_coord.k() * params.batch_stride_D; + } + if (params.mode == GemmUniversalMode::kArray) { + ptr_C = static_cast(params.ptr_C)[tile_work.tiled_coord.k()]; + ptr_D = static_cast(params.ptr_D)[tile_work.tiled_coord.k()]; + } + + // Location of this tile in item-coords + MatrixCoord threadblock_item_begin( + tile_work.tiled_coord.m() * Mma::Shape::kM, + tile_work.tiled_coord.n() * Mma::Shape::kN + ); + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + ptr_C, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Execute the epilogue operator to update the destination tensor. + epilogue( + EpilogueOutputOp(params.output_op), + iterator_D, + accumulator_tile, + iterator_C); + } + + + MCTLASS_DEVICE + void separate_reduction(int reduce_idx) + { + int peer_idx_begin, peer_idx_last, reduce_tile_idx, reduce_fragment_idx; + + // Reduce by sk-tile (every tile contributed to by one or more blocks) + reduce_tile_idx = reduce_idx / Epilogue::kAccumulatorFragments; + reduce_fragment_idx = reduce_idx % Epilogue::kAccumulatorFragments; + + int iter_tile_first = reduce_tile_idx * params.block_mapping.iters_per_tile(); + int iter_tile_last = iter_tile_first + params.block_mapping.iters_per_tile() - 1; + + peer_idx_begin = params.block_mapping.get_sk_block_idx(iter_tile_first); + peer_idx_last = params.block_mapping.get_sk_block_idx(iter_tile_last); + + // Wait for peers to complete + int peer_idx_end = peer_idx_last + 1; + int num_peers = peer_idx_end - peer_idx_begin; + Barrier::wait_eq_reset( + params.barrier_workspace, + thread_idx, + (reduce_tile_idx * Epilogue::kAccumulatorFragments) + reduce_fragment_idx, + num_peers); + + /// The location of this tile (in threadblock-tile coordinates) in the output matrix + GemmCoord tiled_coord = params.block_mapping.get_tile_offset(reduce_tile_idx); + + // Location of this tile in item-coords + MatrixCoord threadblock_item_begin( + tiled_coord.m() * Mma::Shape::kM, + tiled_coord.n() * Mma::Shape::kN + ); + + ElementC *ptr_C = static_cast(params.ptr_C); + ElementC *ptr_D = static_cast(params.ptr_D); + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + ptr_C, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.block_mapping.problem_size.mn(), + thread_idx, + threadblock_item_begin); + + // Execute the epilogue operator to update the destination tensor. + epilogue.reduce( + peer_idx_begin, + peer_idx_end, + reduce_fragment_idx, + params.partials_workspace, + EpilogueOutputOp(params.output_op), + iterator_D, + iterator_C); + } + + + MCTLASS_DEVICE + void process_tile( + TileWorkDesc tile_work, + int block_idx, + int dp_start_block_idx, + int block_iter_begin) + { + // Initialize input iterators + typename Mma::IteratorA iterator_A = init_iterator_A(tile_work, params.mode); + typename Mma::IteratorB iterator_B = init_iterator_B(tile_work, params.mode); + + // Initialize accumulators + AccumulatorTile accumulator_tile; + accumulator_tile.clear(); + + // Initialize MMA abstraction + Mma mma( + shared_storage.main_loop, + thread_idx, + warp_idx, + lane_idx); + + // Perform this tile's range of multiply-accumulate (MAC) iterations + mma(tile_work.k_iters_remaining, accumulator_tile, iterator_A, iterator_B, accumulator_tile); + + if ((ThreadblockSwizzle::kReductionStrategy == ThreadblockSwizzle::kAtomic) || + (params.block_mapping.reduction_blocks == 0) || + (block_idx >= dp_start_block_idx)) + { + // + // Cooperative SK peer reduction or DP block + // + + int first_block_idx = params.block_mapping.get_first_block_idx(tile_work.tile_idx, block_idx); + + if (!tile_work.tile_finished(params)) { + // Non "finishing" SK blocks must share their partial accumulator sums through global scratch workspace + share_accumulators(accumulator_tile, block_idx, first_block_idx); + } + else + { + // DP blocks and "finishing" SK blocks must perform epilogue operations and write the output tile + if (!tile_work.tile_started()) + { + // A "finishing" SK block must first aggregate its accumulator partial sums with those shared by peer threadblocks + acquire_accumulators(accumulator_tile, block_idx, first_block_idx); + } + + do_epilogue(tile_work, accumulator_tile); + } + } + else + { + // + // Separate peer reduction + // + + // Share accumulator partial sums with peer threadblock(s) through scratch workspace + epilogue.share(block_idx, params.partials_workspace, accumulator_tile, tile_work.tile_started()); + + // Signal arrival + Barrier::arrive_range_inc( + params.barrier_workspace, + thread_idx, + tile_work.tile_idx * Epilogue::kAccumulatorFragments, + Epilogue::kAccumulatorFragments); + } + } + + + /// Executes one GEMM + MCTLASS_DEVICE + void gemm() + { + // Initialize block's iteration range + int tile_idx = 0; + int block_iter_begin = 0; + int block_iters_remaining = 0; + + int block_idx = params.block_mapping.get_block_idx(); + + int sk_padding_start_block_idx = params.block_mapping.sk_regions() * params.block_mapping.sk_blocks_per_region(); + int dp_start_block_idx = params.block_mapping.sk_waves * params.block_mapping.avail_sms; + int reduce_start_block_idx = dp_start_block_idx + params.block_mapping.dp_blocks; + int grid_padding_start_block_idx = reduce_start_block_idx + params.block_mapping.reduction_blocks; + + // Initialize tile work descriptor + TileWorkDesc tile_work; + + bool dp_block = (block_idx >= dp_start_block_idx) && (block_idx < reduce_start_block_idx); + bool sk_block = (block_idx < sk_padding_start_block_idx); + bool reduce_block = (block_idx >= reduce_start_block_idx) && + (block_idx < grid_padding_start_block_idx) && + (ThreadblockSwizzle::kReductionStrategy == ThreadblockSwizzle::kMixed); + + if (dp_block) + { + // This is a DP block + int dp_block_idx = block_idx - dp_start_block_idx; + int first_dp_tile = (params.block_mapping.cohort_raster) ? 0 : params.block_mapping.sk_tiles; + + // Blocks in first DP wave get configured number of tiles + tile_idx = first_dp_tile + dp_block_idx; + int tile_allottment = params.block_mapping.dp_first_wave_tiles; + + // Blocks in subsequent DP waves get 1 tile + if (dp_block_idx >= params.block_mapping.avail_sms) { + tile_allottment = 1; + tile_idx += (params.block_mapping.dp_first_wave_tiles - 1) * params.block_mapping.avail_sms; + } + + block_iters_remaining = params.block_mapping.iters_per_tile() * tile_allottment; + + init_dp_tile_work(tile_work, tile_idx); + + // DP blocks exit if out of bounds or overlap an SK tile (only possible during cohort rasterization, where dp_first_wave_tiles must be 1) + if ((tile_idx < params.block_mapping.sk_tiles) || + (tile_work.tiled_coord.m() >= params.block_mapping.tiled_shape().m()) || + (tile_work.tiled_coord.n() >= params.block_mapping.tiled_shape().n())) + { + return; + } + } + else if (sk_block) + { + // This is a SK block + int block_iter_end; + params.block_mapping.get_iter_extents(block_idx, block_iter_begin, block_iter_end); + block_iters_remaining = block_iter_end - block_iter_begin; + + tile_idx = params.block_mapping.get_sk_tile_idx(block_iter_end - 1); + init_sk_tile_work(tile_work, tile_idx, block_iter_begin, block_iter_begin + block_iters_remaining); + } + else + { + if (reduce_block) + { + // This is a reduction threadblock + int reduce_block_idx = block_idx - reduce_start_block_idx; + separate_reduction(reduce_block_idx); + } + + return; + } + + // Iteration-processing loop body + MCTLASS_PRAGMA_NO_UNROLL + while (true) + { + // Perform this block's share of work for this tile + process_tile( + tile_work, + block_idx, + dp_start_block_idx, + block_iter_begin); + + block_iters_remaining -= tile_work.k_iters_remaining; + + if (block_iters_remaining == 0) + { + break; + } + + // Continue to next tile + __syncthreads(); + + if (block_idx >= dp_start_block_idx) + { + // DP block consume their tiles at stride + tile_idx += params.block_mapping.avail_sms; + init_dp_tile_work(tile_work, tile_idx); + } + else + { + // SK blocks consume their tiles in backwards order + tile_idx--; + init_sk_tile_work(tile_work, tile_idx, block_iter_begin, block_iter_begin + block_iters_remaining); + } + } + + } + + +public: + + // + // Device-only API + // + + // Factory invocation + MCTLASS_DEVICE + static void invoke( + Params const ¶ms, + SharedStorage &shared_storage) + { + GemmUniversalStreamk op(params, shared_storage); + op(); + } + + + // Constructor + MCTLASS_DEVICE + GemmUniversalStreamk( + Params const ¶ms, + SharedStorage &shared_storage) + : + params(params), + shared_storage(shared_storage), + thread_idx(threadIdx.x), + warp_idx(__shfl_sync(0xffffffff, threadIdx.x / 64, 0)), // broadcast the warp_id computed by lane 0 to ensure dependent code + lane_idx(threadIdx.x % 64), + epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx) + {} + + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()() + { + // Generic SK code path + gemm(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_with_fused_epilogue.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_with_fused_epilogue.h new file mode 100644 index 0000000..b9f7470 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_with_fused_epilogue.h @@ -0,0 +1,1500 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Gemm kernel with fused reduction operation. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/layout/layout.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/semaphore.h" +#include "mctlass/gemm/kernel/params_universal_base.h" + +#include "mctlass/trace.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + bool IsSingleSource = Epilogue_::kIsSingleSource +> +struct GemmWithFusedEpilogue; + +// GemmWithFusedEpilogue with two sources +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_ ///! Threadblock swizzling function +> +struct GemmWithFusedEpilogue { +public: + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + + static ComplexTransform const kTransformA = Mma::kTransformA; + static ComplexTransform const kTransformB = Mma::kTransformB; + using Operator = typename Mma::Operator; + + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::InstructionShape; + using ArchTag = typename Mma::ArchTag; + + static int const kStages = Mma::kStages; + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 64 * WarpCount::kCount; + + /// Split-K preserves splits that are 128b aligned + static int const kSplitKAlignment = const_max( + 128 / sizeof_bits::value, + 128 / sizeof_bits::value + ); + + // + // Structures + // + + /// Argument structure + struct Arguments : UniversalArgumentsBase{ + + // + // Data members + // + + typename EpilogueOutputOp::Params epilogue; + + void const * ptr_A; + void const * ptr_B; + void const * ptr_C1; + void const * ptr_C2; + void * ptr_D; + + void * ptr_Vector; + void * ptr_Tensor; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C1; + int64_t batch_stride_C2; + int64_t batch_stride_Vector; + int64_t batch_stride_Tensor; + + typename LayoutA::Stride::Index lda; + typename LayoutB::Stride::Index ldb; + typename LayoutC::Stride::Index ldc1; + typename LayoutC::Stride::Index ldc2; + typename LayoutC::Stride::Index ldd; + typename LayoutC::Stride::Index ldr; + typename LayoutC::Stride::Index ldt; + + // + // Methods + // + + Arguments(): + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C1(nullptr), + ptr_C2(nullptr), + ptr_D(nullptr) + {} + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void const * ptr_C1, + void const * ptr_C2, + void * ptr_D, + void * ptr_Vector, + void * ptr_Tensor, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C1, + int64_t batch_stride_C2, + int64_t batch_stride_D, + int64_t batch_stride_Vector, + int64_t batch_stride_Tensor, + typename LayoutA::Stride::Index lda, + typename LayoutB::Stride::Index ldb, + typename LayoutC::Stride::Index ldc1, + typename LayoutC::Stride::Index ldc2, + typename LayoutC::Stride::Index ldd, + typename LayoutC::Stride::Index ldr, + typename LayoutC::Stride::Index ldt) + : + UniversalArgumentsBase(mode, problem_size, batch_count, batch_stride_D), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_C1(ptr_C1), ptr_C2(ptr_C2), ptr_D(ptr_D), + ptr_Vector(ptr_Vector), + ptr_Tensor(ptr_Tensor), + batch_stride_A(batch_stride_A), + batch_stride_B(batch_stride_B), + batch_stride_C1(batch_stride_C1), + batch_stride_C2(batch_stride_C2), + batch_stride_Vector(batch_stride_Vector), + batch_stride_Tensor(batch_stride_Tensor), + lda(lda), ldb(ldb), ldc1(ldc1), ldc2(ldc2), ldd(ldd), ldr(ldr), ldt(ldt) + { + MCTLASS_TRACE_HOST("GemmWithFusedEpilogue::Arguments::Arguments() - problem_size: " << problem_size); + MCTLASS_TRACE_HOST(" ptr_Vector: " << (void *)this->ptr_Vector); + MCTLASS_TRACE_HOST(" ptr_Tensor: " << (void *)this->ptr_Tensor); + MCTLASS_TRACE_HOST(" ldr: " << this->ldr); + MCTLASS_TRACE_HOST(" ldt: " << this->ldt); + } + + /// Returns arguments for the transposed problem + Arguments transposed_problem() const { + Arguments args(*this); + + std::swap(args.problem_size.m(), args.problem_size.n()); + std::swap(args.ptr_A, args.ptr_B); + std::swap(args.lda, args.ldb); + std::swap(args.batch_stride_A, args.batch_stride_B); + + return args; + } + }; + + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params : UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC> + { + using ParamsBase = UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC>; + + // + // Data members + // + + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorB::Params params_B; + typename Epilogue::OutputTileIterator::Params params_C1; + typename Epilogue::OutputTileIterator::Params params_C2; + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::TensorTileIterator::Params params_Tensor; + typename EpilogueOutputOp::Params output_op; + + void * ptr_A; + void * ptr_B; + void * ptr_C1; + void * ptr_C2; + void * ptr_D; + + void * ptr_Vector; + typename LayoutC::Stride::Index ldr; + + void * ptr_Tensor; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C1; + int64_t batch_stride_C2; + int64_t batch_stride_Vector; + int64_t batch_stride_Tensor; + + // + // Host dispatch API + // + + /// Default constructor + Params() = default; + + /// Constructor + Params( + Arguments const &args, /// GEMM application arguments + int device_sms, /// Number of SMs on the device + int sm_occupancy) /// Kernel SM occupancy (in thread blocks) + : + ParamsBase(args, device_sms, sm_occupancy), + params_A(args.lda), + params_B(args.ldb), + params_C1(args.ldc1), + params_C2(args.ldc2), + params_D(args.ldd), + params_Tensor(args.ldt), + output_op(args.epilogue), + ptr_A(const_cast(args.ptr_A)), + ptr_B(const_cast(args.ptr_B)), + ptr_C1(const_cast(args.ptr_C1)), + ptr_C2(const_cast(args.ptr_C2)), + ptr_D(args.ptr_D), + ptr_Vector(args.ptr_Vector), + ldr(args.ldr), + ptr_Tensor(args.ptr_Tensor), + batch_stride_A(args.batch_stride_A), + batch_stride_B(args.batch_stride_B), + batch_stride_C1(args.batch_stride_C1), + batch_stride_C2(args.batch_stride_C2), + batch_stride_Vector(args.batch_stride_Vector), + batch_stride_Tensor(args.batch_stride_Tensor) + { + MCTLASS_TRACE_HOST("GemmWithFusedEpilogue::Params::Params() - problem_size: " << problem_size); + MCTLASS_TRACE_HOST(" ptr_Vector: " << (void *)this->ptr_Vector); + MCTLASS_TRACE_HOST(" ptr_Tensor: " << (void *)this->ptr_Tensor); + MCTLASS_TRACE_HOST(" ldr: " << this->ldr); + MCTLASS_TRACE_HOST(" ldt: " << args.ldt); + } + + /// Lightweight update given a subset of arguments. + MCTLASS_HOST_DEVICE + void update(Arguments const &args) + { + ptr_A = const_cast(args.ptr_A); + ptr_B = const_cast(args.ptr_B); + ptr_C1 = const_cast(args.ptr_C1); + ptr_C2 = const_cast(args.ptr_C2); + ptr_D = args.ptr_D; + + ptr_Vector = args.ptr_Vector; + ldr = args.ldr; + ptr_Tensor = args.ptr_Tensor; + + batch_stride_A = args.batch_stride_A; + batch_stride_B = args.batch_stride_B; + batch_stride_C1 = args.batch_stride_C1; + batch_stride_C2 = args.batch_stride_C2; + batch_stride_Vector = args.batch_stride_Vector; + batch_stride_Tensor = args.batch_stride_Tensor; + this->batch_stride_D = args.batch_stride_D; + + output_op = args.epilogue; + + MCTLASS_TRACE_HOST("GemmWithFusedEpilogue::Params::update()"); + MCTLASS_TRACE_HOST(" ptr_Vector: " << (void *)this->ptr_Vector); + MCTLASS_TRACE_HOST(" ptr_Tensor: " << (void *)this->ptr_Tensor); + MCTLASS_TRACE_HOST(" ldr: " << this->ldr); + } + }; + + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + +public: + + // + // Host dispatch API + // + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size) { + + MCTLASS_TRACE_HOST("GemmWithFusedEpilogue::can_implement()"); + + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + bool isAMisaligned = false; + bool isBMisaligned = false; + bool isCMisaligned = false; + + if (platform::is_same::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } else if (platform::is_same::value) { + isAMisaligned = problem_size.m() % kAlignmentA; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } + + if (platform::is_same::value) { + isBMisaligned = problem_size.n() % kAlignmentB; + } else if (platform::is_same::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } + + if (platform::is_same::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } else if (platform::is_same::value) { + isCMisaligned = problem_size.m() % kAlignmentC; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } + + if (isAMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for A operand"); + return Status::kErrorMisalignedOperand; + } + + if (isBMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for B operand"); + return Status::kErrorMisalignedOperand; + } + + if (isCMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for C operand"); + return Status::kErrorMisalignedOperand; + } + + MCTLASS_TRACE_HOST(" returning kSuccess"); + + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + +public: + + // + // Device-only API + // + + // Factory invocation + MCTLASS_DEVICE + static void invoke( + Params const ¶ms, + SharedStorage &shared_storage) + { + GemmWithFusedEpilogue op; + op(params, shared_storage); + } + + #define SPLIT_K_ENABLED 1 + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + int offset_k = 0; + int problem_size_k = params.problem_size.k(); + + ElementA *ptr_A = static_cast(params.ptr_A); + ElementB *ptr_B = static_cast(params.ptr_B); + + + #if SPLIT_K_ENABLED + // + // Fetch pointers based on mode. + // + if (params.mode == GemmUniversalMode::kGemm || + params.mode == GemmUniversalMode::kGemmSplitKParallel) { + + if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) { + + problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size; + } + + offset_k = threadblock_tile_offset.k() * params.gemm_k_size; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_A += threadblock_tile_offset.k() * params.batch_stride_A; + ptr_B += threadblock_tile_offset.k() * params.batch_stride_B; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_A = static_cast(params.ptr_A)[threadblock_tile_offset.k()]; + ptr_B = static_cast(params.ptr_B)[threadblock_tile_offset.k()]; + } + #endif + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + offset_k, + }; + + mctlass::MatrixCoord tb_offset_B{ + offset_k, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, + ptr_A, + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + params.params_B, + ptr_B, + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma( + gemm_k_iterations, + accumulators, + iterator_A, + iterator_B, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + ElementC *ptr_C1 = static_cast(params.ptr_C1); + ElementC *ptr_C2 = static_cast(params.ptr_C2); + ElementC *ptr_D = static_cast(params.ptr_D); + typename Epilogue::ElementTensor *ptr_Tensor = static_cast(params.ptr_Tensor); + + // Define the reduction output pointer and move to the appropriate place + typename Epilogue::ElementVector *ptr_Vector = + static_cast(params.ptr_Vector); + + // + // Fetch pointers based on mode. + // + + // + // Special path when split-K not enabled. + // + + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() == 1) { + + // Tile iterators loading from source tensors. + typename Epilogue::OutputTileIterator iterator_C1( + params.params_C1, + ptr_C1, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + typename Epilogue::OutputTileIterator iterator_C2( + params.params_C2, + ptr_C2, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Additional tensor to load from + typename Epilogue::TensorTileIterator tensor_iterator( + params.params_Tensor, + // Only the final block outputs Tensor + ptr_Tensor, + params.problem_size.mn(), + thread_idx, + threadblock_offset); + + // Construct the epilogue + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Move to appropriate location for this output tile + if (ptr_Vector) { + ptr_Vector += threadblock_offset.column() + threadblock_tile_offset.m() * params.ldr; + } + + // Execute the epilogue operator to update the destination tensor. + epilogue(output_op, + ptr_Vector, + iterator_D, + accumulators, + iterator_C1, + iterator_C2, + tensor_iterator, + params.problem_size.mn(), + threadblock_offset); + + return; + } + + // + // Slower path when split-K or batching is needed + // + + + #if SPLIT_K_ENABLED + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + if (params.mode == GemmUniversalMode::kGemm) { + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + } + else if (params.mode == GemmUniversalMode::kGemmSplitKParallel) { + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_C1 += threadblock_tile_offset.k() * params.batch_stride_C1; + if (ptr_C2) { + ptr_C2 += threadblock_tile_offset.k() * params.batch_stride_C2; + } + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + if (ptr_Tensor) { + ptr_Tensor += threadblock_tile_offset.k() * params.batch_stride_Tensor; + } + if (ptr_Vector) { + ptr_Vector += threadblock_tile_offset.k() * params.batch_stride_Vector; + } + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_C1 = static_cast(params.ptr_C1)[threadblock_tile_offset.k()]; + if (ptr_C2) { + ptr_C2 = static_cast(params.ptr_C2)[threadblock_tile_offset.k()]; + } + ptr_D = static_cast(params.ptr_D)[threadblock_tile_offset.k()]; + if (ptr_Tensor) { + ptr_Tensor = static_cast(params.ptr_Tensor)[threadblock_tile_offset.k()]; + } + if (ptr_Vector) { + ptr_Vector = static_cast(params.ptr_Vector)[threadblock_tile_offset.k()]; + } + } + #endif + + // Tile iterators loading from source tensors. + typename Epilogue::OutputTileIterator iterator_C1( + params.params_C1, + ptr_C1, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + typename Epilogue::OutputTileIterator iterator_C2( + params.params_C2, + ptr_C2, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Additional tensor to load from + typename Epilogue::TensorTileIterator tensor_iterator( + params.params_Tensor, + // Only the final block outputs Tensor + ((params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) && + (params.grid_tiled_shape.k() != threadblock_tile_offset.k() + 1)) + ? nullptr + : ptr_Tensor, + params.problem_size.mn(), + thread_idx, + threadblock_offset); + + // Construct the epilogue + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + #if SPLIT_K_ENABLED + // Wait on the semaphore - this latency may have been covered by iterator construction + if ((params.mode == GemmUniversalMode::kGemm) && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C1 = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + + } + #endif + + // Move to appropriate location for this output tile + if (ptr_Vector) { + ptr_Vector += threadblock_offset.column() + threadblock_tile_offset.m() * params.ldr; + } + + // Execute the epilogue operator to update the destination tensor. + epilogue(output_op, + // Only the final block uses Vector + ((params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) && + (params.grid_tiled_shape.k() != threadblock_tile_offset.k() + 1)) + ? nullptr + : ptr_Vector, + iterator_D, + accumulators, + iterator_C1, + iterator_C2, + tensor_iterator, + params.problem_size.mn(), + threadblock_offset); + + // + // Release the semaphore + // + + #if SPLIT_K_ENABLED + if ((params.mode == GemmUniversalMode::kGemm) && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + #endif + } +}; + +// GemmWithFusedEpilogue with one source +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_ ///! Threadblock swizzling function +> +struct GemmWithFusedEpilogue { +public: + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + + static ComplexTransform const kTransformA = Mma::kTransformA; + static ComplexTransform const kTransformB = Mma::kTransformB; + using Operator = typename Mma::Operator; + + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::InstructionShape; + using ArchTag = typename Mma::ArchTag; + + static int const kStages = Mma::kStages; + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Split-K preserves splits that are 128b aligned + static int const kSplitKAlignment = const_max( + 128 / sizeof_bits::value, + 128 / sizeof_bits::value + ); + + // + // Structures + // + + /// Argument structure + struct Arguments : UniversalArgumentsBase + { + // + // Data members + // + + typename EpilogueOutputOp::Params epilogue; + + void const * ptr_A; + void const * ptr_B; + void const * ptr_C; + void * ptr_D; + + void * ptr_Vector; + void * ptr_Tensor; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + int64_t batch_stride_Vector; + int64_t batch_stride_Tensor; + + typename LayoutA::Stride::Index lda; + typename LayoutB::Stride::Index ldb; + typename LayoutC::Stride::Index ldc; + typename LayoutC::Stride::Index ldd; + typename LayoutC::Stride::Index ldr; + typename LayoutC::Stride::Index ldt; + + // + // Methods + // + + Arguments(): + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr) + {} + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void const * ptr_C, + void * ptr_D, + void * ptr_Vector, + void * ptr_Tensor, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C, + int64_t batch_stride_D, + int64_t batch_stride_Vector, + int64_t batch_stride_Tensor, + typename LayoutA::Stride::Index lda, + typename LayoutB::Stride::Index ldb, + typename LayoutC::Stride::Index ldc, + typename LayoutC::Stride::Index ldd, + typename LayoutC::Stride::Index ldr, + typename LayoutC::Stride::Index ldt) + : + UniversalArgumentsBase(mode, problem_size, batch_count, batch_stride_D), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), + ptr_Vector(ptr_Vector), + ptr_Tensor(ptr_Tensor), + batch_stride_A(batch_stride_A), + batch_stride_B(batch_stride_B), + batch_stride_C(batch_stride_C), + batch_stride_Vector(batch_stride_Vector), + batch_stride_Tensor(batch_stride_Tensor), + lda(lda), ldb(ldb), ldc(ldc), ldd(ldd), ldr(ldr), ldt(ldt) + { + MCTLASS_TRACE_HOST("GemmWithFusedEpilogue::Arguments::Arguments() - problem_size: " << problem_size); + MCTLASS_TRACE_HOST(" ptr_Vector: " << (void *)this->ptr_Vector); + MCTLASS_TRACE_HOST(" ptr_Tensor: " << (void *)this->ptr_Tensor); + MCTLASS_TRACE_HOST(" ldr: " << this->ldr); + MCTLASS_TRACE_HOST(" ldt: " << this->ldt); + } + + /// Returns arguments for the transposed problem + Arguments transposed_problem() const { + Arguments args(*this); + + std::swap(args.problem_size.m(), args.problem_size.n()); + std::swap(args.ptr_A, args.ptr_B); + std::swap(args.lda, args.ldb); + std::swap(args.batch_stride_A, args.batch_stride_B); + + return args; + } + }; + + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params : UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC> + { + using ParamsBase = UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC>; + + // + // Data members + // + + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorB::Params params_B; + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::TensorTileIterator::Params params_Tensor; + + typename EpilogueOutputOp::Params output_op; + + void * ptr_A; + void * ptr_B; + void * ptr_C; + void * ptr_D; + + void * ptr_Vector; + typename LayoutC::Stride::Index ldr; + + void * ptr_Tensor; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + int64_t batch_stride_Vector; + int64_t batch_stride_Tensor; + + // + // Host dispatch API + // + + /// Default constructor + Params() = default; + + /// Constructor + Params( + Arguments const &args, /// GEMM application arguments + int device_sms, /// Number of SMs on the device + int sm_occupancy) /// Kernel SM occupancy (in thread blocks) + : + ParamsBase(args, device_sms, sm_occupancy), + params_A(args.lda), + params_B(args.ldb), + params_C(args.ldc), + params_D(args.ldd), + params_Tensor(args.ldt), + output_op(args.epilogue), + ptr_A(const_cast(args.ptr_A)), + ptr_B(const_cast(args.ptr_B)), + ptr_C(const_cast(args.ptr_C)), + ptr_D(args.ptr_D), + ptr_Vector(args.ptr_Vector), + ldr(args.ldr), + ptr_Tensor(args.ptr_Tensor), + batch_stride_A(args.batch_stride_A), + batch_stride_B(args.batch_stride_B), + batch_stride_C(args.batch_stride_C), + batch_stride_Vector(args.batch_stride_Vector), + batch_stride_Tensor(args.batch_stride_Tensor) + { + MCTLASS_TRACE_HOST("GemmWithFusedEpilogue::Params::Params() - problem_size: " << problem_size); + MCTLASS_TRACE_HOST(" ptr_Vector: " << (void *)this->ptr_Vector); + MCTLASS_TRACE_HOST(" ptr_Tensor: " << (void *)this->ptr_Tensor); + MCTLASS_TRACE_HOST(" ldr: " << this->ldr); + MCTLASS_TRACE_HOST(" ldt: " << args.ldt); + } + + /// Lightweight update given a subset of arguments. + MCTLASS_HOST_DEVICE + void update(Arguments const &args) + { + ptr_A = const_cast(args.ptr_A); + ptr_B = const_cast(args.ptr_B); + ptr_C = const_cast(args.ptr_C); + ptr_D = args.ptr_D; + + ptr_Vector = args.ptr_Vector; + ldr = args.ldr; + ptr_Tensor = args.ptr_Tensor; + + batch_stride_A = args.batch_stride_A; + batch_stride_B = args.batch_stride_B; + batch_stride_C = args.batch_stride_C; + batch_stride_Vector = args.batch_stride_Vector; + batch_stride_Tensor = args.batch_stride_Tensor; + this->batch_stride_D = args.batch_stride_D; + + output_op = args.epilogue; + + MCTLASS_TRACE_HOST("GemmWithFusedEpilogue::Params::update()"); + MCTLASS_TRACE_HOST(" ptr_Vector: " << (void *)this->ptr_Vector); + MCTLASS_TRACE_HOST(" ptr_Tensor: " << (void *)this->ptr_Tensor); + MCTLASS_TRACE_HOST(" ldr: " << this->ldr); + } + }; + + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + +public: + + // + // Host dispatch API + // + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size) { + + MCTLASS_TRACE_HOST("GemmWithFusedEpilogue::can_implement()"); + + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + bool isAMisaligned = false; + bool isBMisaligned = false; + bool isCMisaligned = false; + + if (platform::is_same::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } else if (platform::is_same::value) { + isAMisaligned = problem_size.m() % kAlignmentA; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } + + if (platform::is_same::value) { + isBMisaligned = problem_size.n() % kAlignmentB; + } else if (platform::is_same::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } + + if (platform::is_same::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } else if (platform::is_same::value) { + isCMisaligned = problem_size.m() % kAlignmentC; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } + + if (isAMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for A operand"); + return Status::kErrorMisalignedOperand; + } + + if (isBMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for B operand"); + return Status::kErrorMisalignedOperand; + } + + if (isCMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for C operand"); + return Status::kErrorMisalignedOperand; + } + + MCTLASS_TRACE_HOST(" returning kSuccess"); + + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + +public: + + // + // Device-only API + // + + // Factory invocation + MCTLASS_DEVICE + static void invoke( + Params const ¶ms, + SharedStorage &shared_storage) + { + GemmWithFusedEpilogue op; + op(params, shared_storage); + } + + #define SPLIT_K_ENABLED 1 + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + int offset_k = 0; + int problem_size_k = params.problem_size.k(); + + ElementA *ptr_A = static_cast(params.ptr_A); + ElementB *ptr_B = static_cast(params.ptr_B); + + + #if SPLIT_K_ENABLED + // + // Fetch pointers based on mode. + // + if (params.mode == GemmUniversalMode::kGemm || + params.mode == GemmUniversalMode::kGemmSplitKParallel) { + + if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) { + + problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size; + } + + offset_k = threadblock_tile_offset.k() * params.gemm_k_size; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_A += threadblock_tile_offset.k() * params.batch_stride_A; + ptr_B += threadblock_tile_offset.k() * params.batch_stride_B; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_A = static_cast(params.ptr_A)[threadblock_tile_offset.k()]; + ptr_B = static_cast(params.ptr_B)[threadblock_tile_offset.k()]; + } + #endif + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + offset_k, + }; + + mctlass::MatrixCoord tb_offset_B{ + offset_k, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, + ptr_A, + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + params.params_B, + ptr_B, + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma( + gemm_k_iterations, + accumulators, + iterator_A, + iterator_B, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + ElementC *ptr_C = static_cast(params.ptr_C); + ElementC *ptr_D = static_cast(params.ptr_D); + typename Epilogue::ElementTensor *ptr_Tensor = static_cast(params.ptr_Tensor); + + // Define the reduction output pointer and move to the appropriate place + typename Epilogue::ElementVector *ptr_Vector = + static_cast(params.ptr_Vector); + + // + // Fetch pointers based on mode. + // + + // + // Special path when split-K not enabled. + // + + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() == 1) { + + // Tile iterators loading from source tensors. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + ptr_C, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Additional tensor to load from + typename Epilogue::TensorTileIterator tensor_iterator( + params.params_Tensor, + // Only the final block outputs Tensor + ptr_Tensor, + params.problem_size.mn(), + thread_idx, + threadblock_offset); + + // Construct the epilogue + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Move to appropriate location for this output tile + if (ptr_Vector) { + ptr_Vector += threadblock_offset.column() + threadblock_tile_offset.m() * params.ldr; + } + + // Execute the epilogue operator to update the destination tensor. + epilogue(output_op, + ptr_Vector, + iterator_D, + accumulators, + iterator_C, + tensor_iterator, + params.problem_size.mn(), + threadblock_offset); + + return; + } + + // + // Slower path when split-K or batching is needed + // + + + #if SPLIT_K_ENABLED + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + if (params.mode == GemmUniversalMode::kGemm) { + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + } + else if (params.mode == GemmUniversalMode::kGemmSplitKParallel) { + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_C += threadblock_tile_offset.k() * params.batch_stride_C; + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + if (ptr_Tensor) { + ptr_Tensor += threadblock_tile_offset.k() * params.batch_stride_Tensor; + } + if (ptr_Vector) { + ptr_Vector += threadblock_tile_offset.k() * params.batch_stride_Vector; + } + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_C = static_cast(params.ptr_C)[threadblock_tile_offset.k()]; + ptr_D = static_cast(params.ptr_D)[threadblock_tile_offset.k()]; + if (ptr_Tensor) { + ptr_Tensor = static_cast(params.ptr_Tensor)[threadblock_tile_offset.k()]; + } + if (ptr_Vector) { + ptr_Vector = static_cast(params.ptr_Vector)[threadblock_tile_offset.k()]; + } + } + #endif + + // Tile iterators loading from source tensors. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + ptr_C, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Additional tensor to load from + typename Epilogue::TensorTileIterator tensor_iterator( + params.params_Tensor, + // Only the final block outputs Tensor + ((params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) && + (params.grid_tiled_shape.k() != threadblock_tile_offset.k() + 1)) + ? nullptr + : ptr_Tensor, + params.problem_size.mn(), + thread_idx, + threadblock_offset); + + // Construct the epilogue + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + #if SPLIT_K_ENABLED + // Wait on the semaphore - this latency may have been covered by iterator construction + if ((params.mode == GemmUniversalMode::kGemm) && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + + } + #endif + + // Move to appropriate location for this output tile + if (ptr_Vector) { + ptr_Vector += threadblock_offset.column() + threadblock_tile_offset.m() * params.ldr; + } + + // Execute the epilogue operator to update the destination tensor. + epilogue(output_op, + // Only the final block uses Vector + ((params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) && + (params.grid_tiled_shape.k() != threadblock_tile_offset.k() + 1)) + ? nullptr + : ptr_Vector, + iterator_D, + accumulators, + iterator_C, + tensor_iterator, + params.problem_size.mn(), + threadblock_offset); + + // + // Release the semaphore + // + + #if SPLIT_K_ENABLED + if ((params.mode == GemmUniversalMode::kGemm) && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + #endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemm_with_k_reduction.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_with_k_reduction.h new file mode 100644 index 0000000..3f22166 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemm_with_k_reduction.h @@ -0,0 +1,700 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/semaphore.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/gemm/kernel/params_universal_base.h" + +#include "mctlass/trace.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename EpilogueGemmKReduction_, ///! Epilogue + typename ThreadblockSwizzle_ ///! Threadblock swizzling function +> +struct GemmWithKReduction { +public: + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using EpilogueGemmKReduction = EpilogueGemmKReduction_; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + using LayoutGemmKReduction = mctlass::layout::PitchLinear; + + static ComplexTransform const kTransformA = Mma::kTransformA; + static ComplexTransform const kTransformB = Mma::kTransformB; + using Operator = typename Mma::Operator; + + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::InstructionShape; + using ArchTag = typename Mma::ArchTag; + + static int const kStages = Mma::kStages; + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Split-K preserves splits that are 128b aligned + static int const kSplitKAlignment = const_max(128 / sizeof_bits::value, 128 / sizeof_bits::value); + + static int const kReduceKForA = Mma::kReduceKForA; + + // + // Structures + // + + /// Argument structure + struct Arguments : UniversalArgumentsBase + { + // + // Data members + // + + typename EpilogueOutputOp::Params epilogue; + + void const * ptr_A; + void const * ptr_B; + void const * ptr_C; + void * ptr_D; + void * ptr_gemm_k_reduction; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + int64_t batch_stride_gemm_k_reduction; + + typename LayoutA::Stride::Index lda; + typename LayoutB::Stride::Index ldb; + typename LayoutC::Stride::Index ldc; + typename LayoutC::Stride::Index ldd; + typename LayoutGemmKReduction::Stride::Index ld_gemm_k_reduction; + + // + // Methods + // + + Arguments() : + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + ptr_gemm_k_reduction(nullptr) + {} + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void const * ptr_C, + void * ptr_D, + void * ptr_gemm_k_reduction, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C, + int64_t batch_stride_D, + int64_t batch_stride_gemm_k_reduction, + typename LayoutA::Stride::Index lda, + typename LayoutB::Stride::Index ldb, + typename LayoutC::Stride::Index ldc, + typename LayoutC::Stride::Index ldd, + typename LayoutGemmKReduction::Stride::Index ld_gemm_k_reduction) + : + UniversalArgumentsBase(mode, problem_size, batch_count, batch_stride_D), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), ptr_gemm_k_reduction(ptr_gemm_k_reduction), + batch_stride_A(batch_stride_A), batch_stride_B(batch_stride_B), batch_stride_C(batch_stride_C), batch_stride_gemm_k_reduction(batch_stride_gemm_k_reduction), + lda(lda), ldb(ldb), ldc(ldc), ldd(ldd), ld_gemm_k_reduction(ld_gemm_k_reduction) + { + MCTLASS_TRACE_HOST("GemmUniversal::Arguments::Arguments() - problem_size: " << problem_size); + } + + /// Returns arguments for the transposed problem + Arguments transposed_problem() const { + Arguments args(*this); + + std::swap(args.problem_size.m(), args.problem_size.n()); + std::swap(args.ptr_A, args.ptr_B); + std::swap(args.lda, args.ldb); + std::swap(args.batch_stride_A, args.batch_stride_B); + + return args; + } + }; + + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params : UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC> + { + using ParamsBase = UniversalParamsBase< + ThreadblockSwizzle, + ThreadblockShape, + ElementA, + ElementB, + ElementC>; + + // + // Data members + // + + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorB::Params params_B; + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::Params params_D; + + typename EpilogueOutputOp::Params output_op; + + void * ptr_A; + void * ptr_B; + void * ptr_C; + void * ptr_D; + void * ptr_gemm_k_reduction; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + int64_t batch_stride_gemm_k_reduction; + + // + // Host dispatch API + // + + /// Default constructor + Params() = default; + + /// Constructor + Params( + Arguments const &args, /// GEMM application arguments + int device_sms, /// Number of SMs on the device + int sm_occupancy) /// Kernel SM occupancy (in thread blocks) + : + ParamsBase(args, device_sms, sm_occupancy), + params_A(args.lda), + params_B(args.ldb), + params_C(args.ldc), + params_D(args.ldd), + output_op(args.epilogue), + ptr_A(const_cast(args.ptr_A)), + ptr_B(const_cast(args.ptr_B)), + ptr_C(const_cast(args.ptr_C)), + batch_stride_A(args.batch_stride_A), + batch_stride_B(args.batch_stride_B), + batch_stride_C(args.batch_stride_C), + batch_stride_gemm_k_reduction(args.batch_stride_gemm_k_reduction), + ptr_D(args.ptr_D), + ptr_gemm_k_reduction(args.ptr_gemm_k_reduction) + {} + + /// Assign and initialize the specified workspace buffer. Assumes + /// the memory allocated to workspace is at least as large as get_workspace_size(). + Status init_workspace( + void *workspace, + mcStream_t stream = nullptr) + { + MCTLASS_TRACE_HOST("GemmUniversal::Params::Params() - problem_size: " << this->problem_size); + + if (this->mode == GemmUniversalMode::kGemmSplitKParallel) { + ptr_D = workspace; + ptr_gemm_k_reduction = static_cast(workspace) + + sizeof(ElementC) * size_t(this->batch_stride_D) * size_t(this->grid_tiled_shape.k()); + + return Status::kSuccess; + } + + return ParamsBase::init_workspace(workspace, stream); + } + + /// Returns the workspace size (in bytes) needed for this problem geometry + size_t get_workspace_size() const + { + size_t workspace_bytes = ParamsBase::get_workspace_size(); + + if (this->mode == GemmUniversalMode::kGemmSplitKParallel) + { + // Split-K parallel always requires a temporary workspace + workspace_bytes += + sizeof(ElementC) * + size_t(batch_stride_gemm_k_reduction) * + size_t(this->grid_tiled_shape.k()); + } + + return workspace_bytes; + } + + /// Lightweight update given a subset of arguments. + void update(Arguments const &args) + { + ptr_A = const_cast(args.ptr_A); + ptr_B = const_cast(args.ptr_B); + ptr_C = const_cast(args.ptr_C); + ptr_D = args.ptr_D; + ptr_gemm_k_reduction = args.ptr_gemm_k_reduction; + + batch_stride_A = args.batch_stride_A; + batch_stride_B = args.batch_stride_B; + batch_stride_C = args.batch_stride_C; + batch_stride_gemm_k_reduction = args.batch_stride_gemm_k_reduction; + this->batch_stride_D = args.batch_stride_D; + + output_op = args.epilogue; + + MCTLASS_TRACE_HOST("GemmUniversal::Params::update()"); + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + +public: + + // + // Host dispatch API + // + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size) { + + MCTLASS_TRACE_HOST("GemmUniversal::can_implement()"); + + static int const kAlignmentA = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = (platform::is_same>::value) + ? 32 + : (platform::is_same>::value) + ? 64 + : Epilogue::OutputTileIterator::kElementsPerAccess; + + bool isAMisaligned = false; + bool isBMisaligned = false; + bool isCMisaligned = false; + + if (platform::is_same::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } else if (platform::is_same::value) { + isAMisaligned = problem_size.m() % kAlignmentA; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isAMisaligned = problem_size.k() % kAlignmentA; + } + + if (platform::is_same::value) { + isBMisaligned = problem_size.n() % kAlignmentB; + } else if (platform::is_same::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isBMisaligned = problem_size.k() % kAlignmentB; + } + + if (platform::is_same::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } else if (platform::is_same::value) { + isCMisaligned = problem_size.m() % kAlignmentC; + } else if (platform::is_same>::value + || platform::is_same>::value) { + isCMisaligned = problem_size.n() % kAlignmentC; + } + + if (isAMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for operand A"); + return Status::kErrorMisalignedOperand; + } + + if (isBMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for operand B"); + return Status::kErrorMisalignedOperand; + } + + if (isCMisaligned) { + MCTLASS_TRACE_HOST(" returning kErrorMisalignedOperand for operand C"); + return Status::kErrorMisalignedOperand; + } + + MCTLASS_TRACE_HOST(" returning kSuccess"); + + return Status::kSuccess; + } + + + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + + +public: + + // + // Device-only API + // + + // Factory invocation + MCTLASS_DEVICE + static void invoke( + Params const ¶ms, + SharedStorage &shared_storage) + { + GemmWithKReduction op; + op(params, shared_storage); + } + + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + int offset_k = 0; + int problem_size_k = params.problem_size.k(); + + ElementA *ptr_A = static_cast(params.ptr_A); + ElementB *ptr_B = static_cast(params.ptr_B); + + // + // Fetch pointers based on mode. + // + if (params.mode == GemmUniversalMode::kGemm || + params.mode == GemmUniversalMode::kGemmSplitKParallel) { + + if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) { + + problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size; + } + + offset_k = threadblock_tile_offset.k() * params.gemm_k_size; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_A += threadblock_tile_offset.k() * params.batch_stride_A; + ptr_B += threadblock_tile_offset.k() * params.batch_stride_B; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_A = static_cast(params.ptr_A)[threadblock_tile_offset.k()]; + ptr_B = static_cast(params.ptr_B)[threadblock_tile_offset.k()]; + } + + __syncthreads(); + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + offset_k, + }; + + mctlass::MatrixCoord tb_offset_B{ + offset_k, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, + ptr_A, + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + params.params_B, + ptr_B, + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + typename Mma::FragmentReduction gemm_k_accumulators; + + gemm_k_accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma( + gemm_k_iterations, + accumulators, + iterator_A, + iterator_B, + accumulators, + gemm_k_accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + ElementC *ptr_C = static_cast(params.ptr_C); + ElementC *ptr_D = static_cast(params.ptr_D); + ElementC *ptr_gemm_k_reduction = static_cast(params.ptr_gemm_k_reduction); + + // + // Fetch pointers based on mode. + // + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + if (params.mode == GemmUniversalMode::kGemm) { + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + } + else if (params.mode == GemmUniversalMode::kGemmSplitKParallel) { + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + ptr_gemm_k_reduction += threadblock_tile_offset.k() * params.batch_stride_gemm_k_reduction; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_C += threadblock_tile_offset.k() * params.batch_stride_C; + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_C = static_cast(params.ptr_C)[threadblock_tile_offset.k()]; + ptr_D = static_cast(params.ptr_D)[threadblock_tile_offset.k()]; + } + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + ptr_C, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + + } + + // Execute the epilogue operator to update the destination tensor. + epilogue( + output_op, + iterator_D, + accumulators, + iterator_C); + + if ((kReduceKForA && threadblock_tile_offset.n() == 0) + || (!kReduceKForA && threadblock_tile_offset.m() == 0)) { + + int warp_idx_mn = warp_idx % (Mma::Base::WarpCount::kM * Mma::Base::WarpCount::kN); + int warp_idx_m = warp_idx_mn % Mma::Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Mma::Base::WarpCount::kM; + + if ((kReduceKForA && warp_idx_n == 0) + || (!kReduceKForA && warp_idx_m == 0)) { + + int reduction_warp_idx = kReduceKForA ? warp_idx_m : warp_idx_n; + int reduction_threadblock_offset = kReduceKForA ? threadblock_tile_offset.m() : + threadblock_tile_offset.n(); + int reduction_vector_size = kReduceKForA ? params.problem_size.m() + : params.problem_size.n(); + EpilogueGemmKReduction epilogue_gemm_k_reduction(thread_idx, + reduction_warp_idx, + lane_idx, + reduction_threadblock_offset, + ptr_gemm_k_reduction); + epilogue_gemm_k_reduction( + reduction_vector_size, + gemm_k_accumulators, + params.mode == GemmUniversalMode::kGemm + && (params.grid_tiled_shape.k() > 1) + && (threadblock_tile_offset.k() > 0)); + } + } + + // + // Release the semaphore + // + + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemv.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemv.h new file mode 100644 index 0000000..60e4f1f --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemv.h @@ -0,0 +1,638 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/tensor_ref.h" + +#include "mctlass/arch/memory.h" +#include "mctlass/arch/cache_operation.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/layout/matrix.h" + +#include "mctlass/numeric_conversion.h" +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementA_, + typename LayoutA_, + typename ElementB_, + typename ElementC_, + typename ElementAccumulator_, + typename EpilogueOutputOp_, + int kElementsPerAccess_ = 1, ///< Number of elements involved in a global access. + int kThreadCount_ = 0, ///< Number of threads in the thread block. + /// It will be calculated automatically if set to 0. + int kThreadsPerRow_ = 0 ///< Number of threads in the k dimension. + /// It will be calculated automatically if set to 0. +> +struct Gemv; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Specializations +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +// GEMV for column-major A matrix +template < + typename ElementA_, + typename ElementB_, + typename ElementC_, + typename ElementAccumulator_, + typename EpilogueOutputOp_, + int kElementsPerAccess_, + int kThreadCount_, + int kThreadsPerRow_ +> +struct Gemv < + ElementA_, + layout::ColumnMajor, + ElementB_, + ElementC_, + ElementAccumulator_, + EpilogueOutputOp_, + kElementsPerAccess_, + kThreadCount_, + kThreadsPerRow_ +>{ +public: + + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using TensorRefA = TensorRef; + + using ElementB = ElementB_; + using ElementC = ElementC_; + + using ElementAccumulator = ElementAccumulator_; + using EpilogueOutputOp = EpilogueOutputOp_; + + static ComplexTransform const kTransformA = ComplexTransform::kNone; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + + // thread block shape (kThreadCount, 1, 1) + static int const kThreadCount = (kThreadCount_ == 0) ? 32 : kThreadCount_; + static int const kThreadsPerRow = kThreadsPerRow_; + + static int const kStages = 1; + + static int const kAlignmentA = 1; + static int const kAlignmentB = 1; + static int const kAlignmentC = 1; + + // + // Structures + // + + /// Argument structure + struct Arguments { + MatrixCoord problem_size; + int32_t batch_count; + typename EpilogueOutputOp::Params output_op; + + TensorRefA ref_A; + + ElementB const *ptr_B; + ElementC const *ptr_C; + ElementC *ptr_D; + + int64_t inc_B; + int64_t inc_C; + int64_t inc_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + int64_t batch_stride_D; + + // + // Methods + // + + Arguments(): batch_count(0) { } + + Arguments( + MatrixCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params output_op, + TensorRefA ref_A, + void const *ptr_B, + void const *ptr_C, + void *ptr_D, + int64_t inc_B, + int64_t inc_C, + int64_t inc_D, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C, + int64_t batch_stride_D + ): + problem_size(problem_size), + batch_count(batch_count), + output_op(output_op), + ref_A(ref_A), + ptr_B(static_cast(ptr_B)), + ptr_C(static_cast(ptr_C)), + ptr_D(static_cast(ptr_D)), + inc_B(inc_B), + inc_C(inc_C), + inc_D(inc_D), + batch_stride_A(batch_stride_A), + batch_stride_B(batch_stride_B), + batch_stride_C(batch_stride_C), + batch_stride_D(batch_stride_D) + { } + + Arguments( + MatrixCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params output_op, + TensorRefA ref_A, + void const *ptr_B, + void const *ptr_C, + void *ptr_D, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C, + int64_t batch_stride_D + ): + Arguments( + problem_size, + batch_count, + output_op, + ref_A, + ptr_B, + ptr_C, + ptr_D, + 1, + 1, + 1, + batch_stride_A, + batch_stride_B, + batch_stride_C, + batch_stride_D) + { } + + Arguments( + MatrixCoord problem_size, + typename EpilogueOutputOp::Params output_op, + TensorRefA ref_A, + void const *ptr_B, + void const *ptr_C, + void *ptr_D, + int64_t inc_B, + int64_t inc_C, + int64_t inc_D + ): + Arguments( + problem_size, + 1, + output_op, + ref_A, + ptr_B, + ptr_C, + ptr_D, + inc_B, + inc_C, + inc_D, + 1, + 1, + 1, + 1) + { } + + Status update(Arguments const &args) { + output_op = args.output_op; + ref_A = ref_A; + ptr_B = args.ptr_B; + ptr_C = args.ptr_C; + ptr_D = args.ptr_D; + + return Status::kSuccess; + } + }; + + using Params = Arguments; + + /// Shared memory storage structure + union SharedStorage { + + }; + +public: + + // + // Methods + // + + MCTLASS_DEVICE + Gemv() { } + + /// Determines whether kernel satisfies alignment + static Status can_implement(mctlass::MatrixCoord const & problem_size) { + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + + /// Executes one GEMV + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Loop over batch indices + for (int batch_idx = blockIdx.z; batch_idx < params.batch_count; batch_idx += gridDim.z) { + + int i = blockIdx.x * kThreadCount + threadIdx.x; + + ElementA const *ptr_A = params.ref_A.data() + i; + ElementB const *ptr_B = params.ptr_B; + + ptr_A += batch_idx * params.batch_stride_A; + ptr_B += batch_idx * params.batch_stride_B; + + ElementAccumulator accum = ElementAccumulator(); + + // Compute inner product + MCTLASS_PRAGMA_NO_UNROLL + for (int k = 0; k < params.problem_size.column(); ++k) { + + // Fetch from A + ElementA a = ElementA(); + if (i < params.problem_size.row()) { + a = *ptr_A; + } + ptr_A += params.ref_A.stride(0); + + // Fetch from B + ElementB b = *ptr_B; + ptr_B += params.inc_B; + + // Math + accum += ElementAccumulator(a) * ElementAccumulator(b); + } + + // + // Epilogue phase + // + + ElementC const *ptr_C = params.ptr_C + i * params.inc_C + batch_idx * params.batch_stride_C; + ElementC *ptr_D = params.ptr_D + i * params.inc_D + batch_idx * params.batch_stride_D; + + EpilogueOutputOp output_op(params.output_op); + + typename EpilogueOutputOp::FragmentAccumulator accum_fragment; + typename EpilogueOutputOp::FragmentOutput source_fragment; + typename EpilogueOutputOp::FragmentOutput output_fragment; + + accum_fragment[0] = accum; + + if (i < params.problem_size.row()) { + if (output_op.is_source_needed()) { + source_fragment[0] = *ptr_C; + output_fragment = output_op(accum_fragment, source_fragment); + } + else { + output_fragment = output_op(accum_fragment); + } + + *ptr_D = output_fragment[0]; + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// GEMV for row-major A matrix +template < + typename ElementA_, + typename ElementB_, + typename ElementC_, + typename ElementAccumulator_, + typename EpilogueOutputOp_, + int kElementsPerAccess_, + int kThreadCount_, + int kThreadsPerRow_ +> +struct Gemv < + ElementA_, + layout::RowMajor, + ElementB_, + ElementC_, + ElementAccumulator_, + EpilogueOutputOp_, + kElementsPerAccess_, + kThreadCount_, + kThreadsPerRow_ +>{ +public: + + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using TensorRefA = TensorRef; + + using ElementB = ElementB_; + using ElementC = ElementC_; + + using ElementAccumulator = ElementAccumulator_; + using EpilogueOutputOp = EpilogueOutputOp_; + + static ComplexTransform const kTransformA = ComplexTransform::kNone; + static ComplexTransform const kTransformB = ComplexTransform::kNone; + + static FloatRoundStyle const Round = mctlass::FloatRoundStyle::round_to_nearest; + + // number of return elements in a global access + static int const kElementsPerAccess = kElementsPerAccess_; + + using FragmentA = Array; + using FragmentB = Array; + using FragmentCompute = Array; + + // thread block shape (kThreadsPerRow, kThreadCount / kThreadsPerRow, 1) + static int const kThreadCount = (kThreadCount_ == 0) ? 128 : kThreadCount_; + static int const kThreadsPerRow = (kThreadsPerRow_ == 0) ? + std::min(static_cast(kThreadCount / (kElementsPerAccess * sizeof(ElementA))), 16) + : kThreadsPerRow_; + + // + // Structures + // + + /// Argument structure + struct Arguments { + MatrixCoord problem_size; + int32_t batch_count; + typename EpilogueOutputOp::Params output_op; + + TensorRefA ref_A; + + ElementB const *ptr_B; + ElementC const *ptr_C; + ElementC *ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + int64_t batch_stride_D; + + // + // Methods + // + + Arguments(): batch_count(0) { } + + Arguments( + MatrixCoord problem_size, + int32_t batch_count, + typename EpilogueOutputOp::Params output_op, + TensorRefA ref_A, + void const *ptr_B, + void const *ptr_C, + void *ptr_D, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C, + int64_t batch_stride_D + ): + problem_size(problem_size), + batch_count(batch_count), + output_op(output_op), + ref_A(ref_A), + ptr_B(static_cast(ptr_B)), + ptr_C(static_cast(ptr_C)), + ptr_D(static_cast(ptr_D)), + batch_stride_A(batch_stride_A), + batch_stride_B(batch_stride_B), + batch_stride_C(batch_stride_C), + batch_stride_D(batch_stride_D) + { } + + Arguments( + MatrixCoord problem_size, + typename EpilogueOutputOp::Params output_op, + TensorRefA ref_A, + void const *ptr_B, + void const *ptr_C, + void *ptr_D + ): + Arguments( + problem_size, + 1, + output_op, + ref_A, + ptr_B, + ptr_C, + ptr_D, + 1, + 1, + 1, + 1) + { } + + Status update(Arguments const &args) { + problem_size = args.problem_size; + batch_count = args.batch_count; + output_op = args.output_op; + ref_A = ref_A; + ptr_B = args.ptr_B; + ptr_C = args.ptr_C; + ptr_D = args.ptr_D; + batch_stride_A = args.batch_stride_A; + batch_stride_B = args.batch_stride_B; + batch_stride_C = args.batch_stride_C; + batch_stride_D = args.batch_stride_D; + + return Status::kSuccess; + } + }; + + using Params = Arguments; + + /// Shared memory storage structure + union SharedStorage { + + }; + +public: + + // + // Methods + // + + MCTLASS_DEVICE + Gemv() {} + + /// Determines whether kernel satisfies alignment + static Status can_implement(mctlass::MatrixCoord const &problem_size) { + if (problem_size.column() % kElementsPerAccess != 0) { + return Status::kErrorMisalignedOperand; + } + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + + /// Executes one GEMV + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Loop over batch indices + for (int batch_idx = blockIdx.z; batch_idx < params.batch_count; batch_idx += gridDim.z) { + int idx_col_k = threadIdx.x; + int idx_row_m = blockIdx.x * blockDim.y + threadIdx.y; + + if (idx_row_m < params.problem_size.row()) { + // problem_size (row = m, column = k) + // matrix A (batch, m, k) + // vector B (batch, 1, k) + // vector C (batch, m, 1) + // vector D (batch, m, 1) + + // move in the batch dimension + ElementA const *ptr_A = params.ref_A.data() + batch_idx * params.batch_stride_A; + ElementB const *ptr_B = params.ptr_B + batch_idx * params.batch_stride_B; + + ElementC const *ptr_C = params.ptr_C + batch_idx * params.batch_stride_C; + ElementC *ptr_D = params.ptr_D + batch_idx * params.batch_stride_D; + + // move in the k dimension + ptr_A += idx_col_k * kElementsPerAccess; + ptr_B += idx_col_k * kElementsPerAccess; + + // move in the m dimension + ptr_A += idx_row_m * params.problem_size.column(); + ptr_C += idx_row_m; + ptr_D += idx_row_m; + + NumericArrayConverter srcA_converter; + NumericArrayConverter srcB_converter; + + ElementAccumulator accum = 0.f; + + FragmentB fragB; + FragmentA fragA; + + int unroll_col_k = 0; + + // rows of the rolling tile + int const tileA_k = kThreadsPerRow * kElementsPerAccess; + + for (; unroll_col_k < params.problem_size.column() / tileA_k * tileA_k; unroll_col_k += tileA_k) { + + // fetch from matrix A + arch::global_load(fragA, (ptr_A + unroll_col_k), true); + + // fetch from vector B + arch::global_load(fragB, (ptr_B + unroll_col_k), true); + + FragmentCompute fragB_Compute = srcB_converter(fragB); + FragmentCompute fragA_Compute = srcA_converter(fragA); + + // Math + MCTLASS_PRAGMA_UNROLL + for (int e = 0; e < kElementsPerAccess; e++) { + accum += fragA_Compute.at(e) * fragB_Compute.at(e); + } + } + + // calculate the rest of K elements + // each thread fetch 1 element each time + for (int k = unroll_col_k + idx_col_k; k < params.problem_size.column(); k += kThreadsPerRow) { + ElementB b = *(ptr_B - idx_col_k * kElementsPerAccess + k); + ElementA a = *(ptr_A - idx_col_k * kElementsPerAccess + k); + + accum += ElementAccumulator(a) * ElementAccumulator(b); + } + + EpilogueOutputOp output_op(params.output_op); + typename EpilogueOutputOp::FragmentOutput source_fragment; + + // prefetch from source matrix C + if (output_op.is_source_needed()) { + source_fragment[0] = *(ptr_C); + } + + typename EpilogueOutputOp::FragmentAccumulator accum_fragment; + typename EpilogueOutputOp::FragmentOutput output_fragment; + + for (int mask = (kThreadsPerRow >> 1); mask > 0; mask >>= 1) { + accum += __shfl_xor_sync(0xFFFFFFFF, accum, mask, 32); + } + + if (idx_col_k == 0) { + accum_fragment[0] = accum; + + if (output_op.is_source_needed()) { + output_fragment = output_op(accum_fragment, source_fragment); + } + else { + output_fragment = output_op(accum_fragment); + } + + *ptr_D = output_fragment[0]; + } + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/gemv_batched_strided.h b/csrc/mctlass/include/mctlass/gemm/kernel/gemv_batched_strided.h new file mode 100644 index 0000000..5e15b03 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/gemv_batched_strided.h @@ -0,0 +1,241 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/aligned_buffer.h" +#include "mctlass/array.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/gemm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +namespace detail +{ + template + struct GemvBatchedStridedEpilogueScaling + { + ElementAlphaBeta const & alpha; + ElementAlphaBeta const & beta; + + MCTLASS_DEVICE + GemvBatchedStridedEpilogueScaling(ElementAlphaBeta& alpha_, ElementAlphaBeta& beta_) : + alpha(alpha_), beta(beta_) + { } + + template + MCTLASS_DEVICE + void operator()(FragmentAccumulator& accumulators, + FragmentCD const& fragment_C, + FragmentCD& fragment_D) const + { + using AccType = typename FragmentAccumulator::value_type; + using CDType = typename FragmentCD::value_type; + + static_assert(FragmentCD::kElements == FragmentAccumulator::kElements, + "Mistmatch in fragment sizes."); + + for (int i = 0; i < FragmentCD::kElements; ++i) + { + if (BetaIsZero) + { + fragment_D[i] = CDType(accumulators[i] * AccType(alpha)); + } + else + { + fragment_D[i] = CDType(accumulators[i] * AccType(alpha) + + AccType(fragment_C[i]) * AccType(beta)); + } + } + } + }; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +MCTLASS_DEVICE void GemvBatchedStridedDevice( + mctlass::gemm::BatchedGemmCoord problem_size, + ElementAlphaBeta alpha, + ElementAlphaBeta beta, + typename GemvKernel::IteratorA::TensorRef ref_A, + typename GemvKernel::IteratorA::TensorRef::LongIndex lda, + typename GemvKernel::IteratorB::TensorRef ref_B, + typename GemvKernel::IteratorB::TensorRef::LongIndex ldb, + typename GemvKernel::IteratorCD::TensorRef ref_C, + typename GemvKernel::IteratorCD::TensorRef::LongIndex ldc, + typename GemvKernel::IteratorCD::TensorRef ref_D, + typename GemvKernel::IteratorCD::TensorRef::LongIndex ldd) +{ + using ThreadBlockGemv = typename GemvKernel::ThreadBlockGemv; + using ThreadBlockSwizzle = typename GemvKernel::ThreadBlockSwizzle; + using EpilogueScale = detail::GemvBatchedStridedEpilogueScaling; + + ThreadBlockSwizzle swizzler; + + // Compute initial location in logical coordinates + BatchedGemmCoord tb_offset = swizzler.get_tile_offset(); + int const batch_idx = swizzler.get_batch_idx(); + + // Offset to the batch + ref_A.add_pointer_offset(batch_idx*lda); + ref_B.add_pointer_offset(batch_idx*ldb); + + // Construct iterators to A and B operands + typename GemvKernel::IteratorA::Params params_A(ref_A.layout()); + typename GemvKernel::IteratorA iterator_A( + params_A, + ref_A.data(), + { 1, problem_size.k() }, + 0, + { 0, 0 }); + + typename GemvKernel::IteratorB::Params params_B(ref_B.layout()); + typename GemvKernel::IteratorB iterator_B( + params_B, + ref_B.data(), + { problem_size.k(), problem_size.n() }, + threadIdx.x, + { 0, tb_offset.n()*ThreadBlockGemv::Shape::kN }); + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + ThreadBlockGemv mma; + + typename ThreadBlockGemv::FragmentC accumulators; + accumulators.clear(); + + // Compute threadblock-scoped gemv + mma(problem_size.mnk(), accumulators, iterator_A, iterator_B, accumulators); + + typename GemvKernel::FragmentCD fragment_CD; + + // Load C (skip if beta is zero) + if (!BetaIsZero) + { + tb_offset = swizzler.get_tile_offset(); + ref_C.add_pointer_offset(batch_idx*ldc); + typename GemvKernel::IteratorCD::Params params_C(ref_C.layout()); + typename GemvKernel::IteratorCD iterator_C( + params_C, + ref_C.data(), + { 1, problem_size.n() }, + threadIdx.x, + { 0, tb_offset.n()*ThreadBlockGemv::Shape::kN }); + iterator_C.load(fragment_CD); + } + + // Apply alpha/beta scaling + EpilogueScale epilogue_scale(alpha, beta); + epilogue_scale(accumulators, fragment_CD, fragment_CD); + + // Store D + tb_offset = swizzler.get_tile_offset(); + ref_D.add_pointer_offset(batch_idx*ldd); + typename GemvKernel::IteratorCD::Params params_D(ref_D.layout()); + typename GemvKernel::IteratorCD iterator_D( + params_D, + ref_D.data(), + { 1, problem_size.n() }, + threadIdx.x, + { 0, tb_offset.n()*ThreadBlockGemv::Shape::kN }); + iterator_D.store(fragment_CD); +} + +template +__global__ void GemvBatchedStrided( + mctlass::gemm::BatchedGemmCoord problem_size, + ElementAlphaBeta alpha, + ElementAlphaBeta beta, + typename GemvKernel::IteratorA::TensorRef ref_A, + typename GemvKernel::IteratorA::TensorRef::LongIndex lda, + typename GemvKernel::IteratorB::TensorRef ref_B, + typename GemvKernel::IteratorB::TensorRef::LongIndex ldb, + typename GemvKernel::IteratorCD::TensorRef ref_C, + typename GemvKernel::IteratorCD::TensorRef::LongIndex ldc, + typename GemvKernel::IteratorCD::TensorRef ref_D, + typename GemvKernel::IteratorCD::TensorRef::LongIndex ldd) +{ + GemvBatchedStridedDevice( + problem_size, alpha, beta, ref_A, lda, ref_B, ldb, ref_C, ldc, ref_D, ldd + ); +} + +template +__global__ void GemvBatchedStrided( + mctlass::gemm::BatchedGemmCoord problem_size, + ElementAlphaBeta alpha, + typename GemvKernel::IteratorA::TensorRef ref_A, + typename GemvKernel::IteratorA::TensorRef::LongIndex lda, + typename GemvKernel::IteratorB::TensorRef ref_B, + typename GemvKernel::IteratorB::TensorRef::LongIndex ldb, + typename GemvKernel::IteratorCD::TensorRef ref_D, + typename GemvKernel::IteratorCD::TensorRef::LongIndex ldd) +{ + GemvBatchedStridedDevice( + problem_size, alpha, ElementAlphaBeta(0), ref_A, lda, ref_B, ldb, ref_D, ldd, ref_D, ldd + ); +} + +template +__global__ void GemvBatchedStrided( + mctlass::gemm::BatchedGemmCoord problem_size, + typename GemvKernel::IteratorA::TensorRef ref_A, + typename GemvKernel::IteratorA::TensorRef::LongIndex lda, + typename GemvKernel::IteratorB::TensorRef ref_B, + typename GemvKernel::IteratorB::TensorRef::LongIndex ldb, + typename GemvKernel::IteratorCD::TensorRef ref_D, + typename GemvKernel::IteratorCD::TensorRef::LongIndex ldd) +{ + using ElementAlphaBeta = typename GemvKernel::IteratorCD::Element; + GemvBatchedStridedDevice( + problem_size, ElementAlphaBeta(1), ElementAlphaBeta(0), ref_A, lda, ref_B, ldb, ref_D, ldd, ref_D, ldd + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/grouped_problem_visitor.h b/csrc/mctlass/include/mctlass/gemm/kernel/grouped_problem_visitor.h new file mode 100644 index 0000000..9a0c4f7 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/grouped_problem_visitor.h @@ -0,0 +1,470 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Base scheduler for grouped problems +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Enumerated type describing the type of scheduling to perform for the ProblemVisitor +enum class GroupScheduleMode { + // Perform all scheduling on device + kDeviceOnly, + // Precompute on the host the full sequence of problems to access + kHostPrecompute +}; + +/// Visitor class to abstract away the algorithm for iterating over tiles +template +struct BaseGroupedProblemVisitor { + using ThreadblockShape = ThreadblockShape_; + + struct ProblemInfo { + static int32_t const kNoPrefetchEntry = -1; + int32_t problem_idx; + int32_t problem_start; + + MCTLASS_DEVICE + ProblemInfo() : problem_idx(kNoPrefetchEntry), problem_start(kNoPrefetchEntry) {} + + //Add HOST function for MACA + ProblemInfo() : problem_idx(kNoPrefetchEntry), problem_start(kNoPrefetchEntry) {} + + MCTLASS_DEVICE + ProblemInfo(int32_t problem_idx_, int32_t problem_start_) : + problem_idx(problem_idx_), problem_start(problem_start_) {} + + //Add HOST function for MACA + ProblemInfo(int32_t problem_idx_, int32_t problem_start_) : + problem_idx(problem_idx_), problem_start(problem_start_) {} + }; + + struct Params { + mctlass::gemm::GemmCoord const *problem_sizes; + int32_t problem_count; + void const *workspace; + int32_t tile_count; + + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + Params(): problem_sizes(nullptr), problem_count(0), workspace(nullptr), tile_count(0) { } + + /// Ctor + MCTLASS_HOST_DEVICE + Params( + mctlass::gemm::GemmCoord const *problem_sizes, + int32_t problem_count, + void const *workspace = nullptr, + int32_t tile_count = 0 + ): + problem_sizes(problem_sizes), + problem_count(problem_count), + workspace(workspace), + tile_count(tile_count) + {} + + }; + + Params params; + int32_t tile_idx; + int32_t problem_tile_start; + int32_t problem_idx; + + // + // Methods + // + MCTLASS_DEVICE + BaseGroupedProblemVisitor( + Params const ¶ms_, + int32_t block_idx + ): + params(params_), + tile_idx(block_idx), + problem_tile_start(0), + problem_idx(0) + {} + + /// Get the grid shape + MCTLASS_HOST_DEVICE + static mctlass::gemm::GemmCoord grid_shape(const mctlass::gemm::GemmCoord& problem) { + return ProblemSizeHelper::grid_shape(problem); + } + + /// Gets the global tile index + MCTLASS_HOST_DEVICE + int32_t tile_index() const { + return tile_idx; + } + + /// Gets the index of the problem + MCTLASS_HOST_DEVICE + int32_t problem_index() const { + return problem_idx; + } + + MCTLASS_HOST_DEVICE + int32_t threadblock_idx() const { + return tile_idx - problem_tile_start; + } + + MCTLASS_DEVICE + void advance(int32_t grid_size) { + tile_idx += grid_size; + } + + MCTLASS_HOST_DEVICE + static void possibly_transpose_problem(mctlass::gemm::GemmCoord& problem) { + ProblemSizeHelper::possibly_transpose_problem(problem); + } + + /// Returns the problem size for the current problem + MCTLASS_HOST_DEVICE + mctlass::gemm::GemmCoord problem_size() const { + GemmCoord problem = params.problem_sizes[problem_idx]; + ProblemSizeHelper::possibly_transpose_problem(problem); + return problem; + } + + MCTLASS_HOST_DEVICE + static int32_t tile_count(const mctlass::gemm::GemmCoord& grid) { + return ProblemSizeHelper::tile_count(grid); + } + + static int32_t group_tile_count(const mctlass::gemm::GemmCoord* host_problem_sizes_ptr, int32_t problem_count) { + int32_t total_tiles = 0; + for (int32_t i = 0; i < problem_count; ++i) { + auto problem = host_problem_sizes_ptr[i]; + possibly_transpose_problem(problem); + auto grid = grid_shape(problem); + total_tiles += tile_count(grid); + } + + return total_tiles; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ProblemSizeHelper, + typename ThreadblockShape, + GroupScheduleMode GroupScheduleMode_, + int PrefetchTileCount, + int ThreadCount +> +struct GroupedProblemVisitor; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// ProblemVisitor that performs all scheduling on device +// +template +struct GroupedProblemVisitor: public BaseGroupedProblemVisitor { + using Base = BaseGroupedProblemVisitor; + using Params = typename Base::Params; + static int const kThreadCount = ThreadCount; + static bool const kRequiresPrecomputation = false; + static int const kThreadsPerWarp = 32; + + struct SharedStorage {}; + + // Final tile of the problem loaded by this thread. Each thread will hold + // a separate value. + int32_t problem_ending_tile; + + SharedStorage &shared_storage; + + // + // Methods + // + MCTLASS_DEVICE + GroupedProblemVisitor( + Params const ¶ms_, + SharedStorage &shared_storage_, + int32_t block_idx + ): Base(params_, block_idx), + problem_ending_tile(0), + shared_storage(shared_storage_) + { + this->problem_idx = -1 * kThreadsPerWarp; + this->problem_tile_start = 0; + } + + MCTLASS_DEVICE + bool next_tile() { + // Check whether the tile to compute is within the range of the current problem. + int32_t problem_tile_end = __shfl_sync(0xffffffff, problem_ending_tile, this->problem_idx % kThreadsPerWarp); + if (this->tile_idx < problem_tile_end) { + return true; + } + + // Check whether the tile to compute is within the current group of problems fetched by the warp. + // The last tile for this group is the final tile of the problem held by the final thread in the warp. + int32_t group_tile_end = __shfl_sync(0xffffffff, problem_ending_tile, kThreadsPerWarp-1); + + // Keep the starting problem for this group in `problem_idx`. This is done to reduce + // register pressure. The starting problem for this group is simply the first problem + // in the group most recently fetched by the warp. + int32_t &group_problem_start = this->problem_idx; + group_problem_start = (this->problem_idx / kThreadsPerWarp) * kThreadsPerWarp; + + // Keep the starting tile for this group in `problem_tile_start`. This is done to reduce + // register pressure. + int32_t &group_tile_start = this->problem_tile_start; + + // Each thread in the warp processes a separate problem to advance until + // reaching a problem whose starting tile is less less than tile_idx. + while (group_tile_end <= this->tile_idx) { + group_problem_start += kThreadsPerWarp; + if (group_problem_start > this->params.problem_count) { + return false; + } + + // Since `group_tile_start` is a reference to `this->problem_tile_start`, this + // also sets `this->problem_tile_start`. The fact that `this->problem_tile_start` + // is also set here is used later in `next_tile`. + group_tile_start = group_tile_end; + + int lane_idx = threadIdx.x % kThreadsPerWarp; + int32_t lane_problem = group_problem_start + lane_idx; + + // Compute the number of tiles in the problem assigned to each thread. + problem_ending_tile = 0; + if (lane_problem < this->params.problem_count) { + mctlass::gemm::GemmCoord problem = this->params.problem_sizes[lane_problem]; + this->possibly_transpose_problem(problem); + mctlass::gemm::GemmCoord grid = this->grid_shape(problem); + problem_ending_tile = this->tile_count(grid); + } + + // Compute a warp-wide inclusive prefix sum to compute the ending tile index of + // each thread's problem. + MCTLASS_PRAGMA_UNROLL + for (int i = 1; i < kThreadsPerWarp; i <<= 1) { + int32_t val = __shfl_up_sync(0xffffffff, problem_ending_tile, i); + if (lane_idx >= i) { + problem_ending_tile += val; + } + } + + // The total tile count for this group is now in the final position of the prefix sum + int32_t tiles_in_group = __shfl_sync(0xffffffff, problem_ending_tile, kThreadsPerWarp-1); + + problem_ending_tile += group_tile_start; + group_tile_end += tiles_in_group; + } + + // The next problem to process is the first one that does not have ending tile position + // that is greater than or equal to tile index. + + int32_t problem_idx_in_group = + __popc(__ballot_sync(0xffffffff, problem_ending_tile <= this->tile_idx)); + this->problem_idx = group_problem_start + problem_idx_in_group; + + // The starting tile for this problem is the ending tile of the previous problem. In cases + // where `problem_idx_in_group` is the first problem in the group, we do not need to reset + // `problem_tile_start`, because it is set to the previous group's ending tile in the while + // loop above. + if (problem_idx_in_group > 0) { + this->problem_tile_start = __shfl_sync(0xffffffff, problem_ending_tile, problem_idx_in_group - 1); + } + + return true; + } + + static size_t get_workspace_size(const mctlass::gemm::GemmCoord* host_problem_sizes_ptr, + int32_t problem_count, + int32_t block_count) { + return 0; + } + + static void host_precompute(const mctlass::gemm::GemmCoord* host_problem_sizes_ptr, + int32_t problem_count, + int32_t block_count, + void* host_workspace_ptr) {} +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Precomputes schedule on host and prefetches into shared memory +// +template +struct GroupedProblemVisitor : public BaseGroupedProblemVisitor { + static_assert(PrefetchTileCount > 0, + "GroupedProblemVisitor with GroupScheduleMode `kHostPrecompute` currently requires prefetching to shared memory"); + + using Base = BaseGroupedProblemVisitor; + using Params = typename Base::Params; + using ProblemInfo = typename Base::ProblemInfo; + static bool const kRequiresPrecomputation = true; + + static int const kPrefetchTileCount = PrefetchTileCount; + static int const kThreadCount = ThreadCount; + + struct SharedStorage { + // Sequence of problem IDs and starting tiles to compute + mctlass::Array prefetched_problems; + }; + + int32_t tiles_computed; + int32_t iterations_per_block; + int32_t block_load_start; + SharedStorage &shared_storage; + ProblemInfo const *problem_info_ptr; + + // + // Methods + // + MCTLASS_DEVICE + GroupedProblemVisitor( + Params const ¶ms_, + SharedStorage &shared_storage_, + int32_t block_idx + ): Base(params_, block_idx), + tiles_computed(0), + shared_storage(shared_storage_), + problem_info_ptr(reinterpret_cast(params_.workspace)) + { + iterations_per_block = (params_.tile_count - 1 + gridDim.x) / gridDim.x; + block_load_start = iterations_per_block * block_idx; + // Start prefetching the first set of tiles to compute + prefetch_tiles(); + } + + MCTLASS_DEVICE + bool next_tile() { + if (this->tile_idx >= this->params.tile_count) { + return false; + } + + int32_t prefetch_idx = (tiles_computed % kPrefetchTileCount); + if (prefetch_idx == 0) { + // Ensure all previous stores to shared memory have been completed + __syncthreads(); + } + + auto problem_info = shared_storage.prefetched_problems[prefetch_idx]; + ++tiles_computed; + + if ((tiles_computed % kPrefetchTileCount) == 0) { + // Begin prefetching next set of tiles. Synchronize first to ensure that + // we don't overwrite the current buffer while someone else is using it. + __syncthreads(); + prefetch_tiles(); + } + + this->problem_idx = problem_info.problem_idx; + this->problem_tile_start = problem_info.problem_start; + + return true; + } + + static size_t get_workspace_size(const mctlass::gemm::GemmCoord* host_problem_sizes_ptr, + int32_t problem_count, + int32_t block_count) { + int32_t total_tiles = Base::group_tile_count(host_problem_sizes_ptr, problem_count); + int32_t entries_per_block = ((total_tiles - 1 + block_count) / block_count); + return sizeof(ProblemInfo) * entries_per_block * block_count; + } +#if !defined(__MACACC_RTC__) + static void host_precompute(const mctlass::gemm::GemmCoord* host_problem_sizes_ptr, + int32_t problem_count, + int32_t block_count, + void* host_workspace_ptr) { + ProblemInfo* host_problem_info_ptr = reinterpret_cast(host_workspace_ptr); + int32_t total_tiles = Base::group_tile_count(host_problem_sizes_ptr, problem_count); + int32_t entries_per_block = (total_tiles - 1 + block_count) / block_count; + + int tile = 0; + int start_tile = 0; + for (int p_idx = 0; p_idx < problem_count; ++p_idx) { + auto problem = host_problem_sizes_ptr[p_idx]; + Base::possibly_transpose_problem(problem); + auto grid = Base::grid_shape(problem); + int tiles = Base::tile_count(grid); + ProblemInfo problem_info(p_idx, start_tile); + for (int i = 0; i < tiles; ++i, ++tile) { + host_problem_info_ptr[(entries_per_block * (tile % block_count)) + (tile / block_count)] = problem_info; + } + start_tile += tiles; + } + } +#endif +private: + MCTLASS_DEVICE + void prefetch_tiles() { + MCTLASS_PRAGMA_UNROLL + for (int32_t i = 0; i < kPrefetchTileCount; i += kThreadCount) { + int32_t offset = threadIdx.x + i; + if (offset < kPrefetchTileCount && (tiles_computed + offset < iterations_per_block)) { + shared_storage.prefetched_problems[offset] = problem_info_ptr[block_load_start + tiles_computed + offset]; + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/params_universal_base.h b/csrc/mctlass/include/mctlass/gemm/kernel/params_universal_base.h new file mode 100644 index 0000000..ea60bf0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/params_universal_base.h @@ -0,0 +1,245 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Base functionality for common types of universal GEMM kernel parameters +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/trace.h" +#include "mctlass/gemm/gemm.h" + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + + +/// Argument structure +struct UniversalArgumentsBase +{ + // + // Data members + // + + GemmUniversalMode mode; + GemmCoord problem_size; + int batch_count; + + int64_t batch_stride_D; + + // + // Methods + // + + UniversalArgumentsBase() : + mode(GemmUniversalMode::kGemm), + batch_count(1), + batch_stride_D(0) + {} + + /// constructs an arguments structure + UniversalArgumentsBase( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_count, + int64_t batch_stride_D) + : + mode(mode), + problem_size(problem_size), + batch_count(batch_count), + batch_stride_D(batch_stride_D) + { + MCTLASS_TRACE_HOST("GemmUniversal::Arguments::Arguments() - problem_size: " << problem_size); + } +}; + + +/// Parameters structure +template < + typename ThreadblockSwizzle, + typename ThreadblockShape, + typename ElementA, + typename ElementB, + typename ElementC> +struct UniversalParamsBase +{ + // + // Data members + // + + GemmCoord problem_size; + GemmCoord grid_tiled_shape; + int swizzle_log_tile; + + GemmUniversalMode mode; + int batch_count; + int gemm_k_size; + + int64_t batch_stride_D; + + int *semaphore; + + + // + // Host dispatch API + // + + /// Default constructor + UniversalParamsBase() = default; + + + /// Constructor + UniversalParamsBase( + UniversalArgumentsBase const &args, /// GEMM application arguments + int device_sms, /// Number of SMs on the device + int sm_occupancy) /// Kernel SM occupancy (in thread blocks) + : + problem_size(args.problem_size), + mode(args.mode), + batch_count(args.batch_count), + batch_stride_D(args.batch_stride_D), + semaphore(nullptr) + { + ThreadblockSwizzle swizzle; + + // Get GEMM volume in thread block tiles + grid_tiled_shape = swizzle.get_tiled_shape( + args.problem_size, + {ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK}, + args.batch_count); + + swizzle_log_tile = swizzle.get_log_tile(grid_tiled_shape); + + // Determine extent of K-dimension assigned to each block + gemm_k_size = args.problem_size.k(); + + if (args.mode == GemmUniversalMode::kGemm || args.mode == GemmUniversalMode::kGemmSplitKParallel) + { + int const kAlignK = const_max(const_max(128 / sizeof_bits::value, 128 / sizeof_bits::value), 1); + + gemm_k_size = round_up(ceil_div(args.problem_size.k(), args.batch_count), kAlignK); + if (gemm_k_size) { + grid_tiled_shape.k() = ceil_div(args.problem_size.k(), gemm_k_size); + } + } + } + + + /// Returns the workspace size (in bytes) needed for this problem geometry + size_t get_workspace_size() const + { + size_t workspace_bytes = 0; + if (mode == GemmUniversalMode::kGemmSplitKParallel) + { + // Split-K parallel always requires a temporary workspace + workspace_bytes = + sizeof(ElementC) * + size_t(batch_stride_D) * + size_t(grid_tiled_shape.k()); + } + else if (mode == GemmUniversalMode::kGemm && grid_tiled_shape.k() > 1) + { + // Serial split-K only requires a temporary workspace if the number of partitions along the + // GEMM K dimension is greater than one. + workspace_bytes = sizeof(int) * size_t(grid_tiled_shape.m()) * size_t(grid_tiled_shape.n()); + } + + return workspace_bytes; + } + + + /// Assign and initialize the specified workspace buffer. Assumes + /// the memory allocated to workspace is at least as large as get_workspace_size(). + Status init_workspace( + void *workspace, + mcStream_t stream = nullptr) + { + semaphore = static_cast(workspace); + // Zero-initialize entire workspace + if (semaphore) + { + size_t workspace_bytes = get_workspace_size(); + + MCTLASS_TRACE_HOST(" Initialize " << workspace_bytes << " workspace bytes"); + + mcError_t result = mcMemsetAsync( + semaphore, + 0, + workspace_bytes, + stream); + + if (result != mcSuccess) { + MCTLASS_TRACE_HOST(" mcMemsetAsync() returned error " << mcGetErrorString(result)); + return Status::kErrorInternal; + } + } + + return Status::kSuccess; + } + + + /// Returns the GEMM volume in thread block tiles + GemmCoord get_tiled_shape() const + { + return grid_tiled_shape; + } + + + /// Returns the total number of thread blocks to launch + int get_grid_blocks() const + { + dim3 grid_dims = get_grid_dims(); + return grid_dims.x * grid_dims.y * grid_dims.z; + } + + + /// Returns the grid extents in thread blocks to launch + dim3 get_grid_dims() const + { + return ThreadblockSwizzle().get_grid_shape(grid_tiled_shape); + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_grouped.h b/csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_grouped.h new file mode 100644 index 0000000..4fca205 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_grouped.h @@ -0,0 +1,697 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Grouped Rank2K kernel. +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/trace.h" +#include "mctlass/gemm/kernel/rank_2k_transpose_operands.h" +#include "mctlass/gemm/kernel/rank_2k_grouped_problem_visitor.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma1_, ///! Threadblock-scoped matrix multiply-accumulate (A*B^T) + typename Mma2_, ///! Threadblock-scoped matrix multiply-accumulate (B*A^T) + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + ComplexTransform OriginalTransformA_, ///! Public-facing transformation on A + ComplexTransform OriginalTransformB_, ///! Public-facing transformation on B + FillMode FillModeC_, ///! Fill Mode for C (kLower or kUpper) + BlasMode BlasMode_, ///! Blas3 computation mode + GroupScheduleMode GroupScheduleMode_, ///! Type of scheduling to perform + bool Transposed = false +> +struct Rank2KGrouped { +public: + + using Mma1 = Mma1_; + using Mma2 = Mma2_; + + static_assert(platform::is_same::value && + platform::is_same::value, + "Kernel-level grouped Rank2K requires that LayoutC be row major."); + + // Define generic Mma for usecases that use Kernel::Mma + using Mma = Mma1_; + + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static GroupScheduleMode const kGroupScheduleMode = GroupScheduleMode_; + static bool const kTransposed = Transposed; + + // Public-facing type definitions related to operand element type, layout, and complex conjugate + // operation. Must interact with the 'kTransposed' notion to reflect the original layout, + // fill mode, etc. passed in. + // + // Recall that a Rank2K operation performs (A x BT) + (B x AT) + // This is performed via: + // Mma1 = (A x BT) + // Mma2 = (B x AT) + // + // However, if C needs to be transposed, then this is changed to the following: + // Mma1 = (B x AT) + // Mma2 = (A x BT) + // + // The transformation above is achieved by swapping the Layouts/Elements/Transforms/etc. + // of A and B as they are passed into the instantiations of Mma1 and Mma2. + // + // Now, given access to only Mma1 and Mma2, as well as whether a transposition has occurred, + // we wish to retrieve the original Layouts/Elements/etc. for A and B that were passed into + // the device-level call. + // + // The logic to do this (which is made clearer by referencing the above instantiations) is as follows: + // LayoutA = kTransposed ? Mma2::LayoutA : Mma1::LayoutA + // LayoutB = kTransposed ? Mma1::LayoutA : Mma2::LayoutA + // + // We achieve this swapping by passing Mma1::*A and Mma2::*B to Rank2KMapArguments: + using MapArgumentsA = kernel::detail::Rank2KMapArguments< + typename Mma1::IteratorA::Element, + typename Mma1::IteratorA::Layout, + Mma1::kTransformA, + Mma1::IteratorA::AccessType::kElements, + typename Mma2::IteratorA::Element, + typename Mma2::IteratorA::Layout, + Mma2::kTransformA, + Mma2::IteratorA::AccessType::kElements, + typename Mma1::LayoutC, + FillModeC_, + kTransposed + >; + + using ElementA = typename MapArgumentsA::ElementA; + using LayoutA = typename MapArgumentsA::LayoutA; + static int const kAlignmentA = MapArgumentsA::kAlignmentA; + + using MapArgumentsB = kernel::detail::Rank2KMapArguments< + typename Mma2::IteratorA::Element, + typename Mma2::IteratorA::Layout, + Mma2::kTransformA, + Mma2::IteratorA::AccessType::kElements, + typename Mma1::IteratorA::Element, + typename Mma1::IteratorA::Layout, + Mma1::kTransformA, + Mma1::IteratorA::AccessType::kElements, + typename Mma2::LayoutC, + FillModeC_, + kTransposed + >; + + using ElementB = typename MapArgumentsB::ElementA; + using LayoutB = typename MapArgumentsB::LayoutA; + static int const kAlignmentB = MapArgumentsB::kAlignmentA; + + // Use the user-provided TransformA and TransformB, rather than those + // resulting from MapArguments, because Mma1 and Mma2 may have different + // complex transforms than those passed in by the user. + // (See kernel/rank_2k_complex.h for an example of this) + static mctlass::ComplexTransform const kTransformA = OriginalTransformA_; + static mctlass::ComplexTransform const kTransformB = OriginalTransformB_; + + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename MapArgumentsA::LayoutC; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + static FillMode const kFillModeC = MapArgumentsA::kFillModeC; + + // Common type definitions for Mma1 and Mma2 + using Operator = typename Mma1::Operator; + using OperatorClass = typename Mma1::Operator::OperatorClass; + using ThreadblockShape = typename Mma1::Shape; + using WarpShape = typename Mma1::Operator::Shape; + using InstructionShape = typename Mma1::Policy::Operator::InstructionShape; + using ArchTag = typename Mma1::ArchTag; + + static int const kStages = Mma1::kStages; + static BlasMode const kBlasMode = BlasMode_; + +private: + static FillMode const kInternalFillModeC = FillModeC_; + +public: + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma1::WarpCount; + // static int const kThreadCount = 32 * WarpCount::kCount; + static int const kThreadCount = 64 * WarpCount::kCount; + + using ProblemVisitor = Rank2KGroupedProblemVisitor< + ThreadblockShape, + kGroupScheduleMode, + kThreadCount, + kThreadCount, + kInternalFillModeC>; + + // + // Structures + // + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmUniversalMode mode; + GemmCoord *problem_sizes; + int problem_count; + int threadblock_count; + + typename EpilogueOutputOp::Params epilogue; + + ElementA ** ptr_A; + ElementB ** ptr_B; + ElementC ** ptr_C; + ElementC ** ptr_D; + + typename LayoutA::Stride::LongIndex *lda; + typename LayoutB::Stride::LongIndex *ldb; + typename LayoutC::Stride::LongIndex *ldc; + typename LayoutC::Stride::LongIndex *ldd; + + // Only used by device-level operator + GemmCoord *host_problem_sizes; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments(): + mode(GemmUniversalMode::kGemm), + problem_count(0), + threadblock_count(0), + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + lda(nullptr), + ldb(nullptr), + ldc(nullptr), + ldd(nullptr), + host_problem_sizes(nullptr) + { + + } + + /// Ctor + MCTLASS_HOST_DEVICE + Arguments( + GemmUniversalMode mode, + GemmCoord *problem_sizes, + int problem_count, + int threadblock_count, + typename EpilogueOutputOp::Params epilogue, + ElementA ** ptr_A, + ElementB ** ptr_B, + ElementC ** ptr_C, + ElementC ** ptr_D, + typename LayoutA::Stride::LongIndex *lda, + typename LayoutB::Stride::LongIndex *ldb, + typename LayoutC::Stride::LongIndex *ldc, + typename LayoutC::Stride::LongIndex *ldd, + GemmCoord *host_problem_sizes=nullptr + ): + mode(mode), + problem_sizes(problem_sizes), + problem_count(problem_count), + threadblock_count(threadblock_count), + epilogue(epilogue), + ptr_A(ptr_A), + ptr_B(ptr_B), + ptr_C(ptr_C), + ptr_D(ptr_D), + lda(lda), + ldb(ldb), + ldc(ldc), + ldd(ldd), + host_problem_sizes(host_problem_sizes) + { + + } + + }; + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params { + + typename ProblemVisitor::Params problem_visitor; + int threadblock_count; + + typename EpilogueOutputOp::Params output_op; + + GemmUniversalMode mode; + int batch_count; + + ElementA ** ptr_A; + ElementB ** ptr_B; + ElementC ** ptr_C; + ElementC ** ptr_D; + + typename LayoutA::Stride::LongIndex *lda; + typename LayoutB::Stride::LongIndex *ldb; + typename LayoutC::Stride::LongIndex *ldc; + typename LayoutC::Stride::LongIndex *ldd; + + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + mode(mctlass::gemm::GemmUniversalMode::kGemm), + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + lda(nullptr), + ldb(nullptr), + ldc(nullptr), + ldd(nullptr) + { } + + MCTLASS_HOST_DEVICE + Params(Arguments const &args, void *workspace = nullptr, int tile_count = 0): + problem_visitor(args.problem_sizes, args.problem_count, workspace, tile_count), + threadblock_count(args.threadblock_count), + output_op(args.epilogue), + ptr_A(args.ptr_A), + ptr_B(args.ptr_B), + ptr_C(args.ptr_C), + ptr_D(args.ptr_D), + lda(args.lda), + ldb(args.ldb), + ldc(args.ldc), + ldd(args.ldd) + { + + } + + MCTLASS_HOST_DEVICE + void update( + Arguments const &args, + void *workspace = nullptr, + int tile_count = 0) { + + problem_visitor = typename ProblemVisitor::Params(args.problem_sizes, args.problem_count, workspace, tile_count); + threadblock_count = args.threadblock_count; + output_op = args.output_op; + ptr_A = args.ptr_A; + ptr_B = args.ptr_B; + ptr_C = args.ptr_C; + ptr_D = args.ptr_D; + } + }; + + /// Shared memory storage structure + struct SharedStorage { + union { + typename Mma1::SharedStorage mma1_main_loop; + typename Mma2::SharedStorage mma2_main_loop; + typename Epilogue::SharedStorage epilogue; + } kernel; + + // ProblemVisitor shared storage can't be overlapped with others + typename ProblemVisitor::SharedStorage problem_visitor; + }; + +public: + + // + // Methods + // + + MCTLASS_DEVICE + Rank2KGrouped() { } + + /// Determines whether kernel satisfies alignment + static Status can_implement(mctlass::gemm::GemmCoord const & problem_size) { + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return Status::kSuccess; + } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // + // Problem visitor. + // + + ProblemVisitor problem_visitor( + params.problem_visitor, + shared_storage.problem_visitor, + blockIdx.x); + + // Outer 'persistent' loop to iterate over tiles + while (problem_visitor.next_tile()) { + + GemmCoord problem_size = problem_visitor.problem_size(); + int32_t problem_idx = problem_visitor.problem_index(); + int32_t threadblock_idx = int32_t(problem_visitor.threadblock_idx()); + + GemmCoord grid_shape = problem_visitor.grid_shape(problem_size); + + mctlass::gemm::GemmCoord threadblock_tile_offset = problem_visitor.threadblock_offset(threadblock_idx); + + + // Early exit if threadblock is out of range + if (grid_shape.m() <= threadblock_tile_offset.m() || + grid_shape.n() <= threadblock_tile_offset.n()) { + // Next tile + problem_visitor.advance(gridDim.x); + continue; + } + + // Skip this tile if Fill Mode is Lower and + // if the entire tile is above the main diagonal (bottom-left corner is at or above the diagonal) + if (kInternalFillModeC == mctlass::FillMode::kLower && + (threadblock_tile_offset.m() + 1) * Mma1::Shape::kM <= threadblock_tile_offset.n() * Mma1::Shape::kN) { + // Next tile + problem_visitor.advance(gridDim.x); + continue; + } + + // Skip this tile if Fill Mode is Upper and + // if the entire tile is below the main diagonal (top-right corner is at or below the diagonal) + if (kInternalFillModeC == mctlass::FillMode::kUpper && + threadblock_tile_offset.m() * Mma1::Shape::kM >= (threadblock_tile_offset.n() + 1) * Mma1::Shape::kN) { + // Next tile + problem_visitor.advance(gridDim.x); + continue; + } + + bool tile_on_diagonal = false; + // Mark tiles that are being crossed by the main diagonal + // (top-right and bottom-left corners are on either side of the diagonal) + if ((threadblock_tile_offset.m() + 1) * Mma1::Shape::kM > threadblock_tile_offset.n() * Mma1::Shape::kN + && threadblock_tile_offset.m() * Mma1::Shape::kM < (threadblock_tile_offset.n() + 1) * Mma1::Shape::kN) { + tile_on_diagonal = true; + } + + int offset_k = 0; + int problem_size_k = problem_size.k(); + + // + // Fetch pointers based on mode. + // + if (params.mode == GemmUniversalMode::kGemm || + params.mode == GemmUniversalMode::kGemmSplitKParallel) { + + if (threadblock_tile_offset.k() + 1 < grid_shape.k()) { + problem_size_k = (threadblock_tile_offset.k() + 1) * problem_size.k(); + } + + offset_k = threadblock_tile_offset.k() * problem_size.k(); + } + + ElementA *ptr_A = reinterpret_cast((kTransposed ? params.ptr_B[problem_idx] : params.ptr_A[problem_idx])); + typename LayoutA::Stride::LongIndex ldm_A = (kTransposed ? params.ldb[problem_idx] : params.lda[problem_idx]); + + ElementB *ptr_B = reinterpret_cast((kTransposed ? params.ptr_A[problem_idx] : params.ptr_B[problem_idx])); + typename LayoutB::Stride::LongIndex ldm_B = (kTransposed ? params.lda[problem_idx] : params.ldb[problem_idx]); + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_MxK{ + threadblock_tile_offset.m() * Mma1::Shape::kM, + offset_k, + }; + + mctlass::MatrixCoord tb_offset_KxN{ + offset_k, + threadblock_tile_offset.n() * Mma1::Shape::kN + }; + + // Assume identity swizzle + MatrixCoord tb_offset( + threadblock_tile_offset.m() * Mma1::Shape::kM, + threadblock_tile_offset.n() * Mma1::Shape::kN + ); + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands for Mma1 + typename Mma1::IteratorA iterator_A( + typename Mma1::IteratorA::Params(ldm_A), + ptr_A, + {problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_MxK); + + typename Mma1::IteratorB iterator_BT( + typename Mma1::IteratorB::Params(ldm_B), + ptr_B, + {problem_size_k, problem_size.n()}, + thread_idx, + tb_offset_KxN); + + // Construct iterators to A and B operands for Mma2 + typename Mma2::IteratorA iterator_B( + typename Mma2::IteratorA::Params(ldm_B), + ptr_B, + {problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_MxK); + + typename Mma2::IteratorB iterator_AT( + typename Mma2::IteratorB::Params(ldm_A), + ptr_A, + {problem_size_k, problem_size.n()}, + thread_idx, + tb_offset_KxN); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply for Mma1 (A x BT) + Mma1 mma1(shared_storage.kernel.mma1_main_loop, thread_idx, warp_idx, lane_idx); + + // Construct thread-scoped matrix multiply for Mma2 (B x AT) + Mma2 mma2(shared_storage.kernel.mma2_main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma1::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - offset_k + Mma1::Shape::kK - 1) / Mma1::Shape::kK; + + // Wait for all threads to finish their epilogue phases from the previous tile. + __syncthreads(); + + // Compute threadblock-scoped matrix multiply-add (A x BT) + mma1( + gemm_k_iterations, + accumulators, + iterator_A, + iterator_BT, + accumulators); + + // HER2K kernel needs Alpha to be complex and is conj(Alpha) is applied to the second HERK. + if (kBlasMode == BlasMode::kHermitian) { + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * grid_shape.m(); + + ElementC *ptr_C = static_cast(params.ptr_C[problem_idx]); + ElementC *ptr_D = static_cast(params.ptr_D[problem_idx]); + + // If TB not on diagonal, FillMode doesn't apply. + FillMode kFillModeTB = tile_on_diagonal ? kInternalFillModeC : FillMode::kNone; + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + typename Epilogue::OutputTileIterator::Params(params.ldc[problem_idx]), + ptr_C, + problem_size.mn(), + thread_idx, + tb_offset, + kFillModeTB + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + typename Epilogue::OutputTileIterator::Params(params.ldd[problem_idx]), + ptr_D, + problem_size.mn(), + thread_idx, + tb_offset, + kFillModeTB + ); + + Epilogue epilogue( + shared_storage.kernel.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Execute the epilogue operator to update the destination tensor. + epilogue( + output_op, + iterator_D, + accumulators, + iterator_C); + + __syncthreads(); + + accumulators.clear(); + } + + // Compute threadblock-scoped matrix multiply-add (B x AT) + mma2( + gemm_k_iterations, + accumulators, + iterator_B, + iterator_AT, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + /* Needed for HER2K where the second HERK is multiplied by conj(alpha) */ + typename EpilogueOutputOp::Params second_her2k_params(conj(params.output_op.alpha), 1); + EpilogueOutputOp output_op_her2k(second_her2k_params); + + // + // Masked tile iterators constructed from members + // + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * grid_shape.m(); + + ElementC *ptr_C = static_cast(params.ptr_C[problem_idx]); + + // HER2K kernel needs Alpha to be complex and is conj(Alpha) is applied to the second HERK. + if (kBlasMode == BlasMode::kHermitian) { + ptr_C = static_cast(params.ptr_D[problem_idx]); + } + + ElementC *ptr_D = static_cast(params.ptr_D[problem_idx]); + + // If TB not on diagonal, FillMode doesn't apply. + FillMode kFillModeTB = tile_on_diagonal ? kInternalFillModeC : FillMode::kNone; + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + typename Epilogue::OutputTileIterator::Params(params.ldc[problem_idx]), + ptr_C, + problem_size.mn(), + thread_idx, + tb_offset, + kFillModeTB + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + typename Epilogue::OutputTileIterator::Params(params.ldd[problem_idx]), + ptr_D, + problem_size.mn(), + thread_idx, + tb_offset, + kFillModeTB + ); + + Epilogue epilogue( + shared_storage.kernel.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Execute the epilogue operator to update the destination tensor. + if (kBlasMode == BlasMode::kSymmetric) { + epilogue( + output_op, + iterator_D, + accumulators, + iterator_C); + } else { + epilogue( + output_op_her2k, + iterator_D, + accumulators, + iterator_C); + } + + // Next tile + problem_visitor.advance(gridDim.x); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_grouped_problem_visitor.h b/csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_grouped_problem_visitor.h new file mode 100644 index 0000000..bfb9418 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_grouped_problem_visitor.h @@ -0,0 +1,376 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Problem visitor for grouped Rank2K operations. + + This problem visitor is specialized for Rank2K operations, for which matrix C is upper/lower + triangular. Using a problem visitor designed for GEMMs for Rank2K problems is inefficient + because threadblocks will be frequently assigned to tiles that exit early (e.g., due to + being assigned to a tile in the upper-triangular portion of a lower-triangular problem). + This can lead to load imbalance among threadblocks, as the GEMM-based scheduler + assigns all threadblocks to nearly the same number of tiles, regardless of whether + those tiles exit early. + + Consider an example of a group of four Rank2Ks with matrix C consisting of a grid of 2x2 tiles. + Consider a grid of 8 threadblocks. The default GEMM scheduler will assign threadblocks to + tiles in the following order: + Rank2K 0 Rank2K 1 Rank2K 2 Rank2K 3 + 0 1 4 5 0 1 4 5 + 2 3 6 7 2 3 6 7 + Assuming that the problems are lower triangular, blocks 1 and 5 are continuously assigned + to inactive tiles. + + This problem visitor aims to assign threadblocks to only those tiles which are in the + upper/lower triangular portion of a given problem. Using the example above, the resulting + assignment would be: + Rank2K 0 Rank2K 1 Rank2K 2 Rank2K 3 + 0 - 3 - 6 - 1 - + 1 2 4 5 7 0 2 3 + + Achieving the schedule above requires a mapping from threadblock ID to tile coordinates (i, j). + We will illustrate this by mapping on a lower-triangular matrix with a 3x3 grid. We first + calculate row and column indices assuming one-indexed rows, tiles, and threadblock IDs, and + then subtract one to convert to zero-indexed. + Col 1 Col 2 Col 3 + ---------------------- + Row 1 | 1 - - + Row 2 | 2 3 - + Row 3 | 4 5 6 + + We next outline this mapping, borrowing from: https://stackoverflow.com/a/40954159 + + Calculating row i given threadblock ID t + ---------------------------------------- + For a given row i, all threadblock IDs t in that row satisfy the following: + t <= 1 + 2 + 3 + ... + (i-1) + i + + The closed-form equation for the right-hand side is: i(i+1)/2. + Using this, we can solve for i given t: + t <= i(i+1)/2 + 2t <= i^2 + i + 2t <= i^2 + i + 0.25 - 0.25 + 2t + 0.25 <= i^2 + i + 0.25 + 2t + 0.25 <= (i + 0.5)^2 + sqrt(2t + 0.25) - 0.5 <= i + + To account for fractional values, we set: + i = ceil(sqrt(2t + 0.25) - 0.5) + + To turn this into a zero-indexed row and work with zero-indexed t, we perform: + i = ceil(sqrt(2(t+1) + 0.25) - 0.5) - 1 + = ceil(sqrt(2t + 2.25) - 0.5) - 1 + + Calculating column j given threadblock ID t and row i + ----------------------------------------------------- + For a given row i, all threadblock IDs t in that row also satisfy the following: + t > 1 + 2 + 3 + ... + (i-2) + (i-1) + --> t > i(i-1)/2 + + Threadblock IDs within a given row are sequential, so the one-indexed column ID + for one-indexed threadblock ID t and row i is: + j = t - (i(i-1)/2) + + The zero-indexed version becomes: + j = (t+1) - (i(i+1)/2) -1 + = t - (i(i+1)/2) + + Accounting for non-square grids + ------------------------------- + Though the overall output problem size for Rank2K problems is guranteed to be square, the + grids used in computing may not be square due to using non-square threadblock shapes. For + example, a threadblock shape of 64x32 operating on a problem of output size 128x128 would + result in a grid of 2x4 tiles. + + This case can be handled by noting that the output resembles a square grid of 2x2 "macro tiles" + each of which contains 2 "true tiles." We can thus first map a threadblock ID to its "macro tile" + using the equations above, and then map it to the "true tile" within its "macro tile." In the example + of a 2x4 grid, this mapping would look as follows: + "Macro grid" "True grid" + {0, 1} - 0 1 - - + {2, 3} {4, 5} 2 3 4 5 + + A zero-indexed threadblock ID t is mapped to its "macro tile ID" t_macro as: + t_macro = t // r + Where r is the ratio of the maximum dimension of the grid to the minimum dimension of the grid + (i.e., r = 4 / 2 = 2 in the previous example). + + One uses t_macro and the calculations above to find the row and column in the square matrix to + obtain i_macro and j_macro (zero-indexed). The mapping from (i_macro, j_macro) --> (i, j) + is simply the following: + if (ThreadblockShape::M > ThreadblockShape::N): + r = ThreadblockShape::M / ThreadblockShape::N + i = i_macro + j = (j_macro * r) + (t % r) + elif (ThreadblockShape::M < ThreadblockShape::N): + r = ThreadblockShape::N / ThreadblockShape::M + i = (i_macro * r) + (t % r) + j = j_macro + else: + i = i_macro + j = j_macro + + Handling cases with grid dimensions that aren't multiples of eachother + ---------------------------------------------------------------------- + Even though threadblock shapes M and N are typically multiples of one another, the grid + for a given problem may not have dimensions of the same ratio as that of the threadblock. + For example, a problem of size 132x132 using a threadblock of shape 64x32 will result + in a grid of 3x5 tiles. In this case, there is not an integer number of "true tiles" + per "macro tile." + + When this scenario arises, we simply pad the larger dimension of the grid such that + there are an integer number of "true tiles" per "macro tile." Thus, the 3x5 grid in + the example above will be treated as a 3x6 grid. Row and column positions for each + tile are calculated as above. Any threadblocks that map to tiles that are outside the + problem range or upper/lower triangular portion (e.g., (2, 5)) will exit early from + this problem and may proceed to the next problem in the group. + + Handling upper-triangular matrices + ---------------------------------- + The only modification needed for upper-triangular matrices is to swap i_macro and j_macro + in the calculations above. +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" + +#include "mctlass/gemm/kernel/grouped_problem_visitor.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +namespace detail { +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Helpers for calculating offsets for Rank2K problem visitor. These helpers specifically pertain +// to the conversion from "macro tiles" to "true tiles" in the description above. +// +template < + typename ThreadblockShape, + typename Enable = void +> +struct Rank2KGroupedProblemVisitorOffsetHelper; + +// Partial specialization for the case where threadblock shape M > threadblock shape N +template < + typename ThreadblockShape +> +struct Rank2KGroupedProblemVisitorOffsetHelper< + ThreadblockShape, + typename platform::enable_if< (ThreadblockShape::kM > ThreadblockShape::kN) >::type +> { + static_assert(ThreadblockShape::kM % ThreadblockShape::kN == 0, + "Rank2KGroupedProblemVisitor with threadblock shape M > threadblock shape N " + "requires that threadblock shape M be a multiple of threadblock shape N."); + + static int32_t const kThreadblockSkewRatio = ThreadblockShape::kM / ThreadblockShape::kN; + + MCTLASS_HOST_DEVICE + static int32_t min_dim(mctlass::gemm::GemmCoord grid) { + return grid.m(); + } + + MCTLASS_HOST_DEVICE + static int32_t macro_row_to_row(int32_t row, int32_t threadblock_id) { + return row; + } + + MCTLASS_HOST_DEVICE + static int32_t macro_col_to_col(int32_t col, int32_t threadblock_id) { + return (col * kThreadblockSkewRatio) + (threadblock_id % kThreadblockSkewRatio); + } +}; + +// Partial specialization for the case where threadblock shape M < threadblock shape N +template < + typename ThreadblockShape +> +struct Rank2KGroupedProblemVisitorOffsetHelper< + ThreadblockShape, + typename platform::enable_if< (ThreadblockShape::kM < ThreadblockShape::kN) >::type +> { + + static_assert(ThreadblockShape::kN % ThreadblockShape::kM == 0, + "Rank2KGroupedProblemVisitor with threadblock shape M < threadblock shape N " + "requires that threadblock shape N be a multiple of threadblock shape M."); + + static int32_t const kThreadblockSkewRatio = ThreadblockShape::kN / ThreadblockShape::kM; + + MCTLASS_HOST_DEVICE + static int32_t min_dim(mctlass::gemm::GemmCoord grid) { + return grid.n(); + } + + MCTLASS_HOST_DEVICE + static int32_t macro_row_to_row(int32_t row, int32_t threadblock_id) { + return (row * kThreadblockSkewRatio) + (threadblock_id % kThreadblockSkewRatio); + } + + MCTLASS_HOST_DEVICE + static int32_t macro_col_to_col(int32_t col, int32_t threadblock_id) { + return col; + } +}; + +// Partial specialization for the case where threadblock shape M == threadblock shape N +// In this case, macro tiles are equivalent to true tiles, so the conversions are +// identity functions. +template < + typename ThreadblockShape +> +struct Rank2KGroupedProblemVisitorOffsetHelper< + ThreadblockShape, + typename platform::enable_if< (ThreadblockShape::kM == ThreadblockShape::kN) >::type +> { + + static int32_t const kThreadblockSkewRatio = 1; + + MCTLASS_HOST_DEVICE + static int32_t min_dim(mctlass::gemm::GemmCoord grid) { + return grid.m(); + } + + MCTLASS_HOST_DEVICE + static int32_t macro_row_to_row(int32_t row, int32_t threadblock_id) { + return row; + } + + MCTLASS_HOST_DEVICE + static int32_t macro_col_to_col(int32_t col, int32_t threadblock_id) { + return col; + } +}; + +// Helper for correctly representing problem sizes in grouped kernels +template +struct Rank2KGroupedProblemSizeHelper { + using OffsetHelper = Rank2KGroupedProblemVisitorOffsetHelper; + + MCTLASS_HOST_DEVICE + static mctlass::gemm::GemmCoord grid_shape(const mctlass::gemm::GemmCoord& problem) { + return mctlass::gemm::GemmCoord( + ((problem.m() - 1 + ThreadblockShape::kM) / ThreadblockShape::kM), + ((problem.n() - 1 + ThreadblockShape::kN) / ThreadblockShape::kN), + 1); + } + + MCTLASS_HOST_DEVICE + static int32_t tile_count(const mctlass::gemm::GemmCoord& grid) { + // Return the number of tiles at or below the diagonal (or at and above + // for mode kUpper). We do this by first calculating this value assuming + // we have a square matrix of tiles of size `dim x dim` where `dim` is the + // minimum among {grid.m(), grid.n()}. We then multiply the resulting value + // by OffsetHelper::kThreadblockSkewRatio to account for cases in which there + // are more tiles in one dimension than the other. + int32_t dim = OffsetHelper::min_dim(grid); + int32_t tiles_on_diagonal = dim; + int32_t tiles_below_diagonal = ((dim * (dim - 1)) / 2); + return (tiles_on_diagonal + tiles_below_diagonal) * OffsetHelper::kThreadblockSkewRatio; + } + + MCTLASS_HOST_DEVICE + static void possibly_transpose_problem(mctlass::gemm::GemmCoord& problem) {} +}; + +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Default problem visitor for fill modes kUpper and kLower. +// +template +struct Rank2KGroupedProblemVisitor : public GroupedProblemVisitor< + detail::Rank2KGroupedProblemSizeHelper, + ThreadblockShape, + GroupScheduleMode_, + PrefetchTileCount, + ThreadCount> { + + static mctlass::FillMode const kFillModeC = FillModeC; + + static_assert(kFillModeC == mctlass::FillMode::kLower || kFillModeC == mctlass::FillMode::kUpper, + "Default Rank2KGroupedProblemVisitor requires fill mode of kLower or kUpper."); + + using ProblemSizeHelper = detail::Rank2KGroupedProblemSizeHelper; + using Base = GroupedProblemVisitor; + using OffsetHelper = typename ProblemSizeHelper::OffsetHelper; + using Params = typename Base::Params; + using SharedStorage = typename Base::SharedStorage; + + // + // Methods + // + MCTLASS_DEVICE + Rank2KGroupedProblemVisitor( + Params const ¶ms_, + SharedStorage &shared_storage_, + int32_t block_idx + ): Base(params_, shared_storage_, block_idx) + {} + + MCTLASS_DEVICE + mctlass::gemm::GemmCoord threadblock_offset(int32_t threadblock_id) const { + int32_t macro_id = threadblock_id / OffsetHelper::kThreadblockSkewRatio; + int32_t macro_row = ceil(mctlass::fast_sqrt((2*macro_id) + 2.25) - 0.5) - 1; + int32_t macro_col = macro_id - (((macro_row+1) * macro_row)/2); + + if (kFillModeC == mctlass::FillMode::kUpper) { + swap(macro_row, macro_col); + } + + int32_t row = OffsetHelper::macro_row_to_row(macro_row, threadblock_id); + int32_t col = OffsetHelper::macro_col_to_col(macro_col, threadblock_id); + + return mctlass::gemm::GemmCoord(row, col, 0); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_transpose_operands.h b/csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_transpose_operands.h new file mode 100644 index 0000000..254d2d1 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_transpose_operands.h @@ -0,0 +1,129 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Transpositions for Rank2K problems. +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementA_, + typename LayoutA_, + ComplexTransform TransformA, + int AlignmentA, + typename ElementB_, + typename LayoutB_, + ComplexTransform TransformB, + int AlignmentB, + typename LayoutC_, + FillMode FillModeC_, + bool Transpose +> +struct Rank2KMapArguments { + using ElementA = ElementA_; + using LayoutA = LayoutA_; + static ComplexTransform const kTransformA = TransformA; + static int const kAlignmentA = AlignmentA; + using ElementB = ElementB_; + using LayoutB = LayoutB_; + static ComplexTransform const kTransformB = TransformB; + static int const kAlignmentB = AlignmentB; + using LayoutC = LayoutC_; + static FillMode const kFillModeC = FillModeC_; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ElementA_, + typename LayoutA_, + ComplexTransform TransformA, + int AlignmentA, + typename ElementB_, + typename LayoutB_, + ComplexTransform TransformB, + int AlignmentB, + typename LayoutC_, + FillMode FillModeC_ +> +struct Rank2KMapArguments< + ElementA_, + LayoutA_, + TransformA, + AlignmentA, + ElementB_, + LayoutB_, + TransformB, + AlignmentB, + LayoutC_, + FillModeC_, + true +> { + using ElementA = ElementB_; + using LayoutA = LayoutB_; + static ComplexTransform const kTransformA = TransformB; + static int const kAlignmentA = AlignmentB; + using ElementB = ElementA_; + using LayoutB = LayoutA_; + static ComplexTransform const kTransformB = TransformA; + static int const kAlignmentB = AlignmentA; + using LayoutC = typename layout::LayoutTranspose::type; + static FillMode const kFillModeC = InvertFillMode::mode; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} +} +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_universal.h b/csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_universal.h new file mode 100644 index 0000000..3fe57d7 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/rank_2k_universal.h @@ -0,0 +1,779 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/semaphore.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma1_, ///! Threadblock-scoped matrix multiply-accumulate (A*B^T) + typename Mma2_, ///! Threadblock-scoped matrix multiply-accumulate (B*A^T) + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + FillMode FillModeC_, ///! Fill Mode for C (kLower or kUpper) + BlasMode BlasMode_ ///! Blas3 computation mode +> +struct Rank2KUniversal { +public: + + using Mma1 = Mma1_; + using Mma2 = Mma2_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma1::IteratorA::Element; + using ElementB = typename Mma1::IteratorB::Element; + + // Mma1 (A x B^T) + using LayoutA = typename Mma1::IteratorA::Layout; + using LayoutBT = typename Mma1::IteratorB::Layout; + static ComplexTransform const kMma1TransformA = Mma1::kTransformA; + static ComplexTransform const kMma1TransformB = Mma1::kTransformB; + + // Mma2 (B x A^T) + using LayoutB = typename Mma2::IteratorA::Layout; + using LayoutAT = typename Mma2::IteratorB::Layout; + static ComplexTransform const kMma2TransformA = Mma2::kTransformA; + static ComplexTransform const kMma2TransformB = Mma2::kTransformB; + + // Common type definitions for Mma1 and Mma2 + using Operator = typename Mma1::Operator; + using OperatorClass = typename Mma1::Operator::OperatorClass; + using ThreadblockShape = typename Mma1::Shape; + using WarpShape = typename Mma1::Operator::Shape; + using InstructionShape = typename Mma1::Policy::Operator::InstructionShape; + using ArchTag = typename Mma1::ArchTag; + + static int const kStages = Mma1::kStages; + static int const kAlignmentA = Mma1::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma1::IteratorB::AccessType::kElements; + + // Output related typedefinitions + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + static FillMode const kFillModeC = FillModeC_; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + static BlasMode const kBlasMode = BlasMode_; + + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma1::WarpCount; + // static int const kThreadCount = 32 * WarpCount::kCount; + static int const kThreadCount = 64 * WarpCount::kCount; + + + // + // Structures + // + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmUniversalMode mode; + GemmCoord problem_size; + int batch_count; + + typename EpilogueOutputOp::Params epilogue; + + void const * ptr_A; + void const * ptr_B; + void const * ptr_C; + void * ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + int64_t batch_stride_D; + + typename LayoutA::Stride::Index lda; + typename LayoutB::Stride::Index ldb; + typename LayoutC::Stride::Index ldc; + typename LayoutC::Stride::Index ldd; + + // + // Methods + // + + Arguments(): + mode(GemmUniversalMode::kGemm), + batch_count(1), + ptr_A(nullptr), ptr_B(nullptr), ptr_C(nullptr), ptr_D(nullptr) { } + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void const * ptr_C, + void * ptr_D, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C, + int64_t batch_stride_D, + typename LayoutA::Stride::Index lda, + typename LayoutB::Stride::Index ldb, + typename LayoutC::Stride::Index ldc, + typename LayoutC::Stride::Index ldd + ): + mode(mode), + problem_size(problem_size), + batch_count(batch_count), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), + batch_stride_A(batch_stride_A), batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D), + lda(lda), ldb(ldb), ldc(ldc), ldd(ldd) { + + } + + /// Returns arguments for a the transposed problem + Arguments transposed_problem() const { + Arguments args(*this); + + std::swap(args.ptr_A, args.ptr_B); + std::swap(args.lda, args.ldb); + std::swap(args.batch_stride_A, args.batch_stride_B); + + return args; + } + + }; + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params { + + mctlass::gemm::GemmCoord problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + + // Mma1 Iterator A and B params + typename Mma1::IteratorA::Params params_A; + typename Mma1::IteratorB::Params params_BT; + + // Mma2 Iterator A and B params + typename Mma2::IteratorA::Params params_B; + typename Mma2::IteratorB::Params params_AT; + + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::Params params_D; + + typename EpilogueOutputOp::Params output_op; + + GemmUniversalMode mode; + int batch_count; + int gemm_k_size; + + void * ptr_A; + void * ptr_B; + void * ptr_C; + void * ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + int64_t batch_stride_D; + + int *semaphore; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + swizzle_log_tile(0), + params_A(0), + params_BT(0), + params_B(0), + params_AT(0), + params_C(0), + params_D(0), + batch_count(0), + gemm_k_size(0), + mode(mctlass::gemm::GemmUniversalMode::kGemm), + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + batch_stride_A(0), + batch_stride_B(0), + batch_stride_C(0), + batch_stride_D(0), + semaphore(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + Arguments const &args, + mctlass::gemm::GemmCoord const & grid_tiled_shape, + int gemm_k_size, + void *workspace = nullptr + ): + problem_size(args.problem_size), + grid_tiled_shape(grid_tiled_shape), + swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)), + params_A(args.lda), + params_BT(args.ldb), + params_B(args.ldb), + params_AT(args.lda), + params_C(args.ldc), + params_D(args.ldd), + output_op(args.epilogue), + mode(args.mode), + batch_count(args.batch_count), + gemm_k_size(gemm_k_size), + ptr_A(const_cast(args.ptr_A)), + ptr_B(const_cast(args.ptr_B)), + ptr_C(const_cast(args.ptr_C)), + ptr_D(const_cast(args.ptr_D)), + batch_stride_A(args.batch_stride_A), + batch_stride_B(args.batch_stride_B), + batch_stride_C(args.batch_stride_C), + batch_stride_D(args.batch_stride_D), + semaphore(static_cast(workspace)) { + } + + MCTLASS_HOST_DEVICE + void update( + Arguments const &args, + void *workspace = nullptr) { + + ptr_A = const_cast(args.ptr_A); + ptr_B = const_cast(args.ptr_B); + ptr_C = const_cast(args.ptr_C); + ptr_D = args.ptr_D; + + output_op = args.epilogue; + + semaphore = static_cast(workspace); + } + + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma1::SharedStorage mma1_main_loop; + typename Mma2::SharedStorage mma2_main_loop; + typename Epilogue::SharedStorage epilogue; + }; + +public: + + // + // Methods + // + + MCTLASS_DEVICE + Rank2KUniversal() { } + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size) { + + static int const kAlignmentA = Mma1::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma1::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + if ((problem_size.m() % kAlignmentA) || (problem_size.k() % kAlignmentA) || + (problem_size.n() % kAlignmentB) || (problem_size.k() % kAlignmentB) || + (problem_size.m() % kAlignmentC) || (problem_size.n() % kAlignmentC)) { + + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + return; + } + + // Early exit if Fill Mode is Lower and + // if the entire tile is above the main diagonal (bottom-left corner is at or above the diagonal) + if (kFillModeC == mctlass::FillMode::kLower && + (threadblock_tile_offset.m() + 1) * Mma1::Shape::kM <= threadblock_tile_offset.n() * Mma1::Shape::kN) { + return; + } + + // Early exit if Fill Mode is Upper and + // if the entire tile is below the main diagonal (top-right corner is at or below the diagonal) + if (kFillModeC == mctlass::FillMode::kUpper && + threadblock_tile_offset.m() * Mma1::Shape::kM >= (threadblock_tile_offset.n() + 1) * Mma1::Shape::kN) { + return; + } + + bool tile_on_diagonal = false; + // Mark tiles that are being crossed by the main diagonal + // (top-right and bottom-left corners are on either side of the diagonal) + if ((threadblock_tile_offset.m() + 1) * Mma1::Shape::kM > threadblock_tile_offset.n() * Mma1::Shape::kN + && threadblock_tile_offset.m() * Mma1::Shape::kM < (threadblock_tile_offset.n() + 1) * Mma1::Shape::kN) { + tile_on_diagonal = true; + } + + int offset_k = 0; + int problem_size_k = params.problem_size.k(); + + ElementA *ptr_A = static_cast(params.ptr_A); + ElementB *ptr_B = static_cast(params.ptr_B); + + // + // Fetch pointers based on mode. + // + if (params.mode == GemmUniversalMode::kGemm || + params.mode == GemmUniversalMode::kGemmSplitKParallel) { + + if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) { + + problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size; + } + + offset_k = threadblock_tile_offset.k() * params.gemm_k_size; + } + + __syncthreads(); + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_MxK{ + threadblock_tile_offset.m() * Mma1::Shape::kM, + offset_k, + }; + + mctlass::MatrixCoord tb_offset_KxN{ + offset_k, + threadblock_tile_offset.n() * Mma1::Shape::kN + }; + + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands for Mma1 + typename Mma1::IteratorA iterator_A( + params.params_A, + ptr_A, + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_MxK); + + typename Mma1::IteratorB iterator_BT( + params.params_BT, + ptr_B, + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_KxN); + + // Construct iterators to A and B operands for Mma2 + typename Mma2::IteratorA iterator_B( + params.params_B, + ptr_B, + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_MxK); + + typename Mma2::IteratorB iterator_AT( + params.params_AT, + ptr_A, + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_KxN); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply for Mma1 (A x BT) + Mma1 mma1(shared_storage.mma1_main_loop, thread_idx, warp_idx, lane_idx); + + // Construct thread-scoped matrix multiply for Mma2 (B x AT) + Mma2 mma2(shared_storage.mma2_main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma1::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - offset_k + Mma1::Shape::kK - 1) / Mma1::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add (A x BT) + mma1( + gemm_k_iterations, + accumulators, + iterator_A, + iterator_BT, + accumulators); + + // HER2K kernel needs Alpha to be complex and is conj(Alpha) is applied to the second HERK. + if (kBlasMode == BlasMode::kHermitian) { + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma1::Shape::kM, + threadblock_tile_offset.n() * Mma1::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + ElementC *ptr_C = static_cast(params.ptr_C); + ElementC *ptr_D = static_cast(params.ptr_D); + + // + // Fetch pointers based on mode. + // + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + if (params.mode == GemmUniversalMode::kGemm) { + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + } + else if (params.mode == GemmUniversalMode::kGemmSplitKParallel) { + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_C += threadblock_tile_offset.k() * params.batch_stride_C; + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_C = static_cast(params.ptr_C)[threadblock_tile_offset.k()]; + ptr_D = static_cast(params.ptr_D)[threadblock_tile_offset.k()]; + } + + + // If CTA not on diagonal, FillMode doesn't apply. + FillMode kFillModeCTA = tile_on_diagonal ? kFillModeC : FillMode::kNone; + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + ptr_C, + params.problem_size.mn(), + thread_idx, + threadblock_offset, + kFillModeCTA + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.problem_size.mn(), + thread_idx, + threadblock_offset, + kFillModeCTA + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + + __threadfence(); + } + + // Execute the epilogue operator to update the destination tensor. + epilogue( + output_op, + iterator_D, + accumulators, + iterator_C); + + // + // Release the semaphore + // + + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + + __syncthreads(); + + accumulators.clear(); + } + + // Compute threadblock-scoped matrix multiply-add (B x AT) + mma2( + gemm_k_iterations, + accumulators, + iterator_B, + iterator_AT, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + /* Needed for HER2K where the second HERK is multiplied by conj(alpha) */ + typename EpilogueOutputOp::Params second_her2k_params(conj(params.output_op.alpha), 1); + EpilogueOutputOp output_op_her2k(second_her2k_params); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma1::Shape::kM, + threadblock_tile_offset.n() * Mma1::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + ElementC *ptr_C = static_cast(params.ptr_C); + + // HER2K kernel needs Alpha to be complex and is conj(Alpha) is applied to the second HERK. + if (kBlasMode == BlasMode::kHermitian) { + ptr_C = static_cast(params.ptr_D); + } + + ElementC *ptr_D = static_cast(params.ptr_D); + + // + // Fetch pointers based on mode. + // + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + if (params.mode == GemmUniversalMode::kGemm) { + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + if (kBlasMode == BlasMode::kSymmetric) { + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } else { + output_op_her2k.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + } + } + else if (params.mode == GemmUniversalMode::kGemmSplitKParallel) { + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_C += threadblock_tile_offset.k() * params.batch_stride_C; + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_C = static_cast(params.ptr_C)[threadblock_tile_offset.k()]; + ptr_D = static_cast(params.ptr_D)[threadblock_tile_offset.k()]; + } + + + // If CTA not on diagonal, FillMode doesn't apply. + FillMode kFillModeCTA = tile_on_diagonal ? kFillModeC : FillMode::kNone; + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + ptr_C, + params.problem_size.mn(), + thread_idx, + threadblock_offset, + kFillModeCTA + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.problem_size.mn(), + thread_idx, + threadblock_offset, + kFillModeCTA + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + + __threadfence(); + } + + // Execute the epilogue operator to update the destination tensor. + if (kBlasMode == BlasMode::kSymmetric) { + epilogue( + output_op, + iterator_D, + accumulators, + iterator_C); + } else { + epilogue( + output_op_her2k, + iterator_D, + accumulators, + iterator_C); + } + + // + // Release the semaphore + // + + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/rank_k_universal.h b/csrc/mctlass/include/mctlass/gemm/kernel/rank_k_universal.h new file mode 100644 index 0000000..4a0665e --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/rank_k_universal.h @@ -0,0 +1,566 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/semaphore.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + FillMode FillModeC_ ///! Fill Mode for C (kLower or kUpper) +> +struct RankKUniversal { +public: + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + static FillMode const kFillModeC = FillModeC_; + + static ComplexTransform const kTransformA = Mma::kTransformA; + static ComplexTransform const kTransformB = Mma::kTransformB; + using Operator = typename Mma::Operator; + + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::InstructionShape; + using ArchTag = typename Mma::ArchTag; + + static int const kStages = Mma::kStages; + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + // static int const kThreadCount = 32 * WarpCount::kCount; + static int const kThreadCount = 64 * WarpCount::kCount; + + /// Split-K preserves splits that are 128b aligned + static int const kSplitKAlignment = 128 / sizeof_bits::value; + + // + // Structures + // + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmUniversalMode mode; + GemmCoord problem_size; + int batch_count; + + typename EpilogueOutputOp::Params epilogue; + + void const * ptr_A; + void const * ptr_C; + void * ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_C; + int64_t batch_stride_D; + + typename LayoutA::Stride::Index lda; + typename LayoutB::Stride::Index ldb; + typename LayoutC::Stride::Index ldc; + typename LayoutC::Stride::Index ldd; + + // + // Methods + // + + Arguments(): + mode(GemmUniversalMode::kGemm), + batch_count(1), + ptr_A(nullptr), ptr_C(nullptr), ptr_D(nullptr) { } + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_C, + void * ptr_D, + int64_t batch_stride_A, + int64_t batch_stride_C, + int64_t batch_stride_D, + typename LayoutA::Stride::Index lda, + typename LayoutC::Stride::Index ldc, + typename LayoutC::Stride::Index ldd + ): + mode(mode), + problem_size(problem_size), + batch_count(batch_count), + epilogue(epilogue), + ptr_A(ptr_A), ptr_C(ptr_C), ptr_D(ptr_D), + batch_stride_A(batch_stride_A), batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D), + lda(lda), ldb(ldb), ldc(ldc), ldd(ldd) { + + } + + }; + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params { + + mctlass::gemm::GemmCoord problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorB::Params params_B; + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::Params params_D; + + typename EpilogueOutputOp::Params output_op; + + GemmUniversalMode mode; + int batch_count; + int gemm_k_size; + + void * ptr_A; + void * ptr_B; + void * ptr_C; + void * ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + int64_t batch_stride_D; + + int *semaphore; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + swizzle_log_tile(0), + params_A(0), + params_B(0), + params_C(0), + params_D(0), + batch_count(0), + gemm_k_size(0), + mode(mctlass::gemm::GemmUniversalMode::kGemm), + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + batch_stride_A(0), + batch_stride_B(0), + batch_stride_C(0), + batch_stride_D(0), + semaphore(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + Arguments const &args, + mctlass::gemm::GemmCoord const & grid_tiled_shape, + int gemm_k_size, + void *workspace = nullptr + ): + problem_size(args.problem_size), + grid_tiled_shape(grid_tiled_shape), + swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)), + params_A(args.lda), + params_B(args.lda), + params_C(args.ldc), + params_D(args.ldd), + output_op(args.epilogue), + mode(args.mode), + batch_count(args.batch_count), + gemm_k_size(gemm_k_size), + ptr_A(const_cast(args.ptr_A)), + ptr_B(const_cast(args.ptr_A)), + ptr_C(const_cast(args.ptr_C)), + ptr_D(const_cast(args.ptr_D)), + batch_stride_A(args.batch_stride_A), + batch_stride_B(args.batch_stride_A), + batch_stride_C(args.batch_stride_C), + batch_stride_D(args.batch_stride_D), + semaphore(static_cast(workspace)) { + } + + MCTLASS_HOST_DEVICE + void update( + Arguments const &args, + void *workspace = nullptr) { + + ptr_A = const_cast(args.ptr_A); + ptr_B = const_cast(args.ptr_A); + ptr_C = const_cast(args.ptr_C); + ptr_D = args.ptr_D; + + output_op = args.epilogue; + + semaphore = static_cast(workspace); + } + + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + +public: + + // + // Methods + // + + MCTLASS_DEVICE + RankKUniversal() { } + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size) { + + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + if ((problem_size.m() % kAlignmentA) || (problem_size.k() % kAlignmentA) || + (problem_size.n() % kAlignmentB) || (problem_size.k() % kAlignmentB) || + (problem_size.m() % kAlignmentC) || (problem_size.n() % kAlignmentC)) { + + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + return; + } + + // Early exit if Fill Mode is Lower and + // if the entire tile is above the main diagonal (bottom-left corner is at or above the diagonal) + if (kFillModeC == mctlass::FillMode::kLower && + (threadblock_tile_offset.m() + 1) * Mma::Shape::kM <= threadblock_tile_offset.n() * Mma::Shape::kN) { + return; + } + + // Early exit if Fill Mode is Upper and + // if the entire tile is below the main diagonal (top-right corner is at or below the diagonal) + if (kFillModeC == mctlass::FillMode::kUpper && + threadblock_tile_offset.m() * Mma::Shape::kM >= (threadblock_tile_offset.n() + 1) * Mma::Shape::kN) { + return; + } + + bool tile_on_diagonal = false; + // Mark tiles that are being crossed by the main diagonal + // (top-right and bottom-left corners are on either side of the diagonal) + if ((threadblock_tile_offset.m() + 1) * Mma::Shape::kM > threadblock_tile_offset.n() * Mma::Shape::kN + && threadblock_tile_offset.m() * Mma::Shape::kM < (threadblock_tile_offset.n() + 1) * Mma::Shape::kN) { + tile_on_diagonal = true; + } + + int offset_k = 0; + int problem_size_k = params.problem_size.k(); + + ElementA *ptr_A = static_cast(params.ptr_A); + ElementB *ptr_B = static_cast(params.ptr_B); + + // + // Fetch pointers based on mode. + // + if (params.mode == GemmUniversalMode::kGemm || + params.mode == GemmUniversalMode::kGemmSplitKParallel) { + + if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) { + + problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size; + } + + offset_k = threadblock_tile_offset.k() * params.gemm_k_size; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_A += threadblock_tile_offset.k() * params.batch_stride_A; + ptr_B += threadblock_tile_offset.k() * params.batch_stride_B; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_A = static_cast(params.ptr_A)[threadblock_tile_offset.k()]; + ptr_B = static_cast(params.ptr_B)[threadblock_tile_offset.k()]; + } + + __syncthreads(); + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + offset_k, + }; + + mctlass::MatrixCoord tb_offset_B{ + offset_k, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, + ptr_A, + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + params.params_B, + ptr_B, + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute threadblock-scoped matrix multiply-add + mma( + gemm_k_iterations, + accumulators, + iterator_A, + iterator_B, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + ElementC *ptr_C = static_cast(params.ptr_C); + ElementC *ptr_D = static_cast(params.ptr_D); + + // + // Fetch pointers based on mode. + // + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + if (params.mode == GemmUniversalMode::kGemm) { + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + } + else if (params.mode == GemmUniversalMode::kGemmSplitKParallel) { + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_C += threadblock_tile_offset.k() * params.batch_stride_C; + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_C = static_cast(params.ptr_C)[threadblock_tile_offset.k()]; + ptr_D = static_cast(params.ptr_D)[threadblock_tile_offset.k()]; + } + + + // If CTA not on diagonal, FillMode doesn't apply. + FillMode kFillModeCTA = tile_on_diagonal ? kFillModeC : FillMode::kNone; + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + ptr_C, + params.problem_size.mn(), + thread_idx, + threadblock_offset, + kFillModeCTA + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.problem_size.mn(), + thread_idx, + threadblock_offset, + kFillModeCTA + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + + __threadfence(); + } + + // Execute the epilogue operator to update the destination tensor. + epilogue( + output_op, + iterator_D, + accumulators, + iterator_C); + + // + // Release the semaphore + // + + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/sm70_gemm.hpp b/csrc/mctlass/include/mctlass/gemm/kernel/sm70_gemm.hpp new file mode 100644 index 0000000..27046e9 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/sm70_gemm.hpp @@ -0,0 +1,250 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/kernel_hardware_info.hpp" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/dispatch_policy.hpp" + +#include "cute/tensor.hpp" + +namespace mctlass::gemm::kernel { + +/////////////////////////////////////////////////////////////////////////////// + +template < + class ProblemShape_, + class CollectiveMainloop_, + class CollectiveEpilogue_, + class GridSwizzle_ +> +class GemmUniversal< + ProblemShape_, + CollectiveMainloop_, + CollectiveEpilogue_, + GridSwizzle_, + cute::enable_if_t>> +{ +public: + // + // Type Aliases + // + using ProblemShape = ProblemShape_; + using GridSwizzle = GridSwizzle_; + static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4, + "ProblemShape{} should be or "); + + // Mainloop derived types + using CollectiveMainloop = CollectiveMainloop_; + using TileShape = typename CollectiveMainloop::TileShape; + using TiledMma = typename CollectiveMainloop::TiledMma; + using ArchTag = typename CollectiveMainloop::ArchTag; + using ElementA = typename CollectiveMainloop::ElementA; + using StrideA = typename CollectiveMainloop::StrideA; + using ElementB = typename CollectiveMainloop::ElementB; + using StrideB = typename CollectiveMainloop::StrideB; + using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy; + using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator; + using MainloopArguments = typename CollectiveMainloop::Arguments; + using MainloopParams = typename CollectiveMainloop::Params; + + // Epilogue derived types + using CollectiveEpilogue = CollectiveEpilogue_; + using ElementC = typename CollectiveEpilogue::ElementC; + using StrideC = typename CollectiveEpilogue::StrideC; + using ElementD = typename CollectiveEpilogue::ElementD; + using StrideD = typename CollectiveEpilogue::StrideD; + using EpilogueArguments = typename CollectiveEpilogue::Arguments; + using EpilogueParams = typename CollectiveEpilogue::Params; + static_assert(cute::is_same_v, + "Mainloop and epilogue do not agree on accumulator value type."); + + static constexpr int SharedStorageSize = cute::max( + sizeof(typename CollectiveMainloop::SharedStorage), + sizeof(typename CollectiveEpilogue::SharedStorage)); + + static constexpr uint32_t MaxThreadsPerBlock = cute::size(TiledMma{}); + static constexpr uint32_t MinBlocksPerMultiprocessor = 1; + + // Device side arguments + struct Arguments { + GemmUniversalMode mode{}; + ProblemShape problem_shape{}; + MainloopArguments mainloop{}; + EpilogueArguments epilogue{}; + KernelHardwareInfo hw_info{}; + }; + + // Kernel entry point API + struct Params { + GemmUniversalMode mode; + ProblemShape problem_shape; + MainloopParams mainloop; + EpilogueParams epilogue; + }; + + // + // Methods + // + + // Convert to underlying arguments. In this case, a simple copy for the aliased type. + static + Params + to_underlying_arguments(Arguments const& args, void* workspace) { + (void) workspace; + return { + args.mode, + args.problem_shape, + CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace), + CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace) + }; + } + + static bool + can_implement(Arguments const& args) { + return args.mode == GemmUniversalMode::kGemm or + (args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4); + } + + static int + get_workspace_size(Arguments const& args) { + return 0; + } + + static dim3 + get_grid_shape(Params const& params) { + int batch_count = 1; + if constexpr (rank(ProblemShape{}) == 4) { + batch_count = cute::size<3>(params.problem_shape); + } + + return dim3( + cute::size(cute::ceil_div(cute::shape<0>(params.problem_shape), cute::shape<0>(TileShape{}))), + cute::size(cute::ceil_div(cute::shape<1>(params.problem_shape), cute::shape<1>(TileShape{}))), + batch_count + ); + } + + static dim3 + get_block_shape() { + return dim3(MaxThreadsPerBlock, 1, 1); + } + + MCTLASS_DEVICE + void + operator()(Params const& params, char* smem_buf) { + using namespace cute; + using X = Underscore; + + // Preconditions + CUTE_STATIC_ASSERT(is_static::value); + + // Separate out problem shape for convenience + // Optionally append _1s until problem shape is rank-4 in case its is only rank-3 (MNK) + auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{}); + auto M = get<0>(problem_shape_MNKL); + auto N = get<1>(problem_shape_MNKL); + auto K = get<2>(problem_shape_MNKL); + auto L = get<3>(problem_shape_MNKL); + + // Preconditions + static_assert(rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideB{}) == 3, "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); + + // Get the appropriate blocks for this thread block -- potential for thread block locality + int thread_idx = int(threadIdx.x); + auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K) + // auto [m_coord, n_coord, l_coord] = blockIdx; + auto m_coord = blockIdx.x; + auto n_coord = blockIdx.y; + auto l_coord = blockIdx.z; + auto blk_coord_mnkl = make_coord(m_coord, n_coord, _, l_coord); // (m,n,k,l) + + // Represent the full tensors + Tensor mA_mkl = make_tensor(make_gmem_ptr(params.mainloop.ptr_A), make_shape(M,K,L), params.mainloop.dA); //(m,k,l) + Tensor mB_nkl = make_tensor(make_gmem_ptr(params.mainloop.ptr_B), make_shape(N,K,L), params.mainloop.dB); //(n,k,l) + + // Get batch slice + Tensor mA_mk = mA_mkl(_,_,l_coord); // (m,k) + Tensor mB_nk = mB_nkl(_,_,l_coord); // (n,k) + + // Slice to get the tiles this thread block is responsible for + Tensor gA = local_tile(mA_mk, blk_shape, take<0,3>(blk_coord_mnkl), Step<_1, X,_1>{}); // (BLK_M,BLK_K,k) + Tensor gB = local_tile(mB_nk, blk_shape, take<0,3>(blk_coord_mnkl), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + + // Compute tile residues for predication + auto m_max_coord = M - size<0>(gA) * get<0>(blk_coord_mnkl); // M - BLK_M * m_coord + auto n_max_coord = N - size<0>(gB) * get<1>(blk_coord_mnkl); // N - BLK_N * n_coord + auto k_residue = K - size<1>(gA) * size<2>(gA); // K - BLK_K * k_coord_max + auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue); + + // Allocate the tiled_mma and the accumulators for the (M,N) blk_shape + TiledMma tiled_mma; + Tensor accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape)); // (MMA,MMA_M,MMA_N) + clear(accumulators); + + auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA)); + int k_tile_count = size<2>(gA); + + // Perform the collective scoped MMA + CollectiveMainloop collective_mma; + collective_mma( + accumulators, + gA, + gB, + accumulators, + k_tile_iter, k_tile_count, + residue_mnk, + thread_idx, + smem_buf + ); + + // Epilogue and write to gD + CollectiveEpilogue epilogue{params.epilogue}; + epilogue( + problem_shape_MNKL, + blk_shape, + blk_coord_mnkl, + accumulators, + tiled_mma, + residue_mnk, + thread_idx, + smem_buf + ); + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::kernel diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma.hpp b/csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma.hpp new file mode 100644 index 0000000..535d270 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma.hpp @@ -0,0 +1,328 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/kernel_hardware_info.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "mctlass/arch/mma_sm90.h" +#include "mctlass/epilogue/collective/detail.hpp" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/dispatch_policy.hpp" +#include "mctlass/gemm/kernel/sm90_tile_scheduler.hpp" + +#include "cute/tensor.hpp" + +/////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::kernel { + +namespace detail { + +// IF_SWAP_AB::value will be true only if: +// class T has member SwapAB and T::SwapAB is true +template +struct IF_SWAP_AB { static constexpr bool value = false; }; + +template +struct IF_SWAP_AB > +{ static constexpr bool value = T::SwapAB; }; + +} // namespace + +/////////////////////////////////////////////////////////////////////////////// + +template < + class ProblemShape_, + class CollectiveMainloop_, + class CollectiveEpilogue_, + class GridSwizzle_ +> +class GemmUniversal< + ProblemShape_, + CollectiveMainloop_, + CollectiveEpilogue_, + GridSwizzle_, + cute::enable_if_t>> +{ +public: + // + // Type Aliases + // + using ProblemShape = ProblemShape_; + using GridSwizzle = GridSwizzle_; + static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4, + "ProblemShape{} should be or "); + + // Mainloop derived types + using CollectiveMainloop = CollectiveMainloop_; + using TileShape = typename CollectiveMainloop::TileShape; + using TiledMma = typename CollectiveMainloop::TiledMma; + using ArchTag = typename CollectiveMainloop::ArchTag; + using ElementA = typename CollectiveMainloop::ElementA; + using StrideA = typename CollectiveMainloop::StrideA; + using ElementB = typename CollectiveMainloop::ElementB; + using StrideB = typename CollectiveMainloop::StrideB; + using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy; + using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator; + using ClusterShape = typename DispatchPolicy::ClusterShape; + using MainloopArguments = typename CollectiveMainloop::Arguments; + using MainloopParams = typename CollectiveMainloop::Params; + static_assert(ArchTag::kMinComputeCapability >= 90); + + // Epilogue derived types + using CollectiveEpilogue = CollectiveEpilogue_; + using ElementC = typename CollectiveEpilogue::ElementC; + using StrideC = typename CollectiveEpilogue::StrideC; + using ElementD = typename CollectiveEpilogue::ElementD; + using StrideD = typename CollectiveEpilogue::StrideD; + using EpilogueArguments = typename CollectiveEpilogue::Params; + using EpilogueParams = typename CollectiveEpilogue::Params; + static_assert(cute::is_same_v, + "Mainloop and epilogue do not agree on accumulator value type."); + + static constexpr int SharedStorageSize = cute::max( + sizeof(typename CollectiveMainloop::SharedStorage), + sizeof(typename CollectiveEpilogue::SharedStorage)); + + static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{}); + static constexpr uint32_t MinBlocksPerMultiprocessor = 1; + + // Device side arguments + struct Arguments { + GemmUniversalMode mode{}; + ProblemShape problem_shape{}; + MainloopArguments mainloop{}; + EpilogueArguments epilogue{}; + KernelHardwareInfo hw_info{}; + }; + + // Kernel entry point API + struct Params { + GemmUniversalMode mode; + ProblemShape problem_shape; + MainloopParams mainloop; + EpilogueParams epilogue; + }; + + // + // Methods + // + + // Convert to underlying arguments. In this case, a simple copy for the aliased type. + static + Params + to_underlying_arguments(Arguments const& args, void* workspace) { + (void) workspace; + auto problem_shape = args.problem_shape; + if constexpr (detail::IF_SWAP_AB::value) { + // swap M/N + get<0>(problem_shape) = get<1>(args.problem_shape); + get<1>(problem_shape) = get<0>(args.problem_shape); + } + return { + args.mode, + problem_shape, + CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace), + CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace) + }; + } + + MCTLASS_HOST_DEVICE static + bool + can_implement(Arguments const& args) { + bool implementable = (args.mode == GemmUniversalMode::kGemm) or + (args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4); + if (!implementable) { + MCTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Size don't meet the requirements.\n"); + return implementable; + } + constexpr int tma_alignment_bits = 128; + constexpr int min_tma_aligned_elements = tma_alignment_bits / mctlass::sizeof_bits::value; + auto M = get<0>(args.problem_shape); + auto N = get<1>(args.problem_shape); + auto K = get<2>(args.problem_shape); + // Contiguous dimension for the TMA tensor should be 128b aligned + implementable = std::is_same_v, layout::RowMajor> ? + K % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0; + implementable = implementable && (std::is_same_v, layout::RowMajor> ? + N % min_tma_aligned_elements == 0 : K % min_tma_aligned_elements == 0); + implementable = implementable && (!mctlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA::value || + (mctlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA::value && + std::is_same_v, layout::RowMajor> ? + N % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0)); + if (!implementable) { + MCTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + return implementable; + } + + constexpr bool is_beta_supported = + CollectiveEpilogue::ThreadEpilogueOp::kScale == mctlass::epilogue::thread::ScaleType::Default; + implementable = is_beta_supported || (args.epilogue.thread.beta == 0 && args.epilogue.thread.beta_ptr == nullptr); + if (!implementable) { + MCTLASS_TRACE_HOST(" CAN IMPLEMENT: Scaling params don't meet ThreadEpilogueOp requirements.\n"); + return implementable; + } + + return implementable; + } + + static int + get_workspace_size(Arguments const& args) { + return 0; + } + + // Computes the kernel launch grid shape based on runtime parameters + static dim3 + get_grid_shape(Params const& params) { + auto cluster_shape = ClusterShape{}; + auto tile_shape = TileShape{}; + auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{}); + return detail::PersistentTileSchedulerSm90::get_tiled_blk_shape_mnl( + problem_shape_MNKL, tile_shape, cluster_shape); + } + + static dim3 + get_block_shape() { + return dim3(MaxThreadsPerBlock, 1, 1); + } + + MCTLASS_DEVICE + void + operator()(Params const& params, char* smem_buf) { + using namespace cute; + using X = Underscore; + + // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a. + #if ! defined(__MACA_ARCH_FEAT_SM90_ALL) + if constexpr(size<0>(typename TiledMma::AtomShape_MNK{}) == 64) { + printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); + return; + } + #endif + + // Preconditions + static_assert(rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideB{}) == 3, "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); + + int thread_idx = int(threadIdx.x); + int warp_idx = canonical_warp_idx(); + int lane_predicate = cute::elect_one_sync(); + + // Issue Tma Descriptor Prefetch from a single thread + if ((warp_idx == 0) && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(params.mainloop); + } + + // Separate out problem shape for convenience + // Optionally append _1s until problem shape is rank-4 in case its is only rank-3 (MNK) + auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{}); + auto M = get<0>(problem_shape_MNKL); + auto N = get<1>(problem_shape_MNKL); + auto K = get<2>(problem_shape_MNKL); + auto L = get<3>(problem_shape_MNKL); + + // TMA requires special handling of strides to deal with coord codomain mapping + // Represent the full tensors -- get these from TMA + Tensor mA_mkl = params.mainloop.tma_load_a.get_tma_tensor(make_shape(M,K,L)); // (m,k,l) + Tensor mB_nkl = params.mainloop.tma_load_b.get_tma_tensor(make_shape(N,K,L)); // (n,k,l) + + // Get the appropriate blocks for this thread block -- potential for thread block locality + auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K) + auto blk_coord = make_coord(_,_,_); // (m,n,k) -- defer the slice + + // Make tiled views + Tensor gA_mkl = local_tile(mA_mkl, blk_shape, blk_coord, Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l) + Tensor gB_nkl = local_tile(mB_nkl, blk_shape, blk_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l) + + // Compute m_coord, n_coord, and l_coord with their post-tiled shapes + auto m_coord = idx2crd(int(blockIdx.x), shape<2>(gA_mkl)); + auto n_coord = idx2crd(int(blockIdx.y), shape<2>(gB_nkl)); + auto l_coord = idx2crd(int(blockIdx.z), shape<4>(gB_nkl)); + auto output_tile_coord = make_coord(m_coord, n_coord, _, l_coord); + + // Slice with m_coord and n_coord + Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k) + Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k) + + // Allocate the tiled_mma and the accumulators for the (M,N) blk_shape + TiledMma tiled_mma; + Tensor accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape)); // (MMA,MMA_M,MMA_N) + + auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA)); + auto k_tile_count = size<2>(gA); + + // Perform the collective scoped MMA + CollectiveMainloop collective_mma; + collective_mma( + gA, params.mainloop.tma_load_a, + gB, params.mainloop.tma_load_b, + accumulators, + k_tile_iter, k_tile_count, + thread_idx, + smem_buf, + params.mainloop + ); + + constexpr int BLK_M_RANK = rank<0>(blk_shape); + bool m_oob = int(blockIdx.x) >= size<2>(gA_mkl); + auto m_max_coord = unwrap(cute::transform(make_seq{}, [&](auto i) { + return m_oob ? 0 : get(M) - get<0,i>(blk_shape) * get(m_coord); + })); + + constexpr int BLK_N_RANK = rank<1>(blk_shape); + bool n_oob = int(blockIdx.y) >= size<2>(gB_nkl); + auto n_max_coord = unwrap(cute::transform(make_seq{}, [&](auto i) { + return n_oob ? 0 : get(N) - get<1,i>(blk_shape) * get(n_coord); + })); + auto residue_mnk = make_tuple(m_max_coord, n_max_coord, Int<0>{}); + + // Epilogue and write to gD + CollectiveEpilogue epilogue{params.epilogue}; + epilogue( + problem_shape_MNKL, + blk_shape, + output_tile_coord, + accumulators, + tiled_mma, + residue_mnk, + thread_idx, + smem_buf + ); + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::kernel diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp b/csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp new file mode 100644 index 0000000..686240b --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp @@ -0,0 +1,441 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/kernel_hardware_info.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "mctlass/arch/reg_reconfig.h" +#include "mctlass/arch/mma_sm90.h" +#include "mctlass/epilogue/collective/detail.hpp" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/dispatch_policy.hpp" +#include "mctlass/gemm/kernel/sm90_tile_scheduler.hpp" +#include "mctlass/pipeline/pipeline.hpp" +#include "cute/tensor.hpp" + +/////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::kernel { + +/////////////////////////////////////////////////////////////////////////////// + +template < + class ProblemShape_, + class CollectiveMainloop_, + class CollectiveEpilogue_, + class GridSwizzle_ +> +class GemmUniversal< + ProblemShape_, + CollectiveMainloop_, + CollectiveEpilogue_, + GridSwizzle_, + cute::enable_if_t>> +{ +public: + // + // Type Aliases + // + using ProblemShape = ProblemShape_; + using GridSwizzle = GridSwizzle_; + static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4, + "ProblemShape{} should be or "); + + // Mainloop derived types + using CollectiveMainloop = CollectiveMainloop_; + using TileShape = typename CollectiveMainloop::TileShape; + using TiledMma = typename CollectiveMainloop::TiledMma; + using ArchTag = typename CollectiveMainloop::ArchTag; + using ElementA = typename CollectiveMainloop::ElementA; + using StrideA = typename CollectiveMainloop::StrideA; + using ElementB = typename CollectiveMainloop::ElementB; + using StrideB = typename CollectiveMainloop::StrideB; + using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy; + using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator; + using ClusterShape = typename DispatchPolicy::ClusterShape; + using MainloopArguments = typename CollectiveMainloop::Arguments; + using MainloopParams = typename CollectiveMainloop::Params; + static_assert(ArchTag::kMinComputeCapability >= 90); + + // Epilogue derived types + using CollectiveEpilogue = CollectiveEpilogue_; + using ElementC = typename CollectiveEpilogue::ElementC; + using StrideC = typename CollectiveEpilogue::StrideC; + using ElementD = typename CollectiveEpilogue::ElementD; + using StrideD = typename CollectiveEpilogue::StrideD; + using EpilogueArguments = typename CollectiveEpilogue::Arguments; + using EpilogueParams = typename CollectiveEpilogue::Params; + static_assert(cute::is_same_v, + "Mainloop and epilogue do not agree on accumulator value type."); + + // Kernel level shared memory storage + struct SharedStorage { + union TensorStorage { + using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage; + using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage; + + MainloopTensorStorage mainloop; + EpilogueTensorStorage epilogue; + } tensors; + + struct PipelineStorage : cute::aligned_struct<16> { + using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage; + using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage; + + alignas(16) MainloopPipelineStorage mainloop; + alignas(16) EpiLoadPipelineStorage epi_load; + } pipelines; + }; + + static constexpr int SharedStorageSize = sizeof(SharedStorage); + + static constexpr uint32_t NumLoadWarpGroups = 1; + static constexpr uint32_t NumMmaWarpGroups = 1; + static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{}) + (NumLoadWarpGroups * NumThreadsPerWarpGroup); + static constexpr uint32_t MinBlocksPerMultiprocessor = 1; + + // Device side arguments + struct Arguments { + GemmUniversalMode mode{}; + ProblemShape problem_shape{}; + MainloopArguments mainloop{}; + EpilogueArguments epilogue{}; + KernelHardwareInfo hw_info{}; + }; + + // Kernel entry point API + struct Params { + GemmUniversalMode mode; + ProblemShape problem_shape; + MainloopParams mainloop; + EpilogueParams epilogue; + }; + + // + // Methods + // + + // Convert to underlying arguments. In this case, a simple copy for the aliased type. + static + Params + to_underlying_arguments(Arguments const& args, void* workspace) { + (void) workspace; + auto problem_shape = args.problem_shape; + if constexpr (detail::IF_SWAP_AB::value) { + // swap M/N + get<0>(problem_shape) = get<1>(args.problem_shape); + get<1>(problem_shape) = get<0>(args.problem_shape); + } + return { + args.mode, + problem_shape, + CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace), + CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace) + }; + } + + MCTLASS_HOST_DEVICE static + bool + can_implement(Arguments const& args) { + bool implementable = (args.mode == GemmUniversalMode::kGemm) or + (args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4); + if (!implementable) { + MCTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Size don't meet the requirements.\n"); + return implementable; + } + constexpr int tma_alignment_bits = 128; + constexpr int min_tma_aligned_elements = tma_alignment_bits / mctlass::sizeof_bits::value; + auto M = get<0>(args.problem_shape); + auto N = get<1>(args.problem_shape); + auto K = get<2>(args.problem_shape); + // Contiguous dimension for the TMA tensor should be 128b aligned + implementable = std::is_same_v, layout::RowMajor> ? + K % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0; + implementable = implementable && (std::is_same_v, layout::RowMajor> ? + N % min_tma_aligned_elements == 0 : K % min_tma_aligned_elements == 0); + implementable = implementable && (!mctlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA::value || + (mctlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA::value && + std::is_same_v, layout::RowMajor> ? + N % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0)); + if (!implementable) { + MCTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + return implementable; + } + + constexpr bool is_beta_supported = not cute::is_void_v && + CollectiveEpilogue::ThreadEpilogueOp::kScale == mctlass::epilogue::thread::ScaleType::Default; + implementable = is_beta_supported || (args.epilogue.thread.beta == 0 && args.epilogue.thread.beta_ptr == nullptr); + if (!implementable) { + MCTLASS_TRACE_HOST(" CAN IMPLEMENT: Scaling params don't meet ThreadEpilogueOp requirements.\n"); + return implementable; + } + + return implementable; + } + + static + int + get_workspace_size(Arguments const& args) { + return 0; + } + + // Computes the kernel launch grid shape based on runtime parameters + static dim3 + get_grid_shape(Params const& params) { + auto cluster_shape = ClusterShape{}; + auto tile_shape = TileShape{}; + auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{}); + return detail::PersistentTileSchedulerSm90::get_tiled_blk_shape_mnl( + problem_shape_MNKL, tile_shape, cluster_shape); + } + + static dim3 + get_block_shape() { + return dim3(MaxThreadsPerBlock, 1, 1); + } + + MCTLASS_DEVICE + void + operator()(Params const& params, char* smem_buf) { + using namespace cute; + using X = Underscore; + + // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a. + #if ! defined(__MACA_ARCH_FEAT_SM90_ALL) + if constexpr(size<0>(typename TiledMma::AtomShape_MNK{}) == 64) { + printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); + return; + } + #endif + + enum class WarpGroupRole { + Producer = 0, + Consumer = 1, + }; + + // Kernel level shared memory storage + SharedStorage& shared_storage = *reinterpret_cast(smem_buf); + + int thread_idx = int(threadIdx.x); + int warp_idx = canonical_warp_idx(); + int warp_group_thread_idx = thread_idx % NumThreadsPerWarpGroup; + auto warp_group_role = WarpGroupRole(canonical_warp_group_idx()); + int lane_predicate = cute::elect_one_sync(); + + // Issue Tma Descriptor Prefetch from a single thread + if ((warp_idx == 0) && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(params.mainloop); + CollectiveEpilogue::prefetch_tma_descriptors(params.epilogue); + } + + // Mainloop Load pipeline + using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline; + typename MainloopPipeline::Params mainloop_pipeline_params; + if (warp_group_role == WarpGroupRole::Producer) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Consumer) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; + } + mainloop_pipeline_params.is_leader = warp_group_thread_idx == 0; + mainloop_pipeline_params.num_consumers = NumThreadsPerWarpGroup; + mainloop_pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytes; + MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params); + + // Epilogue Load pipeline + using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline; + typename EpiLoadPipeline::Params epi_load_pipeline_params; + if (warp_group_role == WarpGroupRole::Producer) { + epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Consumer) { + epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer; + } + epi_load_pipeline_params.dst_blockid = cute::block_rank_in_cluster(); + epi_load_pipeline_params.producer_arv_count = 1; // 1 thread issues TMA load + epi_load_pipeline_params.consumer_arv_count = NumThreadsPerWarpGroup; + epi_load_pipeline_params.transaction_bytes = CollectiveEpilogue::TmaTransactionBytes; + EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params); + + // Epilogue Store pipeline + using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline; + typename EpiStorePipeline::Params epi_store_pipeline_params; + epi_store_pipeline_params.always_wait = true; + EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params); + + // Initialize starting pipeline states for the collectives + // Epilogue store pipe is producer-only (consumer is TMA unit, waits via scoreboarding) + typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state; + typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state; + + // For the DMA Load (producer) we start with an opposite phase + // i.e., we skip all waits since we know that the buffer is indeed empty + PipelineState mainloop_pipe_producer_state = mctlass::make_producer_start_state(); + PipelineState epi_load_pipe_producer_state = mctlass::make_producer_start_state(); + PipelineState epi_store_pipe_producer_state = mctlass::make_producer_start_state(); + + auto cluster_wait_fn = [&] () { + // We need this to guarantee that the Pipeline init is visible + // To all producers and consumer thread blocks in the Cluster + if constexpr (size(ClusterShape{}) > 1) { + cute::cluster_arrive_relaxed(); + return [] () { cute::cluster_wait(); }; + } + else { + __syncthreads(); + return [] () {}; // do nothing + } + } (); + + // Preconditions + static_assert(rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideB{}) == 3, "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); + + // Separate out problem shape for convenience + // Optionally append _1s until problem shape is rank-4 in case its is only rank-3 (MNK) + auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{}); + auto M = get<0>(problem_shape_MNKL); + auto N = get<1>(problem_shape_MNKL); + auto K = get<2>(problem_shape_MNKL); + auto L = get<3>(problem_shape_MNKL); + + // TMA requires special handling of strides to deal with coord codomain mapping + // Represent the full tensors -- get these from TMA + Tensor mA_mkl = params.mainloop.tma_load_a.get_tma_tensor(make_shape(M,K,L)); // (m,k,l) + Tensor mB_nkl = params.mainloop.tma_load_b.get_tma_tensor(make_shape(N,K,L)); // (n,k,l) + + // Get the appropriate blocks for this thread block -- potential for thread block locality + auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K) + TiledMma tiled_mma; + + // Make tiled views, defer the slice + Tensor gA_mkl = local_tile(mA_mkl, blk_shape, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l) + Tensor gB_nkl = local_tile(mB_nkl, blk_shape, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l) + + // Compute m_coord, n_coord, and l_coord with their post-tiled shapes + auto m_coord = idx2crd(int(blockIdx.x), shape<2>(gA_mkl)); + auto n_coord = idx2crd(int(blockIdx.y), shape<2>(gB_nkl)); + auto l_coord = idx2crd(int(blockIdx.z), shape<4>(gB_nkl)); + auto blk_coord = make_coord(m_coord, n_coord, _, l_coord); + + // Slice with m_coord and n_coord + Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k) + Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k) + + // Get pipeline iterators and increments from tensor shapes + auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA)); + auto k_tile_count = size<2>(gA); + auto c_tile_count = CollectiveEpilogue::get_load_pipe_increment(blk_shape); + auto d_tile_count = CollectiveEpilogue::get_store_pipe_increment(blk_shape); + + // Wait for all thread blocks in the Cluster + cluster_wait_fn(); + + // In a warp specialized kernel, collectives expose data movement and compute operations separately + CollectiveMainloop collective_mainloop; + CollectiveEpilogue collective_epilogue{params.epilogue}; + + if (warp_group_role == WarpGroupRole::Producer) { + collective_mainloop.load( + mainloop_pipeline, + mainloop_pipe_producer_state, + gA, params.mainloop.tma_load_a, + gB, params.mainloop.tma_load_b, + k_tile_iter, k_tile_count, + thread_idx, + shared_storage.tensors.mainloop + ); + // Update starting mainloop pipeline state for the pipeline drain + mainloop_pipe_producer_state.advance(k_tile_count); + // Make sure mainloop consumer has been waited upon before issuing epilogue load + collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state); + + if (collective_epilogue.is_source_needed()) { + collective_epilogue.load( + epi_load_pipeline, + epi_load_pipe_producer_state, + problem_shape_MNKL, + blk_shape, + blk_coord, + tiled_mma, + warp_group_thread_idx, + shared_storage.tensors.epilogue + ); + // Update starting load pipeline state for the pipeline drain + epi_load_pipe_producer_state.advance(c_tile_count); + collective_epilogue.load_tail(epi_load_pipeline, epi_load_pipe_producer_state); + } + } + else if (warp_group_role == WarpGroupRole::Consumer) { + Tensor accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape)); // (MMA,MMA_M,MMA_N) + + collective_mainloop.mma( + mainloop_pipeline, + mainloop_pipe_consumer_state, + accumulators, + k_tile_count, + thread_idx, + shared_storage.tensors.mainloop, + params.mainloop + ); + + // Make sure the math instructions are done and free buffers before entering the epilogue + collective_mainloop.mma_tail( + mainloop_pipeline, + mainloop_pipe_consumer_state, + k_tile_count + ); + + // Epilogue and write to gD + collective_epilogue.store( + epi_load_pipeline, + epi_load_pipe_consumer_state, + epi_store_pipeline, + epi_store_pipe_producer_state, + problem_shape_MNKL, + blk_shape, + blk_coord, + accumulators, + tiled_mma, + warp_group_thread_idx, + shared_storage.tensors.epilogue + ); + } + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::kernel diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp b/csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp new file mode 100644 index 0000000..d37990b --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp @@ -0,0 +1,504 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/kernel_hardware_info.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "mctlass/arch/reg_reconfig.h" +#include "mctlass/arch/mma_sm90.h" +#include "mctlass/epilogue/collective/detail.hpp" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/dispatch_policy.hpp" +#include "mctlass/gemm/kernel/sm90_tile_scheduler.hpp" +#include "mctlass/pipeline/pipeline.hpp" +#include "cute/tensor.hpp" + +/////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::kernel { + +/////////////////////////////////////////////////////////////////////////////// + +template < + class ProblemShape_, + class CollectiveMainloop_, + class CollectiveEpilogue_, + class GridSwizzle_ +> +class GemmUniversal< + ProblemShape_, + CollectiveMainloop_, + CollectiveEpilogue_, + GridSwizzle_, + cute::enable_if_t>> +{ +public: + // + // Type Aliases + // + using ProblemShape = ProblemShape_; + using GridSwizzle = GridSwizzle_; + static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4, + "ProblemShape{} should be or "); + + // Mainloop derived types + using CollectiveMainloop = CollectiveMainloop_; + using TileShape = typename CollectiveMainloop::TileShape; + using TiledMma = typename CollectiveMainloop::TiledMma; + using ArchTag = typename CollectiveMainloop::ArchTag; + using ElementA = typename CollectiveMainloop::ElementA; + using StrideA = typename CollectiveMainloop::StrideA; + using ElementB = typename CollectiveMainloop::ElementB; + using StrideB = typename CollectiveMainloop::StrideB; + using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy; + using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator; + using ClusterShape = typename DispatchPolicy::ClusterShape; + using MainloopArguments = typename CollectiveMainloop::Arguments; + using MainloopParams = typename CollectiveMainloop::Params; + + // Epilogue derived types + using CollectiveEpilogue = CollectiveEpilogue_; + using ElementC = typename CollectiveEpilogue::ElementC; + using StrideC = typename CollectiveEpilogue::StrideC; + using ElementD = typename CollectiveEpilogue::ElementD; + using StrideD = typename CollectiveEpilogue::StrideD; + using EpilogueArguments = typename CollectiveEpilogue::Arguments; + using EpilogueParams = typename CollectiveEpilogue::Params; + static_assert(cute::is_same_v, + "Mainloop and epilogue do not agree on accumulator value type."); + + using PersistentTileSchedulerParams = typename detail::PersistentTileSchedulerSm90::Params; + static_assert(ArchTag::kMinComputeCapability >= 90); + + static constexpr uint32_t NumLoadWarpGroups = 1; + static constexpr uint32_t NumMmaWarpGroups = 1; + static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{}) + (NumLoadWarpGroups * NumThreadsPerWarpGroup); + static constexpr uint32_t MinBlocksPerMultiprocessor = 1; + + /// Register requirement for Load and Math WGs + static constexpr uint32_t LoadRegisterRequirement = 40; + static constexpr uint32_t MmaRegisterRequirement = 232; + + // Kernel level shared memory storage + struct SharedStorage { + struct TensorStorage : cute::aligned_struct<128> { + using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage; + using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage; + + MainloopTensorStorage mainloop; + EpilogueTensorStorage epilogue; + } tensors; + + struct PipelineStorage : cute::aligned_struct<16> { + using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage; + using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage; + + alignas(16) MainloopPipelineStorage mainloop; + alignas(16) EpiLoadPipelineStorage epi_load; + } pipelines; + }; + + static constexpr int SharedStorageSize = sizeof(SharedStorage); + + // Device side arguments + struct Arguments { + GemmUniversalMode mode{}; + ProblemShape problem_shape{}; + MainloopArguments mainloop{}; + EpilogueArguments epilogue{}; + KernelHardwareInfo hw_info{}; + }; + + // Kernel entry point API + struct Params { + GemmUniversalMode mode; + ProblemShape problem_shape; + MainloopParams mainloop; + EpilogueParams epilogue; + KernelHardwareInfo hw_info; + PersistentTileSchedulerParams scheduler; + }; + + // + // Methods + // + + // Convert to underlying arguments. In this case, a simple copy for the aliased type. + static + Params + to_underlying_arguments(Arguments const& args, void* workspace) { + MCTLASS_TRACE_HOST("to_underlying_arguments():"); + + (void) workspace; + auto problem_shape = args.problem_shape; + if constexpr (detail::IF_SWAP_AB::value) { + // swap M/N + get<0>(problem_shape) = get<1>(args.problem_shape); + get<1>(problem_shape) = get<0>(args.problem_shape); + } + auto problem_shape_MNKL = append<4>(problem_shape, Int<1>{}); + + // Get SM count if needed, otherwise use user supplied SM count + int sm_count = args.hw_info.sm_count; + if (sm_count <= 0) { + MCTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n" + " For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count."); + sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id); + } + + MCTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count); + return { + args.mode, + problem_shape, + CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace), + CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace), + {args.hw_info.device_id, sm_count}, + detail::PersistentTileSchedulerSm90::to_underlying_arguments(problem_shape_MNKL, TileShape{}, ClusterShape{}) + }; + } + + MCTLASS_HOST_DEVICE static + bool + can_implement(Arguments const& args) { + bool implementable = (args.mode == GemmUniversalMode::kGemm) or + (args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4); + if (!implementable) { + MCTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Size don't meet the requirements.\n"); + return implementable; + } + constexpr int tma_alignment_bits = 128; + constexpr int min_tma_aligned_elements = tma_alignment_bits / mctlass::sizeof_bits::value; + auto M = get<0>(args.problem_shape); + auto N = get<1>(args.problem_shape); + auto K = get<2>(args.problem_shape); + // Contiguous dimension for the TMA tensor should be 128b aligned + implementable = std::is_same_v, layout::RowMajor> ? + K % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0; + implementable = implementable && (std::is_same_v, layout::RowMajor> ? + N % min_tma_aligned_elements == 0 : K % min_tma_aligned_elements == 0); + implementable = implementable && (!mctlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA::value || + (mctlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA::value && + std::is_same_v, layout::RowMajor> ? + N % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0)); + if (!implementable) { + MCTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + return implementable; + } + + constexpr bool is_beta_supported = + CollectiveEpilogue::ThreadEpilogueOp::kScale == mctlass::epilogue::thread::ScaleType::Default; + implementable = is_beta_supported || (args.epilogue.thread.beta == 0 && args.epilogue.thread.beta_ptr == nullptr); + if (!implementable) { + MCTLASS_TRACE_HOST(" CAN IMPLEMENT: Scaling params don't meet ThreadEpilogueOp requirements.\n"); + return implementable; + } + + return implementable; + } + + static + int + get_workspace_size(Arguments const& args) { + return 0; + } + + // Computes the kernel launch grid shape based on runtime parameters + static dim3 + get_grid_shape(Params const& params) { + // Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run concurrently + return detail::PersistentTileSchedulerSm90::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info); + } + + static dim3 + get_block_shape() { + return dim3(MaxThreadsPerBlock, 1, 1); + } + + MCTLASS_DEVICE + void + operator()(Params const& params, char* smem_buf) { + using namespace cute; + using X = Underscore; + + // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a. + #if ! defined(__MACA_ARCH_FEAT_SM90_ALL) + if constexpr(size<0>(typename TiledMma::AtomShape_MNK{}) == 64) { + printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); + return; + } + #endif + + // Preconditions + static_assert(size(TiledMma{}) == 256, "Cooperative kernel must have TiledMMA operating using 256 threads."); + static_assert(size<0>(TileShape{}) >= 128, + "Cooperative kernel requires Tile Size to be greater than or equal to 128 along the M-dimension."); + + static_assert(rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideB{}) == 3, "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); + + /* In the Cooperative kernel, Consumer0 and Consumer1 collaborate on the same tile */ + enum class WarpGroupRole { + Producer = 0, + Consumer0 = 1, + Consumer1 = 2 + }; + + // Kernel level shared memory storage + SharedStorage& shared_storage = *reinterpret_cast(smem_buf); + + int thread_idx = int(threadIdx.x); + int warp_idx = canonical_warp_idx(); + int warp_group_thread_idx = thread_idx % NumThreadsPerWarpGroup; + int mma_thread_idx = thread_idx % size(TiledMma{}); + auto warp_group_role = WarpGroupRole(canonical_warp_group_idx()); + int lane_predicate = cute::elect_one_sync(); + + // Issue Tma Descriptor Prefetch from a single thread + if ((warp_idx == 0) && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(params.mainloop); + CollectiveEpilogue::prefetch_tma_descriptors(params.epilogue); + } + + // Mainloop Load pipeline + using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline; + typename MainloopPipeline::Params mainloop_pipeline_params; + if (warp_group_role == WarpGroupRole::Producer) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; + } + mainloop_pipeline_params.is_leader = warp_group_thread_idx == 0; + mainloop_pipeline_params.num_consumers = size(TiledMma{}); + mainloop_pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytes; + MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params); + + // Epilogue Load pipeline + using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline; + typename EpiLoadPipeline::Params epi_load_pipeline_params; + if (warp_group_role == WarpGroupRole::Producer) { + epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) { + epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer; + } + epi_load_pipeline_params.dst_blockid = cute::block_rank_in_cluster(); + epi_load_pipeline_params.producer_arv_count = 1; // 1 thread issues TMA load + epi_load_pipeline_params.consumer_arv_count = size(TiledMma{}); + epi_load_pipeline_params.transaction_bytes = CollectiveEpilogue::TmaTransactionBytes; + EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params); + + // Epilogue Store pipeline + using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline; + typename EpiStorePipeline::Params epi_store_pipeline_params; + epi_store_pipeline_params.always_wait = true; + EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params); + + // Initialize starting pipeline states for the collectives + // Epilogue store pipe is producer-only (consumer is TMA unit, waits via scoreboarding) + typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state; + typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state; + + // For the DMA Load (producer) we start with an opposite phase + // i.e., we skip all waits since we know that the buffer is indeed empty + PipelineState mainloop_pipe_producer_state = mctlass::make_producer_start_state(); + PipelineState epi_load_pipe_producer_state = mctlass::make_producer_start_state(); + PipelineState epi_store_pipe_producer_state = mctlass::make_producer_start_state(); + + auto cluster_wait_fn = [&] () { + // We need this to guarantee that the Pipeline init is visible + // To all producers and consumer thread blocks in the Cluster + if constexpr (size(ClusterShape{}) > 1) { + cute::cluster_arrive_relaxed(); + return [] () { cute::cluster_wait(); }; + } + else { + __syncthreads(); + return [] () {}; // do nothing + } + } (); + + // Separate out problem shape for convenience + // Optionally append _1s until problem shape is rank-4 in case its is only rank-3 (MNK) + auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{}); + auto M = get<0>(problem_shape_MNKL); + auto N = get<1>(problem_shape_MNKL); + auto K = get<2>(problem_shape_MNKL); + auto L = get<3>(problem_shape_MNKL); + + // TMA requires special handling of strides to deal with coord codomain mapping + // Represent the full tensors -- get these from TMA + Tensor mA_mkl = params.mainloop.tma_load_a.get_tma_tensor(make_shape(M,K,L)); // (m,k,l) + Tensor mB_nkl = params.mainloop.tma_load_b.get_tma_tensor(make_shape(N,K,L)); // (n,k,l) + + // Get the appropriate blocks for this thread block -- potential for thread block locality + TiledMma tiled_mma; + auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K) + + // Make tiled views, defer the slice + Tensor gA_mkl = local_tile(mA_mkl, blk_shape, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l) + Tensor gB_nkl = local_tile(mB_nkl, blk_shape, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l) + + // Get pipeline stage increments from tensor shapes + auto k_tile_count = size<3>(gA_mkl); + auto c_tile_count = CollectiveEpilogue::get_load_pipe_increment(blk_shape); + auto d_tile_count = CollectiveEpilogue::get_store_pipe_increment(blk_shape); + + detail::PersistentTileSchedulerSm90 scheduler; + auto work_tile_info = scheduler.get_current_work(params.scheduler); + + // In a warp specialized kernel, collectives expose data movement and compute operations separately + CollectiveMainloop collective_mainloop; + CollectiveEpilogue collective_epilogue{params.epilogue}; + + // Wait for all thread blocks in the Cluster + cluster_wait_fn(); + + if (warp_group_role == WarpGroupRole::Producer) { + mctlass::arch::warpgroup_reg_dealloc(); + + while (work_tile_info.is_valid_tile) { + // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape + auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl)); + auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl)); + auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); + auto blk_coord = make_coord(m_coord, n_coord, _, l_coord); + + // Slice with our work tile coordinates to construct mainloop tensor views + Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k) + Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k) + + auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA)); + + collective_mainloop.load( + mainloop_pipeline, + mainloop_pipe_producer_state, + gA, params.mainloop.tma_load_a, + gB, params.mainloop.tma_load_b, + k_tile_iter, k_tile_count, + thread_idx, + shared_storage.tensors.mainloop + ); + // Update starting pipeline state for the next tile + mainloop_pipe_producer_state.advance(k_tile_count); + + if (collective_epilogue.is_source_needed()) { + collective_epilogue.load( + epi_load_pipeline, + epi_load_pipe_producer_state, + problem_shape_MNKL, + blk_shape, + blk_coord, + tiled_mma, + warp_group_thread_idx, + shared_storage.tensors.epilogue + ); + // Update starting pipeline state for the next tile + epi_load_pipe_producer_state.advance(c_tile_count); + } + + // Get next work tile + scheduler.advance_to_next_work(); + work_tile_info = scheduler.get_current_work(params.scheduler); + } // Scheduler work fetch loop + + // Make sure all Consumer Warp Groups have been waited upon + collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state); + if (collective_epilogue.is_source_needed()) { + collective_epilogue.load_tail(epi_load_pipeline, epi_load_pipe_producer_state); + } + } // Producer Warp Group End + + else if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) { + mctlass::arch::warpgroup_reg_alloc(); + + while (work_tile_info.is_valid_tile) { + // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape + auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl)); + auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl)); + auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); + auto blk_coord = make_coord(m_coord, n_coord, _, l_coord); + + // Allocate the the accumulators for the (M,N) blk_shape + Tensor accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape)); // (MMA,MMA_M,MMA_N) + + collective_mainloop.mma( + mainloop_pipeline, + mainloop_pipe_consumer_state, + accumulators, + k_tile_count, + mma_thread_idx, + shared_storage.tensors.mainloop, + params.mainloop + ); + + // Make sure the math instructions are done and free buffers before entering the epilogue + collective_mainloop.mma_tail( + mainloop_pipeline, + mainloop_pipe_consumer_state, + k_tile_count + ); + // Update starting mainloop pipeline state for the next tile + mainloop_pipe_consumer_state.advance(k_tile_count); + + // Epilogue and write to gD + collective_epilogue.store( + epi_load_pipeline, + epi_load_pipe_consumer_state, + epi_store_pipeline, + epi_store_pipe_producer_state, + problem_shape_MNKL, + blk_shape, + blk_coord, + accumulators, + tiled_mma, + mma_thread_idx, + shared_storage.tensors.epilogue + ); + // Update starting load/store pipeline states for the next tile + epi_load_pipe_consumer_state.advance(c_tile_count); + epi_store_pipe_producer_state.advance(d_tile_count); + + // Get next work tile + scheduler.advance_to_next_work(); + work_tile_info = scheduler.get_current_work(params.scheduler); + } // Scheduler work fetch loop + } // Consumer Warp Groups End + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::kernel diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp b/csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp new file mode 100644 index 0000000..aadcf40 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp @@ -0,0 +1,536 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/kernel_hardware_info.hpp" +#include "mctlass/fast_math.h" +#include "cute/arch/cluster_sm90.hpp" +#include "mctlass/arch/reg_reconfig.h" +#include "mctlass/arch/mma_sm90.h" +#include "mctlass/epilogue/collective/detail.hpp" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/dispatch_policy.hpp" +#include "mctlass/gemm/kernel/sm90_tile_scheduler.hpp" +#include "mctlass/pipeline/pipeline.hpp" +#include "mctlass/trace.h" + +#include "cute/tensor.hpp" + +/////////////////////////////////////////////////////////////////////////////// + +namespace mctlass::gemm::kernel { + +/////////////////////////////////////////////////////////////////////////////// + +template < + class ProblemShape_, + class CollectiveMainloop_, + class CollectiveEpilogue_, + class GridSwizzle_ +> +class GemmUniversal< + ProblemShape_, + CollectiveMainloop_, + CollectiveEpilogue_, + GridSwizzle_, + cute::enable_if_t>> +{ +public: + // + // Type Aliases + // + using ProblemShape = ProblemShape_; + using GridSwizzle = GridSwizzle_; + static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4, + "ProblemShape{} should be or "); + + // Mainloop derived types + using CollectiveMainloop = CollectiveMainloop_; + using TileShape = typename CollectiveMainloop::TileShape; + using TiledMma = typename CollectiveMainloop::TiledMma; + using ArchTag = typename CollectiveMainloop::ArchTag; + using ElementA = typename CollectiveMainloop::ElementA; + using StrideA = typename CollectiveMainloop::StrideA; + using ElementB = typename CollectiveMainloop::ElementB; + using StrideB = typename CollectiveMainloop::StrideB; + using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy; + using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator; + using ClusterShape = typename DispatchPolicy::ClusterShape; + using MainloopArguments = typename CollectiveMainloop::Arguments; + using MainloopParams = typename CollectiveMainloop::Params; + using PersistentTileSchedulerParams = typename detail::PersistentTileSchedulerSm90::Params; + static_assert(ArchTag::kMinComputeCapability >= 90); + + // Epilogue derived types + using CollectiveEpilogue = CollectiveEpilogue_; + using ElementC = typename CollectiveEpilogue::ElementC; + using StrideC = typename CollectiveEpilogue::StrideC; + using ElementD = typename CollectiveEpilogue::ElementD; + using StrideD = typename CollectiveEpilogue::StrideD; + using EpilogueArguments = typename CollectiveEpilogue::Arguments; + using EpilogueParams = typename CollectiveEpilogue::Params; + static_assert(cute::is_same_v, + "Mainloop and epilogue do not agree on accumulator value type."); + + static constexpr uint32_t NumLoadWarpGroups = 1; + static constexpr uint32_t NumMmaWarpGroups = 2; + static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{}) + (NumMmaWarpGroups * NumThreadsPerWarpGroup); + static constexpr uint32_t MinBlocksPerMultiprocessor = 1; + + /// Register requirement for Load and Math WGs + static constexpr uint32_t LoadRegisterRequirement = 40; + static constexpr uint32_t MmaRegisterRequirement = 232; + + // Order Sequence barrier with two stages: one for Mainloop and one for Epilogue + static constexpr uint32_t StagesPerMathWarpGroup = 2; + using MathWarpGroupOrderBarrier = mctlass::OrderedSequenceBarrier< + StagesPerMathWarpGroup, NumMmaWarpGroups>; + + // Kernel level shared memory storage + struct SharedStorage { + struct TensorStorage : cute::aligned_struct<128> { + using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage; + using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage; + + MainloopTensorStorage mainloop; + EpilogueTensorStorage epilogue; + } tensors; + + struct PipelineStorage : cute::aligned_struct<16> { + using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage; + using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage; + using MathWarpGroupOrderBarrierStorage = typename MathWarpGroupOrderBarrier::SharedStorage; + + alignas(16) MainloopPipelineStorage mainloop; + alignas(16) EpiLoadPipelineStorage epi_load; + alignas(16) MathWarpGroupOrderBarrierStorage math_wg_order; + } pipelines; + }; + + static constexpr int SharedStorageSize = sizeof(SharedStorage); + + // Device side arguments + struct Arguments { + GemmUniversalMode mode{}; + ProblemShape problem_shape{}; + MainloopArguments mainloop{}; + EpilogueArguments epilogue{}; + KernelHardwareInfo hw_info{}; + }; + + // Kernel entry point API + struct Params { + GemmUniversalMode mode; + ProblemShape problem_shape; + MainloopParams mainloop; + EpilogueParams epilogue; + KernelHardwareInfo hw_info; + PersistentTileSchedulerParams scheduler; + }; + + // + // Methods + // + + // Convert to underlying arguments. In this case, a simple copy for the aliased type. + static + Params + to_underlying_arguments(Arguments const& args, void* workspace) { + MCTLASS_TRACE_HOST("to_underlying_arguments():"); + + (void) workspace; + auto problem_shape = args.problem_shape; + if constexpr (detail::IF_SWAP_AB::value) { + // swap M/N + get<0>(problem_shape) = get<1>(args.problem_shape); + get<1>(problem_shape) = get<0>(args.problem_shape); + } + auto problem_shape_MNKL = append<4>(problem_shape, Int<1>{}); + + // Get SM count if needed, otherwise use user supplied SM count + int sm_count = args.hw_info.sm_count; + if (sm_count <= 0) { + MCTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n" + " For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count."); + sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id); + } + + MCTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count); + return { + args.mode, + problem_shape, + CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace), + CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace), + {args.hw_info.device_id, sm_count}, + detail::PersistentTileSchedulerSm90::to_underlying_arguments(problem_shape_MNKL, TileShape{}, ClusterShape{}) + }; + } + + MCTLASS_HOST_DEVICE static + bool + can_implement(Arguments const& args) { + bool implementable = (args.mode == GemmUniversalMode::kGemm) or + (args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4); + if (!implementable) { + MCTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Size don't meet the requirements.\n"); + return implementable; + } + constexpr int tma_alignment_bits = 128; + constexpr int min_tma_aligned_elements = tma_alignment_bits / mctlass::sizeof_bits::value; + auto M = get<0>(args.problem_shape); + auto N = get<1>(args.problem_shape); + auto K = get<2>(args.problem_shape); + // Contiguous dimension for the TMA tensor should be 128b aligned + implementable = std::is_same_v, layout::RowMajor> ? + K % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0; + implementable = implementable && (std::is_same_v, layout::RowMajor> ? + N % min_tma_aligned_elements == 0 : K % min_tma_aligned_elements == 0); + implementable = implementable && (!mctlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA::value || + (mctlass::epilogue::collective::detail::IF_EPILOGUE_USES_TMA::value && + std::is_same_v, layout::RowMajor> ? + N % min_tma_aligned_elements == 0 : M % min_tma_aligned_elements == 0)); + if (!implementable) { + MCTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n"); + return implementable; + } + + constexpr bool is_beta_supported = + CollectiveEpilogue::ThreadEpilogueOp::kScale == mctlass::epilogue::thread::ScaleType::Default; + implementable = is_beta_supported || (args.epilogue.thread.beta == 0 && args.epilogue.thread.beta_ptr == nullptr); + if (!implementable) { + MCTLASS_TRACE_HOST(" CAN IMPLEMENT: Scaling params don't meet ThreadEpilogueOp requirements.\n"); + return implementable; + } + + return implementable; + } + + static + int + get_workspace_size(Arguments const& args) { + return 0; + } + + // Computes the kernel launch grid shape based on runtime parameters + static dim3 + get_grid_shape(Params const& params) { + // Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run concurrently + return detail::PersistentTileSchedulerSm90::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info); + } + + static dim3 + get_block_shape() { + return dim3(MaxThreadsPerBlock, 1, 1); + } + + MCTLASS_DEVICE + void + operator()(Params const& params, char* smem_buf) { + using namespace cute; + using X = Underscore; + + // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a. + #if ! defined(__MACA_ARCH_FEAT_SM90_ALL) + if constexpr(size<0>(typename TiledMma::AtomShape_MNK{}) == 64) { + printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n"); + return; + } + #endif + + // Preconditions + static_assert(rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideB{}) == 3, "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); + static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>."); + + enum class WarpGroupRole { + Producer = 0, + Consumer0 = 1, + Consumer1 = 2 + }; + + // Kernel level shared memory storage + SharedStorage& shared_storage = *reinterpret_cast(smem_buf); + + int thread_idx = int(threadIdx.x); + int warp_idx = canonical_warp_idx(); + int warp_group_thread_idx = thread_idx % NumThreadsPerWarpGroup; + auto warp_group_role = WarpGroupRole(canonical_warp_group_idx()); + int lane_predicate = cute::elect_one_sync(); + + // Issue Tma Descriptor Prefetch from a single thread + if ((warp_idx == 0) && lane_predicate) { + CollectiveMainloop::prefetch_tma_descriptors(params.mainloop); + CollectiveEpilogue::prefetch_tma_descriptors(params.epilogue); + } + + // Mainloop Load pipeline + using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline; + typename MainloopPipeline::Params mainloop_pipeline_params; + if (warp_group_role == WarpGroupRole::Producer) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; + } + mainloop_pipeline_params.is_leader = warp_group_thread_idx == 0; + mainloop_pipeline_params.num_consumers = NumThreadsPerWarpGroup; + mainloop_pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytes; + MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params); + + // Epilogue Load pipeline + using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline; + typename EpiLoadPipeline::Params epi_load_pipeline_params; + if (warp_group_role == WarpGroupRole::Producer) { + epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer; + } + if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) { + epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer; + } + epi_load_pipeline_params.dst_blockid = cute::block_rank_in_cluster(); + epi_load_pipeline_params.producer_arv_count = 1; // 1 thread issues TMA load + epi_load_pipeline_params.consumer_arv_count = NumThreadsPerWarpGroup; + epi_load_pipeline_params.transaction_bytes = CollectiveEpilogue::TmaTransactionBytes; + EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params); + + // Epilogue Store pipeline + using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline; + typename EpiStorePipeline::Params epi_store_pipeline_params; + epi_store_pipeline_params.always_wait = true; + EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params); + + typename MathWarpGroupOrderBarrier::Params params_math_wg_order_barrier; + // DMA Load WG will not participate in these Ordered Barrier syncs + params_math_wg_order_barrier.group_id = canonical_warp_group_idx() - static_cast(WarpGroupRole::Consumer0); + params_math_wg_order_barrier.group_size = NumThreadsPerWarpGroup; // Number of threads / participants in a group + MathWarpGroupOrderBarrier math_wg_order_barrier(shared_storage.pipelines.math_wg_order, params_math_wg_order_barrier); + + // Initialize starting pipeline states for the collectives + // Epilogue store pipe is producer-only (consumer is TMA unit, waits via scoreboarding) + typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state; + typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state; + + // For the DMA Load (producer) we start with an opposite phase + // i.e., we skip all waits since we know that the buffer is indeed empty + PipelineState mainloop_pipe_producer_state = mctlass::make_producer_start_state(); + PipelineState epi_load_pipe_producer_state = mctlass::make_producer_start_state(); + PipelineState epi_store_pipe_producer_state = mctlass::make_producer_start_state(); + + auto cluster_wait_fn = [&] () { + // We need this to guarantee that the Pipeline init is visible + // To all producers and consumer thread blocks in the Cluster + if constexpr (size(ClusterShape{}) > 1) { + cute::cluster_arrive_relaxed(); + return [] () { cute::cluster_wait(); }; + } + else { + __syncthreads(); + return [] () {}; // do nothing + } + } (); + + // Separate out problem shape for convenience + // Optionally append _1s until problem shape is rank-4 in case its is only rank-3 (MNK) + auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{}); + auto M = get<0>(problem_shape_MNKL); + auto N = get<1>(problem_shape_MNKL); + auto K = get<2>(problem_shape_MNKL); + auto L = get<3>(problem_shape_MNKL); + + // TMA requires special handling of strides to deal with coord codomain mapping + // Represent the full tensors -- get these from TMA + Tensor mA_mkl = params.mainloop.tma_load_a.get_tma_tensor(make_shape(M,K,L)); // (m,k,l) + Tensor mB_nkl = params.mainloop.tma_load_b.get_tma_tensor(make_shape(N,K,L)); // (n,k,l) + + // Get the appropriate blocks for this thread block -- potential for thread block locality + TiledMma tiled_mma; + auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K) + + // Make tiled views, defer the slice + Tensor gA_mkl = local_tile(mA_mkl, blk_shape, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l) + Tensor gB_nkl = local_tile(mB_nkl, blk_shape, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l) + + // Get pipeline stage increments from tensor shapes + auto k_tile_count = size<3>(gA_mkl); + auto c_tile_count = CollectiveEpilogue::get_load_pipe_increment(blk_shape); + auto d_tile_count = CollectiveEpilogue::get_store_pipe_increment(blk_shape); + + detail::PersistentTileSchedulerSm90 scheduler; + + if (warp_group_role == WarpGroupRole::Consumer1) { + // Advance 2nd Math WG to the next work tile for the startup + scheduler.advance_to_next_work(); + // Advance 2nd Math WG pipeline states to the end of 1st Math WG + mainloop_pipe_consumer_state.advance(k_tile_count); + epi_load_pipe_consumer_state.advance(c_tile_count); + epi_store_pipe_producer_state.advance(d_tile_count); + } + auto work_tile_info = scheduler.get_current_work(params.scheduler); + + // In a warp specialized kernel, collectives expose data movement and compute operations separately + CollectiveMainloop collective_mainloop; + CollectiveEpilogue collective_epilogue{params.epilogue}; + + // Wait for all thread blocks in the Cluster + cluster_wait_fn(); + + if (warp_group_role == WarpGroupRole::Producer) { + mctlass::arch::warpgroup_reg_dealloc(); + + while (work_tile_info.is_valid_tile) { + // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape + auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl)); + auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl)); + auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); + auto blk_coord = make_coord(m_coord, n_coord, _, l_coord); + + // Slice with our work tile coordinates to construct mainloop tensor views + Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k) + Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k) + + auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA)); + + collective_mainloop.load( + mainloop_pipeline, + mainloop_pipe_producer_state, + gA, params.mainloop.tma_load_a, + gB, params.mainloop.tma_load_b, + k_tile_iter, k_tile_count, + thread_idx, + shared_storage.tensors.mainloop + ); + // Update starting pipeline state for the next tile + mainloop_pipe_producer_state.advance(k_tile_count); + + if (collective_epilogue.is_source_needed()) { + collective_epilogue.load( + epi_load_pipeline, + epi_load_pipe_producer_state, + problem_shape_MNKL, + blk_shape, + blk_coord, + tiled_mma, + warp_group_thread_idx, + shared_storage.tensors.epilogue + ); + // Update starting pipeline state for the next tile + epi_load_pipe_producer_state.advance(c_tile_count); + } + + // Get next work tile + scheduler.advance_to_next_work(); + work_tile_info = scheduler.get_current_work(params.scheduler); + } // Scheduler work fetch loop + + // Make sure all Consumer Warp Groups have been waited upon + collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state); + if (collective_epilogue.is_source_needed()) { + collective_epilogue.load_tail(epi_load_pipeline, epi_load_pipe_producer_state); + } + } // Producer Warp Group End + + else if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) { + mctlass::arch::warpgroup_reg_alloc(); + + while (work_tile_info.is_valid_tile) { + // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape + auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl)); + auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl)); + auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); + auto blk_coord = make_coord(m_coord, n_coord, _, l_coord); + + // Allocate the the accumulators for the (M,N) blk_shape + Tensor accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape)); // (MMA,MMA_M,MMA_N) + + // Order two Math WG's MMA one after the other, helps hide Epilogue + math_wg_order_barrier.wait(); + + collective_mainloop.mma( + mainloop_pipeline, + mainloop_pipe_consumer_state, + accumulators, + k_tile_count, + thread_idx, + shared_storage.tensors.mainloop, + params.mainloop + ); + + // Cue for next Math WG's MMA to start + math_wg_order_barrier.arrive(); + + // Make sure the math instructions are done and free buffers before entering the epilogue + collective_mainloop.mma_tail( + mainloop_pipeline, + mainloop_pipe_consumer_state, + k_tile_count + ); + // Update starting mainloop pipeline state for the next tile + mainloop_pipe_consumer_state.advance(k_tile_count * NumMmaWarpGroups); + + // Order two Math WG's Epilogue one after the other + math_wg_order_barrier.wait(); + + // Epilogue and write to gD + collective_epilogue.store( + epi_load_pipeline, + epi_load_pipe_consumer_state, + epi_store_pipeline, + epi_store_pipe_producer_state, + problem_shape_MNKL, + blk_shape, + blk_coord, + accumulators, + tiled_mma, + warp_group_thread_idx, + shared_storage.tensors.epilogue + ); + // Update starting load/store pipeline states for the next tile + epi_load_pipe_consumer_state.advance(c_tile_count * NumMmaWarpGroups); + epi_store_pipe_producer_state.advance(d_tile_count * NumMmaWarpGroups); + + // Wait for all TMA stores to complete + epi_store_pipeline.producer_tail(epi_store_pipe_producer_state); + + // Cue for next Math WG's Epilogue to start + math_wg_order_barrier.arrive(); + + // Get next work tile + scheduler.advance_to_next_work(NumMmaWarpGroups); + work_tile_info = scheduler.get_current_work(params.scheduler); + } // Scheduler work fetch loop + } // Consumer Warp Groups End + } +}; + +/////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass::gemm::kernel diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/sm90_tile_scheduler.hpp b/csrc/mctlass/include/mctlass/gemm/kernel/sm90_tile_scheduler.hpp new file mode 100644 index 0000000..3602ff0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/sm90_tile_scheduler.hpp @@ -0,0 +1,182 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/fast_math.h" +#include "mctlass/kernel_hardware_info.hpp" +#include "cute/layout.hpp" + +namespace mctlass::gemm::kernel::detail { + +/////////////////////////////////////////////////////////////////////////////// + +// Persistent Thread Block (TB) scheduler +class PersistentTileSchedulerSm90 { + // + // Data members + // + +private: + uint64_t current_work_linear_idx_{static_cast((int(blockIdx.x) * int(gridDim.y)) + int(blockIdx.y))}; + uint64_t grid_blocks_total_{static_cast(int(gridDim.x) * int(gridDim.y))}; + + struct WorkTileInfo { + int32_t M_idx = 0; + int32_t N_idx = 0; + int32_t L_idx = 0; + uint32_t is_valid_tile = false; + }; + + // + // Methods + // + +public: + + struct Params { + FastDivmodU64 divmod_batch_{}; + FastDivmodU64 divmod_grid_y_{}; + FastDivmodU64 divmod_blk_m_{}; + + uint64_t blocks_per_problem_ = 0; + }; + + template + static Params + to_underlying_arguments(ProblemShapeMNKL problem_shape_mnkl, TileShape tile_shape, ClusterShape cluster_shape) { + // We only need the tile and cluster shape during scheduler setup, so let FTAD do the magic + static_assert(is_static::value); + static_assert(is_static::value); + + // Round up to nearest multiple of cluster dim along each mode + auto [problem_blocks_m, problem_blocks_n, problem_blocks_l] = get_tiled_blk_shape_mnl( + problem_shape_mnkl, tile_shape, cluster_shape); + + return { + FastDivmodU64(problem_blocks_m * problem_blocks_n), + FastDivmodU64(size<1>(cluster_shape)), + FastDivmodU64(problem_blocks_m), + problem_blocks_m * problem_blocks_n * problem_blocks_l + }; + } + + PersistentTileSchedulerSm90() = default; + + MCTLASS_DEVICE + WorkTileInfo + get_current_work(Params const& scheduler_params) const { + // Map worker's linear index into the CTA tiled problem shape to the corresponding MNL indices + uint64_t work_idx_l, remainder; + scheduler_params.divmod_batch_(work_idx_l, remainder, current_work_linear_idx_); + + uint64_t blk_per_grid_dim, dontcare; + scheduler_params.divmod_grid_y_(blk_per_grid_dim, dontcare, remainder); + + uint64_t block_idx_m, block_idx_n; + scheduler_params.divmod_blk_m_(block_idx_n, block_idx_m, blk_per_grid_dim); + int32_t work_idx_m = static_cast(block_idx_m); + int32_t work_idx_n = static_cast((block_idx_n * gridDim.y) + blockIdx.y); + + return {work_idx_m, work_idx_n, static_cast(work_idx_l), current_work_linear_idx_ < scheduler_params.blocks_per_problem_}; + } + + MCTLASS_DEVICE + void + advance_to_next_work(uint32_t advance_count = 1) { + current_work_linear_idx_ += grid_blocks_total_ * advance_count; + } + + // Given the inputs, computes the total number of output blocks this problem will compute over + // Note that this is only the logical size of our grid, not the physical grid we will actually launch. + template + MCTLASS_HOST_DEVICE static + dim3 + get_tiled_blk_shape_mnl(ProblemShapeMNKL problem_shape_mnkl, BlockShape blk_shape, ClusterShape cluster_shape) { + // Across M and N is our Cluster tile, so we must round up the blocks to the nearest whole number of Cluster tiles + auto blk_m = cute::size(cute::ceil_div(cute::shape<0>(problem_shape_mnkl), cute::shape<0>(blk_shape))); + auto blk_n = cute::size(cute::ceil_div(cute::shape<1>(problem_shape_mnkl), cute::shape<1>(blk_shape))); + + // Round up to nearest multiple of cluster dim along each mode + int problem_blocks_m = round_up(blk_m, cute::size<0>(cluster_shape)); + int problem_blocks_n = round_up(blk_n, cute::size<1>(cluster_shape)); + + // Cluster tile does not span the batch mode, so no extra rounding up required for it + int problem_blocks_l = int(cute::size<3>(problem_shape_mnkl)); + return {uint32_t(problem_blocks_m), uint32_t(problem_blocks_n), uint32_t(problem_blocks_l)}; + } + + // Given the inputs, computes the physical grid we should launch. + template + MCTLASS_HOST_DEVICE static + dim3 + get_grid_shape(ProblemShapeMNKL problem_shape_mnk, BlockShape blk_shape, ClusterShape cluster_shape, KernelHardwareInfo hw_info) { + int const sm_count = hw_info.sm_count; + MCTLASS_TRACE_HOST("get_grid_shape(): Persistent schedule grid plan using SM count = " << sm_count); + // Compute the total number of output tiles our problem has + auto problem_shape_MNKL = append<4>(problem_shape_mnk, Int<1>{}); + auto [problem_blocks_m, problem_blocks_n, problem_blocks_l] = + get_tiled_blk_shape_mnl(problem_shape_MNKL, blk_shape, cluster_shape); + int problem_blocks_total = problem_blocks_m * problem_blocks_n * problem_blocks_l; + + dim3 launch_grid(1, cute::size<1>(cluster_shape), 1); + + // The else path is generic, however, we can avoid some divs if we know Cluster size is 1 + if constexpr (size(cluster_shape) == 1) { + launch_grid.x = std::min(sm_count, problem_blocks_total); + } + else { + /* + * Optimal grid size calculation is based on + * GH100: 8 GPCs, 72 TPCs (9 TPCs/GPC), 2 SMs/TPC, 144 SMs per full GPU + * Hence, maximum SMs per GPC = 18 + */ + constexpr int max_sm_per_gpc = 18; + // Provided SM count could possibly be less than the assumed maximum SMs per GPC + int const min_num_gpc = sm_count < max_sm_per_gpc ? 1 : sm_count / max_sm_per_gpc; + int const max_blk_occupancy_per_gpc = max_sm_per_gpc - (max_sm_per_gpc % size(cluster_shape)); + int blk_per_device = min_num_gpc * max_blk_occupancy_per_gpc; + + // The calculation below allows for larger grid size launch for different GPUs. + int const num_gpc_residual = sm_count < max_sm_per_gpc ? 0 : sm_count % max_sm_per_gpc; + int const max_blk_occupancy_per_residual_gpc = num_gpc_residual - (num_gpc_residual % size(cluster_shape)); + blk_per_device += max_blk_occupancy_per_residual_gpc; + + blk_per_device = sm_count < blk_per_device ? sm_count : blk_per_device; + + launch_grid.x = std::min( + blk_per_device / size<1>(cluster_shape), + problem_blocks_total / size<1>(cluster_shape)); + } + return launch_grid; + } +}; + +} // namespace mctlass::gemm::kernel::detail diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/sparse_gemm.h b/csrc/mctlass/include/mctlass/gemm/kernel/sparse_gemm.h new file mode 100644 index 0000000..691cd78 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/sparse_gemm.h @@ -0,0 +1,400 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/semaphore.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + bool SplitKSerial ///! If true, code supporting split-K via serial reduction is enabled. +> +struct SparseGemm { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using OutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static bool const kSplitKSerial = SplitKSerial; + + static int const kSparse = Mma::kSparse; + static int const kMetaSizeInBits = Mma::kMetaSizeInBits; + static int const kMaxID2 = Mma::kMaxID2; + static int const kElementsPerElementE = Mma::kElementsPerElementE; + + using ElementE = typename Mma::ElementE; + using LayoutE = typename Mma::LayoutE; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Parameters structure + struct Params { + mctlass::gemm::GemmCoord problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorA::TensorRef ref_A; + typename Mma::IteratorB::Params params_B; + typename Mma::IteratorB::TensorRef ref_B; + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::TensorRef ref_C; + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::OutputTileIterator::TensorRef ref_D; + typename Mma::IteratorE::Params params_E; + typename Mma::IteratorE::TensorRef ref_E; + typename OutputOp::Params output_op; + int *semaphore; + int gemm_k_iterations; + int gemm_k_size; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): swizzle_log_tile(0), semaphore(0), gemm_k_iterations(0), gemm_k_size(0) { } + + MCTLASS_HOST_DEVICE + Params( + mctlass::gemm::GemmCoord const & problem_size, + mctlass::gemm::GemmCoord const & grid_tiled_shape, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::OutputTileIterator::TensorRef ref_C, + typename Epilogue::OutputTileIterator::TensorRef ref_D, + typename Mma::IteratorE::TensorRef ref_E, + typename OutputOp::Params output_op = typename OutputOp::Params(), + int *workspace = nullptr + ): + problem_size(problem_size), + grid_tiled_shape(grid_tiled_shape), + swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)), + params_A(ref_A.layout()), + ref_A(ref_A), + params_B(ref_B.layout()), + ref_B(ref_B), + params_C(ref_C.layout()), + ref_C(ref_C), + params_D(ref_D.layout()), + ref_D(ref_D), + params_E(ref_E.layout()), + ref_E(ref_E), + output_op(output_op) { + + int total_gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + int gemm_k_iterations = (total_gemm_k_iterations + grid_tiled_shape.k() - 1) / grid_tiled_shape.k(); + + gemm_k_size = gemm_k_iterations * Mma::Shape::kK; + + semaphore = workspace; + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + SparseGemm() { } + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::OutputTileIterator::TensorRef ref_C, + typename Epilogue::OutputTileIterator::TensorRef ref_D, + typename Mma::IteratorE::TensorRef ref_E) { + + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + static int const kAlignmentE = Mma::IteratorE::AccessType::kElements; + + if (!TensorRef_aligned(ref_A, kAlignmentA)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_B, kAlignmentB)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_C, kAlignmentC)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_D, kAlignmentC)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_E, kAlignmentE)) { + return Status::kErrorMisalignedOperand; + } + + if ((problem_size.m() % kAlignmentA) || ((problem_size.k() / kSparse) % kAlignmentA) || + (problem_size.n() % kAlignmentB) || (problem_size.k() % kAlignmentB) || + (problem_size.m() % kAlignmentC) || (problem_size.n() % kAlignmentC) || + (problem_size.m() % kAlignmentE) || ((problem_size.k() / kSparse) % kAlignmentE)) { + + return Status::kErrorMisalignedOperand; + } + + // The k dimension has to be the multiple of the Threadblock k because out + // of bound meta data would be initialized to 0 by acync.zfill but 0 is not + // a valid meta data. + if (problem_size.k() % Mma::Shape::kK) { + return Status::kErrorMisalignedOperand; + } + + // M dimension has to be multiple of 32 (sparse float) or 16 (sparse int) + // because of the row reordering of operand E + static int const kAlignmentM = (sizeof(ElementE) == 2) ? 32 : 16; + + if (problem_size.m() % kAlignmentM) { + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.k() * params.gemm_k_size / kSparse, + }; + + mctlass::MatrixCoord tb_offset_B{ + threadblock_tile_offset.k() * params.gemm_k_size, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + mctlass::MatrixCoord tb_offset_E{ + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.k() * params.gemm_k_size / kSparse, + }; + + // Problem size is a function of threadblock index in the K dimension + int problem_size_k = min( + params.problem_size.k(), + (threadblock_tile_offset.k() + 1) * params.gemm_k_size); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - tb_offset_B.row() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A, B, and E operands + typename Mma::IteratorA iterator_A( + params.params_A, + params.ref_A.data(), + {params.problem_size.m(), problem_size_k / kSparse}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + params.params_B, + params.ref_B.data(), + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B); + + typename Mma::IteratorE iterator_E( + params.params_E, params.ref_E.data(), + {params.problem_size.m(), + problem_size_k / kSparse / kElementsPerElementE}, + thread_idx, tb_offset_E); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + if (!kSplitKSerial || gemm_k_iterations > 0) { + // Compute threadblock-scoped matrix multiply-add + mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, iterator_E, accumulators); + } + + // + // Epilogue + // + + OutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + // If performing a reduction via split-K, fetch the initial synchronization + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + params.ref_C.data(), + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + params.ref_D.data(), + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + + __threadfence(); + } + + // Execute the epilogue operator to update the destination tensor. + epilogue(output_op, iterator_D, accumulators, iterator_C); + + // + // Release the semaphore + // + + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + __threadfence(); + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/sparse_gemm_row_broadcast.h b/csrc/mctlass/include/mctlass/gemm/kernel/sparse_gemm_row_broadcast.h new file mode 100644 index 0000000..acf6e83 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/sparse_gemm_row_broadcast.h @@ -0,0 +1,400 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/semaphore.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + bool SplitKSerial ///! If true, code supporting split-K via serial reduction is enabled. +> +struct SparseGemmRowBroadcast { + + using Mma = Mma_; + using Epilogue = Epilogue_; + using OutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + static bool const kSplitKSerial = SplitKSerial; + + static int const kSparse = Mma::kSparse; + static int const kMetaSizeInBits = Mma::kMetaSizeInBits; + static int const kMaxID2 = Mma::kMaxID2; + static int const kElementsPerElementE = Mma::kElementsPerElementE; + + using ElementE = typename Mma::ElementE; + using LayoutE = typename Mma::LayoutE; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 32 * WarpCount::kCount; + + /// Parameters structure + struct Params { + mctlass::gemm::GemmCoord problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorA::TensorRef ref_A; + typename Mma::IteratorB::Params params_B; + typename Mma::IteratorB::TensorRef ref_B; + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::TensorRef ref_C; + typename Epilogue::OutputTileIterator::Params params_D; + typename Epilogue::OutputTileIterator::TensorRef ref_D; + typename Mma::IteratorE::Params params_E; + typename Mma::IteratorE::TensorRef ref_E; + typename OutputOp::Params output_op; + int *semaphore; + int gemm_k_iterations; + int gemm_k_size; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): swizzle_log_tile(0), semaphore(0), gemm_k_iterations(0), gemm_k_size(0) { } + + MCTLASS_HOST_DEVICE + Params( + mctlass::gemm::GemmCoord const & problem_size, + mctlass::gemm::GemmCoord const & grid_tiled_shape, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::OutputTileIterator::TensorRef ref_C, + typename Epilogue::OutputTileIterator::TensorRef ref_D, + typename Mma::IteratorE::TensorRef ref_E, + typename OutputOp::Params output_op = typename OutputOp::Params(), + int *workspace = nullptr + ): + problem_size(problem_size), + grid_tiled_shape(grid_tiled_shape), + swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)), + params_A(ref_A.layout()), + ref_A(ref_A), + params_B(ref_B.layout()), + ref_B(ref_B), + params_C(ref_C.layout()), + ref_C(ref_C), + params_D(ref_D.layout()), + ref_D(ref_D), + params_E(ref_E.layout()), + ref_E(ref_E), + output_op(output_op) { + + int total_gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK; + int gemm_k_iterations = (total_gemm_k_iterations + grid_tiled_shape.k() - 1) / grid_tiled_shape.k(); + + gemm_k_size = gemm_k_iterations * Mma::Shape::kK; + + semaphore = workspace; + } + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + SparseGemmRowBroadcast() { } + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size, + typename Mma::IteratorA::TensorRef ref_A, + typename Mma::IteratorB::TensorRef ref_B, + typename Epilogue::OutputTileIterator::TensorRef ref_C, + typename Epilogue::OutputTileIterator::TensorRef ref_D, + typename Mma::IteratorE::TensorRef ref_E) { + + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + static int const kAlignmentE = Mma::IteratorE::AccessType::kElements; + + if (!TensorRef_aligned(ref_A, kAlignmentA)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_B, kAlignmentB)) { + return Status::kErrorMisalignedOperand; + } + + // if (!TensorRef_aligned(ref_C, kAlignmentC)) { + // return Status::kErrorMisalignedOperand; + // } + + if (!TensorRef_aligned(ref_D, kAlignmentC)) { + return Status::kErrorMisalignedOperand; + } + + if (!TensorRef_aligned(ref_E, kAlignmentE)) { + return Status::kErrorMisalignedOperand; + } + + if ((problem_size.m() % kAlignmentA) || ((problem_size.k() / kSparse) % kAlignmentA) || + (problem_size.n() % kAlignmentB) || (problem_size.k() % kAlignmentB) || + (problem_size.m() % kAlignmentC) || (problem_size.n() % kAlignmentC) || + (problem_size.m() % kAlignmentE) || ((problem_size.k() / kSparse) % kAlignmentE)) { + + return Status::kErrorMisalignedOperand; + } + + // The k dimension has to be the multiple of the Threadblock k because out + // of bound meta data would be initialized to 0 by acync.zfill but 0 is not + // a valid meta data. + if (problem_size.k() % Mma::Shape::kK) { + return Status::kErrorMisalignedOperand; + } + + // M dimension has to be multiple of 32 (sparse float) or 16 (sparse int) + // because of the row reordering of operand E + static int const kAlignmentM = (sizeof(ElementE) == 2) ? 32 : 16; + + if (problem_size.m() % kAlignmentM) { + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.k() * params.gemm_k_size / kSparse, + }; + + mctlass::MatrixCoord tb_offset_B{ + threadblock_tile_offset.k() * params.gemm_k_size, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + mctlass::MatrixCoord tb_offset_E{ + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.k() * params.gemm_k_size / kSparse, + }; + + // Problem size is a function of threadblock index in the K dimension + int problem_size_k = min( + params.problem_size.k(), + (threadblock_tile_offset.k() + 1) * params.gemm_k_size); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - tb_offset_B.row() + Mma::Shape::kK - 1) / Mma::Shape::kK; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Construct iterators to A, B, and E operands + typename Mma::IteratorA iterator_A( + params.params_A, + params.ref_A.data(), + {params.problem_size.m(), problem_size_k / kSparse}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + params.params_B, + params.ref_B.data(), + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B); + + typename Mma::IteratorE iterator_E( + params.params_E, params.ref_E.data(), + {params.problem_size.m(), + problem_size_k / kSparse / kElementsPerElementE}, + thread_idx, tb_offset_E); + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + if (!kSplitKSerial || gemm_k_iterations > 0) { + // Compute threadblock-scoped matrix multiply-add + mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, iterator_E, accumulators); + } + + // + // Epilogue + // + + OutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + // If performing a reduction via split-K, fetch the initial synchronization + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + params.ref_C.data(), + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + params.ref_D.data(), + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + + __threadfence(); + } + + // Execute the epilogue operator to update the destination tensor. + epilogue(output_op, iterator_D, accumulators, iterator_C); + + // + // Release the semaphore + // + + if (kSplitKSerial && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + __threadfence(); + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/symm_universal.h b/csrc/mctlass/include/mctlass/gemm/kernel/symm_universal.h new file mode 100644 index 0000000..495a6eb --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/symm_universal.h @@ -0,0 +1,699 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/semaphore.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma1_, ///! Threadblock-scoped triangular matrix multiply-accumulate (A*B or B*A) + typename Mma2_, ///! Threadblock-scoped triangular matrix multiply-accumulate (AT*B or B*AT) + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + SideMode SideMode_, ///! Side Mode for the kernel (kLeft or kRight) + FillMode FillMode_ ///! Fill Mode for triangular matrix (kLower or kUpper) +> +struct SymmUniversal { +public: + + using Mma1 = Mma1_; + using Mma2 = Mma2_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma1::IteratorA::Element; + using ElementB = typename Mma1::IteratorB::Element; + + // Mma1 (TRMM - with diagonal: C_tmp = alpha * A * B) + using LayoutA = typename Mma1::IteratorA::Layout; + using LayoutBT = typename Mma1::IteratorB::Layout; + static ComplexTransform const kMma1TransformA = Mma1::kTransformA; + static ComplexTransform const kMma1TransformB = Mma1::kTransformB; + + // Mma2 (TRMM - withOUT diagonal: alpha * AT * B) + using LayoutB = typename Mma2::IteratorA::Layout; + using LayoutAT = typename Mma2::IteratorB::Layout; + static ComplexTransform const kMma2TransformA = Mma2::kTransformA; + static ComplexTransform const kMma2TransformB = Mma2::kTransformB; + + // Common type definitions for Mma1 and Mma2 + using Operator = typename Mma1::Operator; + using OperatorClass = typename Mma1::Operator::OperatorClass; + using ThreadblockShape = typename Mma1::Shape; + using WarpShape = typename Mma1::Operator::Shape; + using InstructionShape = typename Mma1::Policy::Operator::InstructionShape; + using ArchTag = typename Mma1::ArchTag; + + static int const kStages = Mma1::kStages; + static int const kAlignmentA = Mma1::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma1::IteratorB::AccessType::kElements; + + // Output related typedefinitions + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + static SideMode const kSideModeA = SideMode_; + static FillMode const kFillModeA = FillMode_; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma1::WarpCount; + // static int const kThreadCount = 32 * WarpCount::kCount; + static int const kThreadCount = 64 * WarpCount::kCount; + + + // + // Structures + // + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmUniversalMode mode; + GemmCoord problem_size; + int batch_count; + + typename EpilogueOutputOp::Params epilogue; + + void const * ptr_A; + void const * ptr_B; + void const * ptr_C; + void * ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + int64_t batch_stride_D; + + typename LayoutA::Stride::Index lda; + typename LayoutB::Stride::Index ldb; + typename LayoutC::Stride::Index ldc; + typename LayoutC::Stride::Index ldd; + + // + // Methods + // + + Arguments(): + mode(GemmUniversalMode::kGemm), + batch_count(1), + ptr_A(nullptr), ptr_B(nullptr), ptr_C(nullptr), ptr_D(nullptr) { } + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void const * ptr_C, + void * ptr_D, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_C, + int64_t batch_stride_D, + typename LayoutA::Stride::Index lda, + typename LayoutB::Stride::Index ldb, + typename LayoutC::Stride::Index ldc, + typename LayoutC::Stride::Index ldd + ): + mode(mode), + problem_size(problem_size), + batch_count(batch_count), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), + batch_stride_A(batch_stride_A), batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D), + lda(lda), ldb(ldb), ldc(ldc), ldd(ldd) { + + } + + /// Returns arguments for the transposed problem sizes + Arguments transposed_problem_size() const { + Arguments args(*this); + + std::swap(args.problem_size.m(), args.problem_size.n()); + + return args; + } + + /// Returns arguments for the transposed matrices + Arguments swapped_matrices() const { + Arguments args(*this); + + std::swap(args.ptr_A, args.ptr_B); + std::swap(args.lda, args.ldb); + std::swap(args.batch_stride_A, args.batch_stride_B); + + return args; + } + }; + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params { + + mctlass::gemm::GemmCoord problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + + // Mma1 Iterator A and B params + typename Mma1::IteratorA::Params params_A_mma1; + typename Mma1::IteratorB::Params params_B_mma1; + + // Mma2 Iterator A and B params + typename Mma2::IteratorA::Params params_A_mma2; + typename Mma2::IteratorB::Params params_B_mma2; + + typename Epilogue::OutputTileIterator::Params params_C; + typename Epilogue::OutputTileIterator::Params params_D; + + typename EpilogueOutputOp::Params output_op; + + GemmUniversalMode mode; + int batch_count; + int gemm_k_size; + + void * ptr_A; + void * ptr_B; + void * ptr_C; + void * ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_C; + int64_t batch_stride_D; + + int *semaphore; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + swizzle_log_tile(0), + params_A_mma1(0), + params_B_mma1(0), + params_A_mma2(0), + params_B_mma2(0), + params_C(0), + params_D(0), + batch_count(0), + gemm_k_size(0), + mode(mctlass::gemm::GemmUniversalMode::kGemm), + ptr_A(nullptr), + ptr_B(nullptr), + ptr_C(nullptr), + ptr_D(nullptr), + batch_stride_A(0), + batch_stride_B(0), + batch_stride_C(0), + batch_stride_D(0), + semaphore(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + Arguments const &args, + mctlass::gemm::GemmCoord const & grid_tiled_shape, + int gemm_k_size, + void *workspace = nullptr + ): + problem_size(args.problem_size), + grid_tiled_shape(grid_tiled_shape), + swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)), + params_A_mma1(args.lda), + params_B_mma1(args.ldb), + params_A_mma2(args.lda), + params_B_mma2(args.ldb), + params_C(args.ldc), + params_D(args.ldd), + output_op(args.epilogue), + mode(args.mode), + batch_count(args.batch_count), + gemm_k_size(gemm_k_size), + ptr_A(const_cast(args.ptr_A)), + ptr_B(const_cast(args.ptr_B)), + ptr_C(const_cast(args.ptr_C)), + ptr_D(const_cast(args.ptr_D)), + batch_stride_A(args.batch_stride_A), + batch_stride_B(args.batch_stride_B), + batch_stride_C(args.batch_stride_C), + batch_stride_D(args.batch_stride_D), + semaphore(static_cast(workspace)) { + } + + MCTLASS_HOST_DEVICE + void update( + Arguments const &args, + void *workspace = nullptr) { + + ptr_A = const_cast(args.ptr_A); + ptr_B = const_cast(args.ptr_B); + ptr_C = const_cast(args.ptr_C); + ptr_D = args.ptr_D; + + output_op = args.epilogue; + + semaphore = static_cast(workspace); + } + + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma1::SharedStorage mma1_main_loop; + typename Mma2::SharedStorage mma2_main_loop; + typename Epilogue::SharedStorage epilogue; + }; + +public: + + // + // Methods + // + + MCTLASS_DEVICE + SymmUniversal() { } + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size) { + + static int const kAlignmentA = Mma1::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma1::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + if ((problem_size.m() % kAlignmentA) || (problem_size.k() % kAlignmentA) || + (problem_size.n() % kAlignmentB) || (problem_size.k() % kAlignmentB) || + (problem_size.m() % kAlignmentC) || (problem_size.n() % kAlignmentC)) { + + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + + /// Executes two GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + return; + } + + int offset_k = 0; + int problem_size_k = params.problem_size.k(); + + ElementA *ptr_A = static_cast(params.ptr_A); + ElementB *ptr_B = static_cast(params.ptr_B); + + // + // Fetch pointers based on mode. + // + if (params.mode == GemmUniversalMode::kGemm || + params.mode == GemmUniversalMode::kGemmSplitKParallel) { + + if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) { + + problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size; + } + + offset_k = threadblock_tile_offset.k() * params.gemm_k_size; + } + + __syncthreads(); + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_MxK_mma1{ + threadblock_tile_offset.m() * Mma1::Shape::kM, + offset_k, + }; + + mctlass::MatrixCoord tb_offset_KxN_mma1{ + offset_k, + threadblock_tile_offset.n() * Mma1::Shape::kN + }; + + mctlass::MatrixCoord tb_offset_MxK_mma2{ + threadblock_tile_offset.m() * Mma1::Shape::kM, + offset_k, + }; + + mctlass::MatrixCoord tb_offset_KxN_mma2{ + offset_k, + threadblock_tile_offset.n() * Mma1::Shape::kN + }; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply for Mma1 + Mma1 mma1(shared_storage.mma1_main_loop, thread_idx, warp_idx, lane_idx); + + // Construct thread-scoped matrix multiply for Mma2 + Mma2 mma2(shared_storage.mma2_main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma1::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - offset_k + Mma1::Shape::kK - 1) / Mma1::Shape::kK; + int gemm_k_iterations_mma1 = gemm_k_iterations; + int gemm_k_iterations_mma2 = gemm_k_iterations; + + + /****************************************************************************************************** + * SYMM (Side Mode, Fill Mode) is made of two TRMMs: + First TRMM (Mma1: Side Mode, Fill Mode, Non-Unit Diag): (A * B) or (B * A) + Second TRMM (Mma2: Side Mode, Inverted Fill Mode, Unit Diag): (AT * B) or (B * AT) + + * For the first TRMM (Mma1) of SYMM, the following method is used to calculate the k-iterations: + First two cases: (Left Side, Lower Fill) and (Right Side, Upper Fill) are transpose of each other + - (Left Side, Lower Fill): calculate bottom of the CTA tile, then find the k-iterations + needed to process all elements till that coordinate. + - (Right Side, Upper Fill): calculate right end of the CTA tile, then find the k-iterations + needed to process all elements till that coordinate. + + Last two cases: (Left Side, Upper Fill) and (Right Side, Lower Fill) are transpose of each other + - (Left Side, Upper Fill): calculate the top of the CTA tile, then find k-iterations + that can be skipped for all elements of this tile. + - (Right Side, Lower Fill): calculate the left start of the CTA tile, then find k-iterations + that can be skipped for all elements of this tile. + + * For the second TRMM (Mma2) of SYMM, the k-iterations and threadblock offsets are calculated + the same way as the first TRMM (Mma1) of same side mode but with inverted fill mode. + For example, if the first TRMM is left sided with lower fill, the second TRMM would be + left sided with upper fill. + ********************************************************************************************************/ + + if (kSideModeA == SideMode::kLeft && kFillModeA == FillMode::kLower) { + + int k_iterations_till_diagonal_mma1 = ((threadblock_tile_offset.m() + 1) * Mma1::Shape::kM + Mma1::Shape::kK - 1) / Mma1::Shape::kK; + if (k_iterations_till_diagonal_mma1 < gemm_k_iterations) { + gemm_k_iterations_mma1 = k_iterations_till_diagonal_mma1; + } + + int k_iterations_till_diagonal_mma2 = ((threadblock_tile_offset.m()) * Mma1::Shape::kM) / Mma1::Shape::kK; + if (k_iterations_till_diagonal_mma2 != 0) { + tb_offset_MxK_mma2 += mctlass::MatrixCoord({0, k_iterations_till_diagonal_mma2 * Mma1::Shape::kK}); + tb_offset_KxN_mma2 += mctlass::MatrixCoord({k_iterations_till_diagonal_mma2 * Mma1::Shape::kK, 0}); + gemm_k_iterations_mma2 -= k_iterations_till_diagonal_mma2; + } + + } else if (kSideModeA == SideMode::kRight && kFillModeA == FillMode::kUpper) { + + int k_iterations_till_diagonal_mma1 = ((threadblock_tile_offset.n() + 1) * Mma1::Shape::kN + Mma1::Shape::kK - 1) / Mma1::Shape::kK; + if (k_iterations_till_diagonal_mma1 < gemm_k_iterations) { + gemm_k_iterations_mma1 = k_iterations_till_diagonal_mma1; + } + + int k_iterations_till_diagonal_mma2 = ((threadblock_tile_offset.n()) * Mma1::Shape::kN) / Mma1::Shape::kK; + if (k_iterations_till_diagonal_mma2 != 0) { + tb_offset_MxK_mma2 += mctlass::MatrixCoord({0, k_iterations_till_diagonal_mma2 * Mma1::Shape::kK}); + tb_offset_KxN_mma2 += mctlass::MatrixCoord({k_iterations_till_diagonal_mma2 * Mma1::Shape::kK, 0}); + gemm_k_iterations_mma2 -= k_iterations_till_diagonal_mma2; + } + + } else if (kSideModeA == SideMode::kLeft && kFillModeA == FillMode::kUpper) { + + int k_iterations_till_diagonal_mma1 = ((threadblock_tile_offset.m()) * Mma1::Shape::kM) / Mma1::Shape::kK; + if (k_iterations_till_diagonal_mma1 != 0) { + tb_offset_MxK_mma1 += mctlass::MatrixCoord({0, k_iterations_till_diagonal_mma1 * Mma1::Shape::kK}); + tb_offset_KxN_mma1 += mctlass::MatrixCoord({k_iterations_till_diagonal_mma1 * Mma1::Shape::kK, 0}); + gemm_k_iterations_mma1 -= k_iterations_till_diagonal_mma1; + } + + int k_iterations_till_diagonal_mma2 = ((threadblock_tile_offset.m() + 1) * Mma1::Shape::kM + Mma1::Shape::kK - 1) / Mma1::Shape::kK; + if (k_iterations_till_diagonal_mma2 < gemm_k_iterations) { + gemm_k_iterations_mma2 = k_iterations_till_diagonal_mma2; + } + + } else if (kSideModeA == SideMode::kRight && kFillModeA == FillMode::kLower) { + + int k_iterations_till_diagonal_mma1 = ((threadblock_tile_offset.n()) * Mma1::Shape::kN) / Mma1::Shape::kK; + + if (k_iterations_till_diagonal_mma1 != 0) { + tb_offset_MxK_mma1 += mctlass::MatrixCoord({0, k_iterations_till_diagonal_mma1 * Mma1::Shape::kK}); + tb_offset_KxN_mma1 += mctlass::MatrixCoord({k_iterations_till_diagonal_mma1 * Mma1::Shape::kK, 0}); + gemm_k_iterations_mma1 -= k_iterations_till_diagonal_mma1; + } + + int k_iterations_till_diagonal_mma2 = ((threadblock_tile_offset.n() + 1) * Mma1::Shape::kN + Mma1::Shape::kK - 1) / Mma1::Shape::kK; + if (k_iterations_till_diagonal_mma2 < gemm_k_iterations) { + gemm_k_iterations_mma2 = k_iterations_till_diagonal_mma2; + } + + } + + // Construct iterators to A and B operands for Mma1 + typename Mma1::IteratorA iterator_A_mma1( + params.params_A_mma1, + ptr_A, + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_MxK_mma1); + + typename Mma1::IteratorB iterator_B_mma1( + params.params_B_mma1, + ptr_B, + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_KxN_mma1); + + // Construct iterators to A and B operands for Mma2 + typename Mma2::IteratorA iterator_A_mma2( + params.params_A_mma2, + ptr_A, + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_MxK_mma2); + + typename Mma2::IteratorB iterator_B_mma2( + params.params_B_mma2, + ptr_B, + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_KxN_mma2); + + // Compute threadblock-scoped matrix multiply-add (A x B) or (B x A) + mma1( + gemm_k_iterations_mma1, + accumulators, + iterator_A_mma1, + iterator_B_mma1, + accumulators); + + // Compute threadblock-scoped matrix multiply-add (AT x B) or (B x AT) + mma2( + gemm_k_iterations_mma2, + accumulators, + iterator_A_mma2, + iterator_B_mma2, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = + threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma1::Shape::kM, + threadblock_tile_offset.n() * Mma1::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + ElementC *ptr_C = static_cast(params.ptr_C); + ElementC *ptr_D = static_cast(params.ptr_D); + + // + // Fetch pointers based on mode. + // + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + if (params.mode == GemmUniversalMode::kGemm) { + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + } + else if (params.mode == GemmUniversalMode::kGemmSplitKParallel) { + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_C += threadblock_tile_offset.k() * params.batch_stride_C; + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_C = static_cast(params.ptr_C)[threadblock_tile_offset.k()]; + ptr_D = static_cast(params.ptr_D)[threadblock_tile_offset.k()]; + } + + // Tile iterator loading from source tensor. + typename Epilogue::OutputTileIterator iterator_C( + params.params_C, + ptr_C, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + + __threadfence(); + } + + // Execute the epilogue operator to update the destination tensor. + epilogue( + output_op, + iterator_D, + accumulators, + iterator_C); + + // + // Release the semaphore + // + + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/kernel/trmm_universal.h b/csrc/mctlass/include/mctlass/gemm/kernel/trmm_universal.h new file mode 100644 index 0000000..80b5040 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/kernel/trmm_universal.h @@ -0,0 +1,599 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief + +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/fast_math.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/complex.h" +#include "mctlass/semaphore.h" +#include "mctlass/core_io.h" +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Mma_, ///! Threadblock-scoped matrix multiply-accumulate + typename Epilogue_, ///! Epilogue + typename ThreadblockSwizzle_, ///! Threadblock swizzling function + SideMode SideMode_, ///! Side Mode for the kernel (kLeft or kRight) + FillMode FillMode_, ///! Fill Mode for triangular matrix (kLower or kUpper) + DiagType DiagType_ ///! Diag Type for triangular matrix (kNonUnit or kUnit) +> +struct TrmmUniversal { +public: + + using Mma = Mma_; + using Epilogue = Epilogue_; + using EpilogueOutputOp = typename Epilogue::OutputOp; + using ThreadblockSwizzle = ThreadblockSwizzle_; + + using ElementA = typename Mma::IteratorA::Element; + using LayoutA = typename Mma::IteratorA::Layout; + using ElementB = typename Mma::IteratorB::Element; + using LayoutB = typename Mma::IteratorB::Layout; + using ElementC = typename Epilogue::OutputTileIterator::Element; + using LayoutC = typename Epilogue::OutputTileIterator::Layout; + static SideMode const kSideMode = SideMode_; + static FillMode const kFillMode = FillMode_; + static DiagType const kDiagType = DiagType_; + + static ComplexTransform const kTransformA = Mma::kTransformA; + static ComplexTransform const kTransformB = Mma::kTransformB; + using Operator = typename Mma::Operator; + + using OperatorClass = typename Mma::Operator::OperatorClass; + using ThreadblockShape = typename Mma::Shape; + using WarpShape = typename Mma::Operator::Shape; + using InstructionShape = typename Mma::Policy::Operator::InstructionShape; + using ArchTag = typename Mma::ArchTag; + + static int const kStages = Mma::kStages; + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + /// Warp count (concept: GemmShape) + using WarpCount = typename Mma::WarpCount; + static int const kThreadCount = 64 * WarpCount::kCount; + + /// Split-K preserves splits that are 128b aligned + static int const kSplitKAlignment = const_max(128 / sizeof_bits::value, 128 / sizeof_bits::value); + + // + // Structures + // + + /// Argument structure + struct Arguments { + + // + // Data members + // + + GemmUniversalMode mode; + GemmCoord problem_size; + int batch_count; + + typename EpilogueOutputOp::Params epilogue; + + void const * ptr_A; + void const * ptr_B; + void * ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_D; + + typename LayoutA::Stride::Index lda; + typename LayoutB::Stride::Index ldb; + typename LayoutC::Stride::Index ldd; + + // + // Methods + // + + Arguments(): + mode(GemmUniversalMode::kGemm), + batch_count(1), + ptr_A(nullptr), ptr_B(nullptr), ptr_D(nullptr) { } + + /// constructs an arguments structure + Arguments( + GemmUniversalMode mode, + GemmCoord problem_size, + int batch_count, + typename EpilogueOutputOp::Params epilogue, + void const * ptr_A, + void const * ptr_B, + void * ptr_D, + int64_t batch_stride_A, + int64_t batch_stride_B, + int64_t batch_stride_D, + typename LayoutA::Stride::Index lda, + typename LayoutB::Stride::Index ldb, + typename LayoutC::Stride::Index ldd + ): + mode(mode), + problem_size(problem_size), + batch_count(batch_count), + epilogue(epilogue), + ptr_A(ptr_A), ptr_B(ptr_B), ptr_D(ptr_D), + batch_stride_A(batch_stride_A), batch_stride_B(batch_stride_B), batch_stride_D(batch_stride_D), + lda(lda), ldb(ldb), ldd(ldd) { + } + + /// Returns arguments for the transposed problem sizes + Arguments transposed_problem_size() const { + Arguments args(*this); + + std::swap(args.problem_size.m(), args.problem_size.n()); + + return args; + } + + /// Returns arguments for the transposed matrices + Arguments swapped_matrices() const { + Arguments args(*this); + + std::swap(args.ptr_A, args.ptr_B); + std::swap(args.lda, args.ldb); + std::swap(args.batch_stride_A, args.batch_stride_B); + + return args; + } + }; + + // + // Structure for precomputing values in host memory and passing to kernels + // + + /// Parameters structure + struct Params { + + mctlass::gemm::GemmCoord problem_size; + mctlass::gemm::GemmCoord grid_tiled_shape; + int swizzle_log_tile; + + typename Mma::IteratorA::Params params_A; + typename Mma::IteratorB::Params params_B; + typename Epilogue::OutputTileIterator::Params params_D; + + typename EpilogueOutputOp::Params output_op; + + GemmUniversalMode mode; + int batch_count; + int gemm_k_size; + + void * ptr_A; + void * ptr_B; + void * ptr_D; + + int64_t batch_stride_A; + int64_t batch_stride_B; + int64_t batch_stride_D; + + int *semaphore; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params(): + swizzle_log_tile(0), + params_A(0), + params_B(0), + params_D(0), + batch_count(0), + gemm_k_size(0), + mode(mctlass::gemm::GemmUniversalMode::kGemm), + ptr_A(nullptr), + ptr_B(nullptr), + ptr_D(nullptr), + batch_stride_A(0), + batch_stride_B(0), + batch_stride_D(0), + semaphore(nullptr) { } + + MCTLASS_HOST_DEVICE + Params( + Arguments const &args, + mctlass::gemm::GemmCoord const & grid_tiled_shape, + int gemm_k_size, + void *workspace = nullptr + ): + problem_size(args.problem_size), + grid_tiled_shape(grid_tiled_shape), + swizzle_log_tile(ThreadblockSwizzle().get_log_tile(grid_tiled_shape)), + params_A(args.lda), + params_B(args.ldb), + params_D(args.ldd), + output_op(args.epilogue), + mode(args.mode), + batch_count(args.batch_count), + gemm_k_size(gemm_k_size), + ptr_A(const_cast(args.ptr_A)), + ptr_B(const_cast(args.ptr_B)), + ptr_D(args.ptr_D), + batch_stride_A(args.batch_stride_A), + batch_stride_B(args.batch_stride_B), + batch_stride_D(args.batch_stride_D), + semaphore(static_cast(workspace)) { + } + + MCTLASS_HOST_DEVICE + void update( + Arguments const &args, + void *workspace = nullptr) { + + ptr_A = const_cast(args.ptr_A); + ptr_B = const_cast(args.ptr_B); + ptr_D = args.ptr_D; + + batch_stride_A = args.batch_stride_A; + batch_stride_B = args.batch_stride_B; + batch_stride_D = args.batch_stride_D; + + output_op = args.epilogue; + + semaphore = static_cast(workspace); + } + + }; + + /// Shared memory storage structure + union SharedStorage { + typename Mma::SharedStorage main_loop; + typename Epilogue::SharedStorage epilogue; + }; + +public: + + // + // Methods + // + + MCTLASS_DEVICE + TrmmUniversal() { } + + /// Determines whether kernel satisfies alignment + static Status can_implement( + mctlass::gemm::GemmCoord const & problem_size) { + + static int const kAlignmentA = Mma::IteratorA::AccessType::kElements; + static int const kAlignmentB = Mma::IteratorB::AccessType::kElements; + static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess; + + if ((problem_size.m() % kAlignmentA) || (problem_size.k() % kAlignmentA) || + (problem_size.n() % kAlignmentB) || (problem_size.k() % kAlignmentB) || + (problem_size.m() % kAlignmentC) || (problem_size.n() % kAlignmentC)) { + + return Status::kErrorMisalignedOperand; + } + + return Status::kSuccess; + } + + static Status can_implement(Arguments const &args) { + return can_implement(args.problem_size); + } + + /// Executes one GEMM + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + // Compute threadblock location + ThreadblockSwizzle threadblock_swizzle; + + mctlass::gemm::GemmCoord threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + // Early exit if CTA is out of range + if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() || + params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) { + + return; + } + + int offset_k = 0; + int problem_size_k = params.problem_size.k(); + + ElementA *ptr_A = static_cast(params.ptr_A); + ElementB *ptr_B = static_cast(params.ptr_B); + + // + // Fetch pointers based on mode. + // + if (params.mode == GemmUniversalMode::kGemm || + params.mode == GemmUniversalMode::kGemmSplitKParallel) { + + if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) { + + problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size; + } + + offset_k = threadblock_tile_offset.k() * params.gemm_k_size; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_A += threadblock_tile_offset.k() * params.batch_stride_A; + ptr_B += threadblock_tile_offset.k() * params.batch_stride_B; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_A = static_cast(params.ptr_A)[threadblock_tile_offset.k()]; + ptr_B = static_cast(params.ptr_B)[threadblock_tile_offset.k()]; + } + + __syncthreads(); + + // Compute initial location in logical coordinates + mctlass::MatrixCoord tb_offset_A{ + threadblock_tile_offset.m() * Mma::Shape::kM, + offset_k, + }; + + mctlass::MatrixCoord tb_offset_B{ + offset_k, + threadblock_tile_offset.n() * Mma::Shape::kN + }; + + // Compute position within threadblock + int thread_idx = threadIdx.x; + + // Broadcast the warp_id computed by lane 0 to ensure dependent code + // is compiled as warp-uniform. + int warp_idx = canonical_warp_idx(); + + int lane_idx = threadIdx.x % 64; + + // + // Main loop + // + + // Construct thread-scoped matrix multiply + Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx); + + typename Mma::FragmentC accumulators; + + accumulators.clear(); + + // Compute threadblock-scoped matrix multiply-add + int gemm_k_iterations = (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK; + + /****************************************************************************************************** + First two cases: (Left Side, Lower Fill) and (Right Side, Upper Fill) are transpose of each other + - (Left Side, Lower Fill): calculate bottom of the CTA tile, then find the k-iterations + needed to process all elements till that coordinate. + - (Right Side, Upper Fill): calculate right end of the CTA tile, then find the k-iterations + needed to process all elements till that coordinate. + + Last two cases: (Left Side, Upper Fill) and (Right Side, Lower Fill) are transpose of each other + - (Left Side, Upper Fill): calculate the top of the CTA tile, then find k-iterations + that can be skipped for all elements of this tile. + - (Right Side, Lower Fill): calculate the left start of the CTA tile, then find k-iterations + that can be skipped for all elements of this tile. + ********************************************************************************************************/ + + if (kSideMode == SideMode::kLeft && kFillMode == FillMode::kLower) { + + int k_iterations_till_diagonal = ((threadblock_tile_offset.m() + 1) * Mma::Shape::kM + Mma::Shape::kK - 1) / Mma::Shape::kK; + if (k_iterations_till_diagonal < gemm_k_iterations) { + gemm_k_iterations = k_iterations_till_diagonal; + } + + } else if (kSideMode == SideMode::kRight && kFillMode == FillMode::kUpper) { + + int k_iterations_till_diagonal = ((threadblock_tile_offset.n() + 1) * Mma::Shape::kN + Mma::Shape::kK - 1) / Mma::Shape::kK; + if (k_iterations_till_diagonal < gemm_k_iterations) { + gemm_k_iterations = k_iterations_till_diagonal; + } + + } else if (kSideMode == SideMode::kLeft && kFillMode == FillMode::kUpper) { + + int k_iterations_till_diagonal = ((threadblock_tile_offset.m()) * Mma::Shape::kM) / Mma::Shape::kK; + + if (k_iterations_till_diagonal != 0) { + tb_offset_A += mctlass::MatrixCoord({0, k_iterations_till_diagonal * Mma::Shape::kK}); + tb_offset_B += mctlass::MatrixCoord({k_iterations_till_diagonal * Mma::Shape::kK, 0}); + gemm_k_iterations -= k_iterations_till_diagonal; + } + + } else if (kSideMode == SideMode::kRight && kFillMode == FillMode::kLower) { + + int k_iterations_till_diagonal = ((threadblock_tile_offset.n()) * Mma::Shape::kN) / Mma::Shape::kK; + + if (k_iterations_till_diagonal != 0) { + tb_offset_A += mctlass::MatrixCoord({0, k_iterations_till_diagonal * Mma::Shape::kK}); + tb_offset_B += mctlass::MatrixCoord({k_iterations_till_diagonal * Mma::Shape::kK, 0}); + gemm_k_iterations -= k_iterations_till_diagonal; + } + + } + + // Construct iterators to A and B operands + typename Mma::IteratorA iterator_A( + params.params_A, + ptr_A, + {params.problem_size.m(), problem_size_k}, + thread_idx, + tb_offset_A); + + typename Mma::IteratorB iterator_B( + params.params_B, + ptr_B, + {problem_size_k, params.problem_size.n()}, + thread_idx, + tb_offset_B); + + // Compute threadblock-scoped matrix multiply-add + mma( + gemm_k_iterations, + accumulators, + iterator_A, + iterator_B, + accumulators); + + // + // Epilogue + // + + EpilogueOutputOp output_op(params.output_op); + + // + // Masked tile iterators constructed from members + // + + threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile); + + //assume identity swizzle + MatrixCoord threadblock_offset( + threadblock_tile_offset.m() * Mma::Shape::kM, + threadblock_tile_offset.n() * Mma::Shape::kN + ); + + int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m(); + + ElementC *ptr_D = static_cast(params.ptr_D); + + // + // Fetch pointers based on mode. + // + + // Construct the semaphore. + Semaphore semaphore(params.semaphore + block_idx, thread_idx); + + if (params.mode == GemmUniversalMode::kGemm) { + + // If performing a reduction via split-K, fetch the initial synchronization + if (params.grid_tiled_shape.k() > 1) { + + // Fetch the synchronization lock initially but do not block. + semaphore.fetch(); + + // Indicate which position in a serial reduction the output operator is currently updating + output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k()); + } + } + else if (params.mode == GemmUniversalMode::kGemmSplitKParallel) { + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kBatched) { + ptr_D += threadblock_tile_offset.k() * params.batch_stride_D; + } + else if (params.mode == GemmUniversalMode::kArray) { + ptr_D = static_cast(params.ptr_D)[threadblock_tile_offset.k()]; + } + + + // Tile iterator loading from source tensor (although irrelevant to this kernel as beta is zero). + typename Epilogue::OutputTileIterator iterator_C( + params.params_D, + ptr_D, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + // Tile iterator writing to destination tensor. + typename Epilogue::OutputTileIterator iterator_D( + params.params_D, + ptr_D, + params.problem_size.mn(), + thread_idx, + threadblock_offset + ); + + Epilogue epilogue( + shared_storage.epilogue, + thread_idx, + warp_idx, + lane_idx); + + // Wait on the semaphore - this latency may have been covered by iterator construction + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + // For subsequent threadblocks, the source matrix is held in the 'D' tensor. + if (threadblock_tile_offset.k()) { + iterator_C = iterator_D; + } + + semaphore.wait(threadblock_tile_offset.k()); + + __threadfence(); + } + + + // Execute the epilogue operator to update the destination tensor. + epilogue( + output_op, + iterator_D, + accumulators, + iterator_C); + + // + // Release the semaphore + // + + if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) { + + int lock = 0; + if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) { + + // The final threadblock resets the semaphore for subsequent grids. + lock = 0; + } + else { + // Otherwise, the semaphore is incremented + lock = threadblock_tile_offset.k() + 1; + } + + semaphore.release(lock); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/thread/mma.h b/csrc/mctlass/include/mctlass/gemm/thread/mma.h new file mode 100644 index 0000000..6410416 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/thread/mma.h @@ -0,0 +1,90 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates exposing architecture support for warp-level multiply-add operations +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/arch/mma.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape, + /// Data type of A elements + typename ElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Data type of B elements + typename ElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Element type of C matrix + typename ElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Concept: arch::OpMultiplyAdd or arch::Mma<> + typename Operator = arch::OpMultiplyAdd, + /// Used for partial specialization + typename Enable = bool +> +struct Mma; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Overloads specialized for existing architectures +// + +#include "mctlass/gemm/thread/mma_sm50.h" +#include "mctlass/gemm/thread/mma_sm60.h" +#include "mctlass/gemm/thread/mma_sm61.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/thread/mma_sm50.h b/csrc/mctlass/include/mctlass/gemm/thread/mma_sm50.h new file mode 100644 index 0000000..5ee533c --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/thread/mma_sm50.h @@ -0,0 +1,536 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates exposing architecture support for multiply-add operations +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/arch/mma.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/thread/mma.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Gemplate that handles all packed matrix layouts +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: layout::MapFunc) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: layout::MapFunc) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: layout::MapFunc) + typename LayoutC_, + /// Operator used to compute GEMM + typename Operator_ +> +struct MmaGeneric { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// Data type of operand A + using ElementA = ElementA_; + + /// Layout of A matrix (concept: layout::MapFunc) + using LayoutA = LayoutA_; + + /// Data type of operand B + using ElementB = ElementB_; + + /// Layout of B matrix (concept: layout::MapFunc) + using LayoutB = LayoutB_; + + /// Element type of operand C + using ElementC = ElementC_; + + /// Layout of C matrix (concept: layout::MapFunc) + using LayoutC = LayoutC_; + + /// Underlying mathematical operator + using Operator = Operator_; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Instruction + using MmaOp = arch::Mma< + gemm::GemmShape<1,1,1>, + 1, + ElementA, LayoutA, + ElementB, LayoutB, + ElementC, LayoutC, + Operator>; + + static bool const kMultipleOf2 = ((Shape::kM % 2 == 0) && (Shape::kN % 2 == 0)); + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + TensorRef a_ref( + reinterpret_cast(&A), LayoutA::packed({Shape::kM, Shape::kK})); + + TensorRef b_ref( + reinterpret_cast(&B), LayoutB::packed({Shape::kK, Shape::kN})); + + TensorRef d_ref( + reinterpret_cast(&D), LayoutC::packed(make_Coord(Shape::kM, Shape::kN))); + + MmaOp mma_op; + + // Copy accumulators + D = C; + + // Compute matrix product + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Shape::kK; ++k) { + //#if defined(__MACA_ARCH__) && (__TLASS_ARCH__ >= 860) + #if 0 + if (kMultipleOf2 && platform::is_same::value && platform::is_same::value && platform::is_same::value) { + + //2x2 zigzag - m and n loops to increment by 2. Inner loop to process 4 multiply-adds in a 2x2 tile. + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Shape::kN; n+=2) { + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Shape::kM; m+=2) { + + int m_serpentine = (n % 4) ? (Shape::kM - 2 - m) : m; + + //top-left element in 2x2 tile + { + MatrixCoord mn(m_serpentine, n); + MatrixCoord mk(m_serpentine, k); + MatrixCoord kn(k, n); + Array d; + Array a; + Array b; + d[0] = d_ref.at(mn); + a[0] = a_ref.at(mk); + b[0] = b_ref.at(kn); + mma_op(d, a, b, d); + d_ref.at(mn) = d[0]; + } + + //bottom-left element in 2x2 tile + { + MatrixCoord mn(m_serpentine+1, n); + MatrixCoord mk(m_serpentine+1, k); + MatrixCoord kn(k, n); + Array d; + Array a; + Array b; + d[0] = d_ref.at(mn); + a[0] = a_ref.at(mk); + b[0] = b_ref.at(kn); + mma_op(d, a, b, d); + d_ref.at(mn) = d[0]; + } + + //bottom-right element in 2x2 tile + { + MatrixCoord mn(m_serpentine+1, n+1); + MatrixCoord mk(m_serpentine+1, k); + MatrixCoord kn(k, n+1); + Array d; + Array a; + Array b; + d[0] = d_ref.at(mn); + a[0] = a_ref.at(mk); + b[0] = b_ref.at(kn); + mma_op(d, a, b, d); + d_ref.at(mn) = d[0]; + } + + //top-right element in 2x2 tile + { + MatrixCoord mn(m_serpentine, n+1); + MatrixCoord mk(m_serpentine, k); + MatrixCoord kn(k, n+1); + Array d; + Array a; + Array b; + d[0] = d_ref.at(mn); + a[0] = a_ref.at(mk); + b[0] = b_ref.at(kn); + mma_op(d, a, b, d); + d_ref.at(mn) = d[0]; + } + } + } + } else + #endif + { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Shape::kN; ++n) { + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Shape::kM; ++m) { + + int m_serpentine = (n % 2) ? (Shape::kM - 1 - m) : m; + + MatrixCoord mn(m_serpentine, n); + MatrixCoord mk(m_serpentine, k); + MatrixCoord kn(k, n); + + Array d; + Array a; + Array b; + + d[0] = d_ref.at(mn); + a[0] = a_ref.at(mk); + b[0] = b_ref.at(kn); + + mma_op(d, a, b, d); + + d_ref.at(mn) = d[0]; + } + } + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +/// Matrix multiply-add operation - assumes operand B is not changing +struct MmaComplexF32_Column { + + using Shape = gemm::GemmShape<1, 1, 1>; + using ElementC = complex; + + MCTLASS_HOST_DEVICE + void operator()( + Array, 1> &d, + Array, 1> const &a, + Array, 1> const &b, + Array, 1> const &c + ) { + + d[0].real() = a[0].real() * b[0].real() + c[0].real(); + d[0].imag() = a[0].real() * b[0].imag() + d[0].imag(); + d[0].real() = -a[0].imag() * b[0].imag() + d[0].real(); + d[0].imag() = a[0].imag() * b[0].real() + c[0].imag(); + } +}; + +/// Matrix multiply-add operation - assumes operand A is not changing +struct MmaComplexF32_Corner { + + using Shape = gemm::GemmShape<1, 1, 1>; + using ElementC = complex; + + MCTLASS_HOST_DEVICE + void operator()( + Array, 1> &d, + Array, 1> const &a, + Array, 1> const &b, + Array, 1> const &c + ) { + + d[0].real() = -a[0].imag() * b[0].imag() + d[0].real(); + d[0].imag() = a[0].real() * b[0].imag() + d[0].imag(); + d[0].real() = a[0].real() * b[0].real() + c[0].real(); + d[0].imag() = a[0].imag() * b[0].real() + c[0].imag(); + } +}; + +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Gemplate that handles all packed matrix layouts +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Layout of A matrix (concept: layout::MapFunc) + typename LayoutA_, + /// Layout of B matrix (concept: layout::MapFunc) + typename LayoutB_, + /// Layout of C matrix (concept: layout::MapFunc) + typename LayoutC_ +> +struct MmaGeneric< + Shape_, + complex, + LayoutA_, + complex, + LayoutB_, + complex, + LayoutC_, + arch::OpMultiplyAdd> { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// Data type of operand A + using ElementA = complex; + + /// Layout of A matrix (concept: layout::MapFunc) + using LayoutA = LayoutA_; + + /// Data type of operand B + using ElementB = complex; + + /// Layout of B matrix (concept: layout::MapFunc) + using LayoutB = LayoutB_; + + /// Element type of operand C + using ElementC = complex; + + /// Layout of C matrix (concept: layout::MapFunc) + using LayoutC = LayoutC_; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Instruction + using MmaOp = arch::Mma< + gemm::GemmShape<1,1,1>, + 1, + ElementA, LayoutA, + ElementB, LayoutB, + ElementC, LayoutC, + Operator>; + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + TensorRef a_ref( + reinterpret_cast(&A), LayoutA::packed({Shape::kM, Shape::kK})); + + TensorRef b_ref( + reinterpret_cast(&B), LayoutB::packed({Shape::kK, Shape::kN})); + + TensorRef d_ref( + reinterpret_cast(&D), LayoutC::packed(make_Coord(Shape::kM, Shape::kN))); + + detail::MmaComplexF32_Column mma_column; + detail::MmaComplexF32_Corner mma_corner; + + // Copy accumulators + D = C; + + // Compute matrix product + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Shape::kK; ++k) { + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Shape::kN; ++n) { + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Shape::kM; ++m) { + + int m_serpentine = (n % 2) ? (Shape::kM - 1 - m) : m; + + MatrixCoord mn(m_serpentine, n); + MatrixCoord mk(m_serpentine, k); + MatrixCoord kn(k, n); + + Array d; + Array a; + Array b; + + d[0] = d_ref.at(mn); + a[0] = a_ref.at(mk); + b[0] = b_ref.at(kn); + + if ((m == 0 && n) || m == Shape::kM - 1) { + mma_corner(d, a, b, d); + } + else { + mma_column(d, a, b, d); + } + + d_ref.at(mn) = d[0]; + } + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Gemplate that handles conventional layouts for FFMA and DFMA GEMM +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: layout::MapFunc) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: layout::MapFunc) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: layout::MapFunc) + typename LayoutC_ +> +struct Mma< + Shape_, + ElementA_, + LayoutA_, + ElementB_, + LayoutB_, + ElementC_, + LayoutC_, + arch::OpMultiplyAdd, + bool> { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// Data type of operand A + using ElementA = ElementA_; + + /// Layout of A matrix (concept: layout::MapFunc) + using LayoutA = LayoutA_; + + /// Data type of operand B + using ElementB = ElementB_; + + /// Layout of B matrix (concept: layout::MapFunc) + using LayoutB = LayoutB_; + + /// Element type of operand C + using ElementC = ElementC_; + + /// Layout of C matrix (concept: layout::MapFunc) + using LayoutC = LayoutC_; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename MmaGeneric< + Shape, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + Operator>::MmaOp; + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + MmaGeneric< + Shape, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + Operator> mma; + + mma(D, A, B, C); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/thread/mma_sm60.h b/csrc/mctlass/include/mctlass/gemm/thread/mma_sm60.h new file mode 100644 index 0000000..593609c --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/thread/mma_sm60.h @@ -0,0 +1,1178 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates exposing architecture support for multiply-add operations +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/thread/mma.h" +#include "mctlass/functional.h" +#include "mctlass/reduction/thread/reduce.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +/// Structure to compute the matrix product for HFMA +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape, + + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + + /// Type of GEMM inner vs outer product + bool +> +struct Mma_HFMA2; + + +///////////////////////////// +// Specialization for NNN // +///////////////////////////// + +template +struct Mma_HFMA2 < + Shape_, + layout::ColumnMajor, + layout::ColumnMajor, + layout::ColumnMajor, + true + > { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + static_assert( + !(Shape::kM % 2), + "Mma_HFMA2 requires the M dimension to be divisible by 2." + ); + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + /// Initialize output with input + D = C; + + /// Use 1x1x1 HFMA2 sequence for bulk of computation + using Mma = arch::Mma< + gemm::GemmShape<2,1,1>, + 1, + half_t, + layout::ColumnMajor, + half_t, + layout::ColumnMajor, + half_t, + layout::ColumnMajor, + arch::OpMultiplyAdd>; + + Array *ptr_D = reinterpret_cast *>(&D); + Array const *ptr_A = reinterpret_cast const *>(&A); + Array const *ptr_B = reinterpret_cast const *>(&B); + + Mma mma; + + MCTLASS_PRAGMA_UNROLL + for(auto k=0; k < Shape::kK / Mma::Shape::kK; k++){ + + MCTLASS_PRAGMA_UNROLL + for(auto m=0; m < Shape::kM / Mma::Shape::kM; m++){ + + MCTLASS_PRAGMA_UNROLL + for(auto n=0; n < Shape::kN / Mma::Shape::kN; n++){ + + Array tmp; + Array *ptr_tmp = &tmp; + ptr_tmp[0] = ptr_D[n*Shape::kM/2 + m]; + + mma( + tmp, + ptr_A[k*Shape::kM/2 + m], + ptr_B[n*Shape::kK + k], + tmp); + + ptr_D[n*Shape::kM/2 + m] = ptr_tmp[0]; + } + } + } + } +}; + +///////////////////////////// +// Specialization for NNT // +///////////////////////////// + +template +struct Mma_HFMA2< + Shape_, + layout::ColumnMajor, + layout::ColumnMajor, + layout::RowMajor, + true + > { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + static_assert( + !(Shape::kN % 2), + "Mma_HFMA2 requires the N dimension to be divisible by 2." + ); + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + /// Initialize output with input + D = C; + + /// Use 1x2x1 HFMA2 sequence for bulk of computation + using Mma = arch::Mma< + gemm::GemmShape<1,2,1>, + 1, + half_t, + layout::ColumnMajor, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + arch::OpMultiplyAdd>; + + Array *ptr_D = reinterpret_cast *>(&D); + Array const *ptr_A = reinterpret_cast const *>(&A); + Array const *ptr_B = reinterpret_cast const *>(&B); + + Mma mma; + + MCTLASS_PRAGMA_UNROLL + for(auto k=0; k < Shape::kK / Mma::Shape::kK; k++){ + + MCTLASS_PRAGMA_UNROLL + for(auto n=0; n < Shape::kN / Mma::Shape::kN; n++){ + + MCTLASS_PRAGMA_UNROLL + for(auto m=0; m < Shape::kM / Mma::Shape::kM; m++){ + + Array tmp; + Array *ptr_tmp = &tmp; + ptr_tmp[0] = ptr_D[m*Shape::kN/2 + n]; + + Array tmp_B; + tmp_B[0] = ptr_B->at(2*n*Shape::kK + k); + tmp_B[1] = ptr_B->at((2*n+1)*Shape::kK + k); + + mma( + tmp, + ptr_A[k*Shape::kM + m], + tmp_B, + tmp); + + ptr_D[m*Shape::kN/2 + n] = ptr_tmp[0]; + } + } + } + } +}; + + +///////////////////////////// +// Specialization for NTN // +///////////////////////////// + +template +struct Mma_HFMA2 < + Shape_, + layout::ColumnMajor, + layout::RowMajor, + layout::ColumnMajor, + true + > { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + static_assert( + !(Shape::kM % 2), + "Mma_HFMA2 requires the GEMM M dimension to be divisible by 2." + ); + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + /// Initialize output with input + D = C; + + using Mma = arch::Mma< + gemm::GemmShape<2,1,1>, + 1, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + arch::OpMultiplyAdd>; + + Array *ptr_D = reinterpret_cast *>(&D); + Array const *ptr_A = reinterpret_cast const *>(&A); + Array const *ptr_B = reinterpret_cast const *>(&B); + + Mma mma; + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Shape::kK / Mma::Shape::kK; ++k) { + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Shape::kM / Mma::Shape::kM; ++m) { + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Shape::kN / Mma::Shape::kN; ++n) { + + Array tmp; + Array *ptr_tmp = &tmp; + + ptr_tmp[0] = ptr_D[m + n * Shape::kM/2]; + + mma( + tmp, + ptr_A[m + k * Shape::kM/2], + ptr_B[k * Shape::kN + n], + tmp); + + ptr_D[m + n * Shape::kM/2] = ptr_tmp[0]; + } + } + } + } +}; + +///////////////////////////// +// Specialization for NTT // +///////////////////////////// + +template +struct Mma_HFMA2< + Shape_, + layout::ColumnMajor, + layout::RowMajor, + layout::RowMajor, + true + > { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + static_assert( + !(Shape::kN % 2), + "Mma_HFMA2 requires the N dimension to be divisible by 2." + ); + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + /// Initialize output with input + D = C; + + /// Use 1x2x1 HFMA2 sequence for bulk of computation + using Mma = arch::Mma< + gemm::GemmShape<1,2,1>, + 1, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + half_t, + layout::RowMajor, + arch::OpMultiplyAdd>; + + Array *ptr_D = reinterpret_cast *>(&D); + Array const *ptr_A = reinterpret_cast const *>(&A); + Array const *ptr_B = reinterpret_cast const *>(&B); + + Mma mma; + + MCTLASS_PRAGMA_UNROLL + for(auto k=0; k < Shape::kK / Mma::Shape::kK; k++){ + + MCTLASS_PRAGMA_UNROLL + for(auto n=0; n < Shape::kN / Mma::Shape::kN; n++){ + + MCTLASS_PRAGMA_UNROLL + for(auto m=0; m < Shape::kM / Mma::Shape::kM; m++){ + + Array tmp; + Array *ptr_tmp = &tmp; + ptr_tmp[0] = ptr_D[m*Shape::kN/2 + n]; + + mma( + tmp, + ptr_A[k*Shape::kM + m], + ptr_B[k*Shape::kN/2 + n], + tmp); + + ptr_D[m*Shape::kN/2 + n] = ptr_tmp[0]; + } + } + } + } +}; + + +///////////////////////////// +// Specialization for TNN // +///////////////////////////// + +template +struct Mma_HFMA2 < + Shape_, + layout::RowMajor, + layout::ColumnMajor, + layout::ColumnMajor, + true + > { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + static_assert( + !(Shape::kM % 2), + "Mma_HFMA2 requires the M dimension to be divisible by 2." + ); + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + /// Initialize output with input + D = C; + + /// Use 1x1x1 HFMA2 sequence for bulk of computation + using Mma = arch::Mma< + gemm::GemmShape<2,1,1>, + 1, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + half_t, + layout::ColumnMajor, + arch::OpMultiplyAdd>; + + Array *ptr_D = reinterpret_cast *>(&D); + Array const *ptr_A = reinterpret_cast const *>(&A); + Array const *ptr_B = reinterpret_cast const *>(&B); + + Mma mma; + + MCTLASS_PRAGMA_UNROLL + for(auto k=0; k < Shape::kK / Mma::Shape::kK; k++){ + + MCTLASS_PRAGMA_UNROLL + for(auto m=0; m < Shape::kM / Mma::Shape::kM; m++){ + + MCTLASS_PRAGMA_UNROLL + for(auto n=0; n < Shape::kN / Mma::Shape::kN; n++){ + + Array tmp; + Array *ptr_tmp = &tmp; + ptr_tmp[0] = ptr_D[n*Shape::kM/2 + m]; + + Array tmp_A; + tmp_A[0] = ptr_A->at(2*m*Shape::kK + k); + tmp_A[1] = ptr_A->at((2*m+1)*Shape::kK + k); + + mma( + tmp, + tmp_A, + ptr_B[n*Shape::kK + k], + tmp); + + ptr_D[n*Shape::kM/2 + m] = ptr_tmp[0]; + } + } + } + } +}; + +///////////////////////////// +// Specialization for TNT // +///////////////////////////// + +template +struct Mma_HFMA2 < + Shape_, + layout::RowMajor, + layout::ColumnMajor, + layout::RowMajor, + true + > { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + static_assert( + !(Shape::kN % 2), + "Mma_HFMA2 requires the N dimension to be divisible by 2." + ); + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + /// Initialize output with input + D = C; + + /// Use 1x2x1 HFMA2 sequence for bulk of computation + using Mma = arch::Mma< + gemm::GemmShape<1,2,1>, + 1, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + half_t, + layout::RowMajor, + arch::OpMultiplyAdd>; + + Array *ptr_D = reinterpret_cast *>(&D); + Array const *ptr_A = reinterpret_cast const *>(&A); + Array const *ptr_B = reinterpret_cast const *>(&B); + + Mma mma; + + MCTLASS_PRAGMA_UNROLL + for(auto k=0; k < Shape::kK / Mma::Shape::kK; k++){ + + MCTLASS_PRAGMA_UNROLL + for(auto n=0; n < Shape::kN / Mma::Shape::kN; n++){ + + MCTLASS_PRAGMA_UNROLL + for(auto m=0; m < Shape::kM / Mma::Shape::kM; m++){ + + Array tmp; + Array *ptr_tmp = &tmp; + ptr_tmp[0] = ptr_D[m*Shape::kN/2 + n]; + + Array tmp_B; + tmp_B[0] = ptr_B->at(2*n*Shape::kK + k); + tmp_B[1] = ptr_B->at((2*n+1)*Shape::kK + k); + + mma( + tmp, + ptr_A[m*Shape::kK + k], + tmp_B, + tmp); + + ptr_D[m*Shape::kN/2 + n] = ptr_tmp[0]; + } + } + } + } +}; + +///////////////////////////// +// Specialization for TTN // +///////////////////////////// + +template +struct Mma_HFMA2 < + Shape_, + layout::RowMajor, + layout::RowMajor, + layout::ColumnMajor, + true + > { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + static_assert( + !(Shape::kM % 2), + "Mma_HFMA2 requires the M dimension to be divisible by 2." + ); + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + /// Initialize output with input + D = C; + + /// Use 1x2x1 HFMA2 sequence for bulk of computation + using Mma = arch::Mma< + gemm::GemmShape<2,1,1>, + 1, + half_t, + layout::RowMajor, + half_t, + layout::RowMajor, + half_t, + layout::ColumnMajor, + arch::OpMultiplyAdd>; + + Array *ptr_D = reinterpret_cast *>(&D); + Array const *ptr_A = reinterpret_cast const *>(&A); + Array const *ptr_B = reinterpret_cast const *>(&B); + + Mma mma; + + MCTLASS_PRAGMA_UNROLL + for(auto k=0; k < Shape::kK / Mma::Shape::kK; k++){ + + MCTLASS_PRAGMA_UNROLL + for(auto m=0; m < Shape::kM / Mma::Shape::kM; m++){ + + MCTLASS_PRAGMA_UNROLL + for(auto n=0; n < Shape::kN / Mma::Shape::kN; n++){ + + Array tmp; + Array *ptr_tmp = &tmp; + ptr_tmp[0] = ptr_D[n*Shape::kM/2 + m]; + + Array tmp_A; + tmp_A[0] = ptr_A->at(2*m*Shape::kK + k); + tmp_A[1] = ptr_A->at((2*m+1)*Shape::kK + k); + + mma( + tmp, + tmp_A, + ptr_B[k*Shape::kN + n], + tmp); + + ptr_D[n*Shape::kM/2 + m] = ptr_tmp[0]; + } + } + } + } +}; + + +///////////////////////////// +// Specialization for TTT // +///////////////////////////// + +template +struct Mma_HFMA2< + Shape_, + layout::RowMajor, + layout::RowMajor, + layout::RowMajor, + true + > { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + static_assert( + !(Shape::kN % 2), + "Mma_HFMA2 requires the N dimension to be divisible by 2." + ); + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + /// Initialize output with input + D = C; + + /// Use 1x2x1 HFMA2 sequence for bulk of computation + using Mma = arch::Mma< + gemm::GemmShape<1,2,1>, + 1, + half_t, + layout::RowMajor, + half_t, + layout::RowMajor, + half_t, + layout::RowMajor, + arch::OpMultiplyAdd>; + + Array *ptr_D = reinterpret_cast *>(&D); + Array const *ptr_A = reinterpret_cast const *>(&A); + Array const *ptr_B = reinterpret_cast const *>(&B); + + Mma mma; + + MCTLASS_PRAGMA_UNROLL + for(auto k=0; k < Shape::kK / Mma::Shape::kK; k++){ + + MCTLASS_PRAGMA_UNROLL + for(auto n=0; n < Shape::kN / Mma::Shape::kN; n++){ + + MCTLASS_PRAGMA_UNROLL + for(auto m=0; m < Shape::kM / Mma::Shape::kM; m++){ + + Array tmp; + Array *ptr_tmp = &tmp; + ptr_tmp[0] = ptr_D[m*Shape::kN/2 + n]; + + mma( + tmp, + ptr_A[m*Shape::kK + k], + ptr_B[k*Shape::kN/2 + n], + tmp); + + ptr_D[m*Shape::kN/2 + n] = ptr_tmp[0]; + } + } + } + } +}; + +///////////////////////////////////////////////////////////////////// +// Specialization for TNT + Inner Product or 1x1x2K + LayoutC = T // +///////////////////////////////////////////////////////////////////// + +template +struct Mma_HFMA2< + Shape_, + LayoutA, + LayoutB, + layout::RowMajor, + false + > { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + static_assert( + !(Shape::kK % 2), + "Mma_HFMA2 requires the K dimension to be divisible by 2." + ); + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + /// Initialize output with input + D = C; + + /// Use 1x1x2 HFMA2 sequence for bulk of computation + using GemmShape = gemm::GemmShape<1,1,2>; + + Array *ptr_D = reinterpret_cast *>(&D); + Array const *ptr_A = reinterpret_cast const *>(&A); + Array const *ptr_B = reinterpret_cast const *>(&B); + + // Inner product is calculated using MACs, followed by final reduction + multiply_add> mac; + mctlass::reduction::thread::Reduce< plus, Array > reduce; + + MCTLASS_PRAGMA_UNROLL + for(auto n=0; n < Shape::kN / GemmShape::kN; n++){ + + MCTLASS_PRAGMA_UNROLL + for(auto m=0; m < Shape::kM / GemmShape::kM; m++){ + + Array tmp_C; + tmp_C.clear(); + Array *ptr_tmp_C = reinterpret_cast *>(&tmp_C); + ptr_tmp_C[0] = ptr_D[n*Shape::kM + m]; + + MCTLASS_PRAGMA_UNROLL + for(auto k=0; k < Shape::kK / GemmShape::kK; k++){ + tmp_C = mac(ptr_A[m*Shape::kK/2 + k], ptr_B[n*Shape::kK/2 + k], tmp_C); + } + + Array res; + Array *ptr_res = &res; + res = reduce(tmp_C); + + ptr_D[m*Shape::kN + n] = ptr_res[0]; + } + } + } +}; + +///////////////////////////////////////////////////////////////////// +// Specialization for TNN + Inner Product or 1x1x2K + LayoutC = N // +///////////////////////////////////////////////////////////////////// + +template +struct Mma_HFMA2< + Shape_, + LayoutA, + LayoutB, + layout::ColumnMajor, + false + > { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + static_assert( + !(Shape::kK % 2), + "Mma_HFMA2 requires the K dimension to be divisible by 2." + ); + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + /// Initialize output with input + D = C; + + /// Use 1x1x2 HFMA2 sequence for bulk of computation + using GemmShape= gemm::GemmShape<1,1,2>; + + Array *ptr_D = reinterpret_cast *>(&D); + Array const *ptr_A = reinterpret_cast const *>(&A); + Array const *ptr_B = reinterpret_cast const *>(&B); + + // Inner product is calculated using MACs, followed by final reduction + multiply_add> mac; + mctlass::reduction::thread::Reduce< plus, Array > reduce; + + MCTLASS_PRAGMA_UNROLL + for(auto n=0; n < Shape::kN / GemmShape::kN; n++){ + + MCTLASS_PRAGMA_UNROLL + for(auto m=0; m < Shape::kM / GemmShape::kM; m++){ + + Array tmp_C; + tmp_C.clear(); + Array *ptr_tmp_C = reinterpret_cast *>(&tmp_C); + ptr_tmp_C[0] = ptr_D[n*Shape::kM + m]; + + MCTLASS_PRAGMA_UNROLL + for(auto k=0; k < Shape::kK / GemmShape::kK; k++){ + + tmp_C = mac(ptr_A[m*Shape::kK/2 + k], ptr_B[n*Shape::kK/2 + k], tmp_C); + + } + + Array res; + Array *ptr_res = &res; + res = reduce(tmp_C); + + ptr_D[n*Shape::kM + m] = ptr_res[0]; + } + } + } +}; + +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, typename LayoutA, typename LayoutB, typename LayoutC +> +struct Mma< + Shape_, + half_t, + LayoutA, + half_t, + LayoutB, + half_t, + LayoutC, + arch::OpMultiplyAdd + > { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// Data type of operand A + using ElementA = half_t; + + /// Data type of operand B + using ElementB = half_t; + + /// Element type of operand C + using ElementC = half_t; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + static bool const a_row_major = platform::is_same< LayoutA, layout::RowMajor>::value; + static bool const b_column_major = platform::is_same< LayoutB, layout::ColumnMajor>::value; + static bool const c_row_major = platform::is_same< LayoutC, layout::RowMajor>::value; + static bool const c_column_major = platform::is_same< LayoutC, layout::ColumnMajor>::value; + + static bool const m_mod2 = !(Shape::kM % 2); + static bool const n_mod2 = !(Shape::kN % 2); + static bool const k_mod2 = !(Shape::kK % 2); + + // HFMA based MMA optimizations are of 2 types : + // 1. Inner product + // 2. Outer product + // It is chosen based on LayoutC (for outer product gemm) or + // Using LayoutA and LayoutB or shape=1x1x2K (for inner product gemms) + // If all fails, we choose the generic MMA + static bool const use_outer_prod = (c_column_major && m_mod2) || (c_row_major && n_mod2); + static bool const use_inner_prod = (a_row_major && b_column_major && k_mod2) || (Shape::kM==1 && Shape::kN==1 && k_mod2); + static bool const use_optimized = (use_outer_prod || use_inner_prod); + + using ArchMmaOperator = typename platform::conditional< use_optimized, + detail::Mma_HFMA2, + MmaGeneric + >::type; + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + ArchMmaOperator mma; + + mma(D, A, B, C); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + + /// Determines whether to enable thread::Gemm<> specializations compatible with SM50 + template < + typename LayoutA, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB> + struct EnableMma_Crow_SM60 { + + static bool const kIsConventionalLayout = + (platform::is_same::value || + platform::is_same::value) && + (platform::is_same::value || + platform::is_same::value); + + static bool const value = kIsConventionalLayout; + }; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Computes matrix product when C is row-major +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + typename LayoutA_, + typename LayoutB_ +> +struct Mma< + Shape_, + half_t, + LayoutA_, + half_t, + LayoutB_, + half_t, + layout::RowMajor, + arch::OpMultiplyAdd, + typename platform::enable_if::value>::type>{ + + using Shape = Shape_; + using ElementA = half_t; + using LayoutA = LayoutA_; + using ElementB = half_t; + using LayoutB = LayoutB_; + using ElementC = half_t; + using LayoutC = layout::RowMajor; + using Operator = arch::OpMultiplyAdd; + + using TransposeMma = Mma< + GemmShapeTranspose, + half_t, + typename layout::LayoutTranspose::type, + half_t, + typename layout::LayoutTranspose::type, + half_t, + layout::ColumnMajor, + arch::OpMultiplyAdd, + bool>; + + using FragmentA = Array; + using FragmentB = Array; + using FragmentC = Array; + + using ArchMmaOperator = typename TransposeMma::ArchMmaOperator; + + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + TransposeMma mma; + + mma(D, B, A, C); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/thread/mma_sm61.h b/csrc/mctlass/include/mctlass/gemm/thread/mma_sm61.h new file mode 100644 index 0000000..0b85f23 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/thread/mma_sm61.h @@ -0,0 +1,284 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates exposing architecture support for multiply-add operations +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/thread/mma.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Gemplate that handles conventional layouts for IDP4A +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_ +> +struct Mma< + Shape_, + int8_t, + layout::RowMajor, + int8_t, + layout::ColumnMajor, + int32_t, + LayoutC_, + arch::OpMultiplyAdd, + bool> { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// Data type of operand A + using ElementA = int8_t; + + /// Layout of A matrix (concept: layout::MapFunc) + using LayoutA = layout::RowMajor; + + /// Data type of operand B + using ElementB = int8_t; + + /// Layout of B matrix (concept: layout::MapFunc) + using LayoutB = layout::ColumnMajor; + + /// Element type of operand C + using ElementC = int32_t; + + /// Layout of C matrix (concept: layout::MapFunc) + using LayoutC = LayoutC_; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Underlying matrix multiply operator (concept: arch::Mma) + // Use 1x1x4 IDP4A sequence for bulk of computation + using ArchMmaOperator = arch::Mma< + gemm::GemmShape<1,1,4>, + 1, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + arch::OpMultiplyAdd>; + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + TensorRef d( + reinterpret_cast(&D), LayoutC::packed({ Shape::kM, Shape::kN })); + + // Copy accumulators + D = C; + + /// Use 1x1x4 IDP4A sequence for bulk of computation + ArchMmaOperator mma; + + // Compute matrix product + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Shape::kK / ArchMmaOperator::Shape::kK; ++k) { + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Shape::kN; ++n) { + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Shape::kM; ++m) { + MatrixCoord mn(m, n); + + Array const *ptr_A = reinterpret_cast const *>(&A); + Array const *ptr_B = reinterpret_cast const *>(&B); + + Array tmp = reinterpret_cast &>(d.at(mn)); + + mma( + tmp, + ptr_A[m * Shape::kK / ArchMmaOperator::Shape::kK + k], + ptr_B[n * Shape::kK / ArchMmaOperator::Shape::kK + k], + tmp); + + d.at(mn) = reinterpret_cast(tmp); + } + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Gemplate that handles conventional layouts for IDP4A +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_ +> +struct Mma< + Shape_, + int8_t, + layout::ColumnMajor, + int8_t, + layout::RowMajor, + int32_t, + LayoutC_, + arch::OpMultiplyAdd, + int8_t> { + + /// Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + /// Data type of operand A + using ElementA = int8_t; + + /// Layout of A matrix (concept: layout::MapFunc) + using LayoutA = layout::ColumnMajor; + + /// Data type of operand B + using ElementB = int8_t; + + /// Layout of B matrix (concept: layout::MapFunc) + using LayoutB = layout::RowMajor; + + /// Element type of operand C + using ElementC = int32_t; + + /// Layout of C matrix (concept: layout::MapFunc) + using LayoutC = LayoutC_; + + /// Underlying mathematical operator + using Operator = arch::OpMultiplyAdd; + + /// A operand storage + using FragmentA = Array; + + /// B operand storage + using FragmentB = Array; + + /// C operand storage + using FragmentC = Array; + + /// Underlying matrix multiply operator (concept: arch::Mma) + /// Use 1x1x4 IDP4A sequence for bulk of computation + using ArchMmaOperator = arch::Mma< + gemm::GemmShape<1,1,4>, + 1, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + arch::OpMultiplyAdd>; + + // + // Methods + // + + /// Computes a matrix product D = A * B + C + MCTLASS_HOST_DEVICE + void operator()( + FragmentC & D, + FragmentA const & A, + FragmentB const & B, + FragmentC const & C) { + + TensorRef d( + reinterpret_cast(&D), LayoutC::packed({ Shape::kM, Shape::kN })); + + // Copy accumulators + D = C; + + /// Underlying matrix multiply operator + ArchMmaOperator mma; + + Array const *ptr_A = reinterpret_cast const *>(&A); + Array const *ptr_B = reinterpret_cast const *>(&B); + + // Compute matrix product + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Shape::kK / ArchMmaOperator::Shape::kK; ++k) { + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Shape::kN; ++n) { + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Shape::kM; ++m) { + MatrixCoord mn(m, n); + + Array tmp = reinterpret_cast &>(d.at(mn)); + + mma( + tmp, + ptr_A[m + k * Shape::kM], + ptr_B[n + k * Shape::kN], + tmp); + + d.at(mn) = reinterpret_cast(tmp); + } + } + } + } +}; + +} // namespace thread +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_ell_mma.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_ell_mma.h new file mode 100644 index 0000000..b3a32ee --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_ell_mma.h @@ -0,0 +1,734 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Default template for a Blocked-Ell MMA. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator_2dthreadtile.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/gemm/threadblock/default_mma_core_wmma.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + +#include "mctlass/gemm/threadblock/ell_mma_pipelined.h" +#include "mctlass/gemm/threadblock/ell_mma_multistage.h" +#include "mctlass/transform/threadblock/ell_predicated_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false + > +struct DefaultEllMma; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output (OperatorClass Simt) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Operation performed by GEMM + typename Operator> +struct DefaultEllMma { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, + arch::OpClassSimt, 2, Operator>; + + // Define iterators over tiles from the A operand + using IteratorA = + mctlass::transform::threadblock::EllPredicatedTileIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, typename MmaCore::IteratorThreadMapA, kAlignmentA>; + + // Define iterators over tiles from the B operand + using IteratorB = + mctlass::transform::threadblock::EllPredicatedTileIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, typename MmaCore::IteratorThreadMapB, kAlignmentB>; + + // Define the threadblock-scoped pipelined matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::EllMmaPipelined< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, + layout::RowMajor, typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output (OperatorClass TensorOp) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Operation performed by GEMM + typename Operator + > +struct DefaultEllMma { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, 2, Operator>; + + // Define iterators over tiles from the A operand + using IteratorA = + mctlass::transform::threadblock::EllPredicatedTileIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, typename MmaCore::IteratorThreadMapA, kAlignmentA>; + + // Define iterators over tiles from the B operand + using IteratorB = + mctlass::transform::threadblock::EllPredicatedTileIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, typename MmaCore::IteratorThreadMapB, kAlignmentB>; + + // Define the threadblock-scoped pipelined matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::EllMmaPipelined< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, + layout::RowMajor, typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// +/// Specialization for row-major output (OperatorClass TensorOp) +template < + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Operation performed by GEMM + typename Operator + > +struct DefaultEllMma { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, float, LayoutA, float, + LayoutB, float, layout::RowMajor, arch::OpClassTensorOp, 2, + arch::OpMultiplyAddFastF16>; + + // Define iterators over tiles from the A operand + using IteratorA = + mctlass::transform::threadblock::EllPredicatedTileIterator< + mctlass::MatrixShape, + float, LayoutA, 1, typename MmaCore::IteratorThreadMapA, kAlignmentA>; + + // Define iterators over tiles from the B operand + using IteratorB = + mctlass::transform::threadblock::EllPredicatedTileIterator< + mctlass::MatrixShape, + float, LayoutB, 0, typename MmaCore::IteratorThreadMapB, kAlignmentB>; + + // Define the threadblock-scoped pipelined matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::EllMmaPipelined< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, float, + layout::RowMajor, typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for column-major-interleaved output +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Operation performed by GEMM + typename Operator, + /// Number of Interleaved K + int InterleavedK> +struct DefaultEllMma, OperatorClass, + ArchTag, ThreadblockShape, WarpShape, InstructionShape, 2, + Operator, true> { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, + layout::ColumnMajorInterleaved, OperatorClass, 2, Operator, + true>; + + static_assert(kAlignmentA == 128 / sizeof_bits::value, + "Alignment must match thread data map's vector length"); + + static_assert(kAlignmentB ==128 / sizeof_bits::value, + "Alignment must match thread data map's vector length"); + + // Define iterators over tiles from the A operand + using IteratorA = mctlass::transform::threadblock::EllPredicatedTileIterator< + mctlass::MatrixShape, ElementA, + LayoutA, 1, typename MmaCore::IteratorThreadMapA>; + + // Define iterators over tiles from the B operand + using IteratorB = mctlass::transform::threadblock::EllPredicatedTileIterator< + mctlass::MatrixShape, ElementB, + LayoutB, 0, typename MmaCore::IteratorThreadMapB>; + + // Define the threadblock-scoped pipelined matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::EllMmaPipelined< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, + layout::ColumnMajorInterleaved, + typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator + > +struct DefaultEllMma { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, arch::OpClassSimt, + Stages, Operator>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::EllPredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::EllPredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::EllMmaMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output (OperatorClass TensorOp) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator + > +struct DefaultEllMma { + static mctlass::arch::CacheOperation::Kind const CacheOpA = + ((sizeof_bits::value * kAlignmentA) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * kAlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, Operator, false, CacheOpA, CacheOpB>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::EllPredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::EllPredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::EllMmaMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for column-major-interleaved output +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// Number of Interleaved K + int InterleavedK> +struct DefaultEllMma, OperatorClass, + ArchTag, ThreadblockShape, WarpShape, InstructionShape, + Stages, Operator, true> { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, + layout::ColumnMajorInterleaved, OperatorClass, Stages, + Operator, true>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::EllPredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::EllPredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::EllMmaMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for SIMT IDP4A Kernels +template < + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Operation performed by GEMM + typename Operator, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape> +struct DefaultEllMma, 2, + Operator, false> { + using InstructionShape = GemmShape<1, 1, 4>; + using ElementA = int8_t; + using ElementB = int8_t; + using OperatorClass = arch::OpClassSimt; + + static const bool transposeA = mctlass::platform::is_same< LayoutA, layout::ColumnMajor >::value; + static const bool transposeB = mctlass::platform::is_same< LayoutB, layout::RowMajor >::value; + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, + OperatorClass, 2, Operator>; + + // Define iterators over tiles from the A operand + using IteratorA = + mctlass::transform::threadblock::PredicatedTileIterator2dThreadTile< + mctlass::MatrixShape, + ElementA, LayoutA, 1, typename MmaCore::IteratorThreadMapA, transposeA>; + + // Define iterators over tiles from the B operand + using IteratorB = + mctlass::transform::threadblock::PredicatedTileIterator2dThreadTile< + mctlass::MatrixShape, + ElementB, LayoutB, 0, typename MmaCore::IteratorThreadMapB, transposeB>; + + // Define the threadblock-scoped pipelined matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::EllMmaPipelined< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, + layout::RowMajor, typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +/// Specialization for Wmma TensorOp operator with 2 staged pipeline +template < + ///< Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Operation performed by GEMM + typename Operator> +struct DefaultEllMma { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, LayoutC, + arch::OpClassWmmaTensorOp, 2, Operator>; + + // Define iterators over tiles from the A operand + using IteratorA = + mctlass::transform::threadblock::EllPredicatedTileIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, typename MmaCore::IteratorThreadMapA, kAlignmentA>; + + // Define iterators over tiles from the B operand + using IteratorB = + mctlass::transform::threadblock::EllPredicatedTileIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, typename MmaCore::IteratorThreadMapB, kAlignmentB>; + + // Define the threadblock-scoped pipelined matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::EllMmaPipelined< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, + LayoutC, typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for Wmma TensorOp operator with 1 staged pipeline +template < + ///< Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Operation performed by GEMM + typename Operator> +struct DefaultEllMma { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, LayoutC, + arch::OpClassWmmaTensorOp, 1, Operator>; + + // Define iterators over tiles from the A operand + using IteratorA = + mctlass::transform::threadblock::EllPredicatedTileIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, typename MmaCore::IteratorThreadMapA, kAlignmentA>; + + // Define iterators over tiles from the B operand + using IteratorB = + mctlass::transform::threadblock::EllPredicatedTileIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, typename MmaCore::IteratorThreadMapB, kAlignmentB>; + + // Define the threadblock-scoped singlestage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaSingleStage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, + LayoutC, typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// +#endif //MCTLASS_ARCH_WMMA_ENABLED + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_gemv_core.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_gemv_core.h new file mode 100644 index 0000000..40069f9 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_gemv_core.h @@ -0,0 +1,151 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic properties needed by CTA-level batched GEMV assuming expectations about data + layout of the global memory fragments, data types, and internal tile sizes. + + Partial specializations for threadblock::Mma operations targeting SIMT instructions. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/layout/matrix.h" + +#include "mctlass/platform/platform.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/thread/mma.h" + +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" +#include "mctlass/transform/pitch_linear_thread_map.h" + +#include "mctlass/gemm/threadblock/gemv.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// +namespace mctlass { +namespace gemm { +namespace threadblock { + +/// Template defininng default vector-matrix multiply operators inferred from threadblock tile size, +/// global memory data layout. +template < + typename Shape_, /// Shape of the threadblock vector-matrix multiply operator + typename ThreadShape_, /// Shape of per-thread vector-matrix multiply operator + typename ElementA_, /// Element data type of A operand + typename LayoutA_, /// Layout of operand A + typename ElementB_, /// Element data type of B operand + typename LayoutB_, /// Layout of operand B + typename ElementC_, /// Data type of accumulator + typename LayoutC_ /// Layout of accumulator +> +struct DefaultGemvCore { + + using Shape = Shape_; + using ThreadShape = ThreadShape_; + + using LayoutA = LayoutA_; + using LayoutB = LayoutB_; + using LayoutC = LayoutC_; + + using ElementA = ElementA_; + using ElementB = ElementB_; + using ElementC = ElementC_; + + static int const kThreadsPerN = Shape::kN / ThreadShape::kN; + + using IteratorPolicyA = typename platform::conditional< + platform::is_same::value, + mctlass::transform::PitchLinearTilePolicyStripminedThreadContiguous< + layout::PitchLinearShape, 1, ThreadShape::kK>, + mctlass::transform::PitchLinearTilePolicyStripminedThreadStrided< + layout::PitchLinearShape, 1, ThreadShape::kM>>::type; + + using IteratorA = mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, ElementA, LayoutA, 1, IteratorPolicyA>; + + using IteratorPolicyB = typename platform::conditional< + platform::is_same::value, + mctlass::transform::PitchLinearTilePolicyStripminedThreadContiguous< + layout::PitchLinearShape, kThreadsPerN, ThreadShape::kN>, + mctlass::transform::PitchLinearTilePolicyStripminedThreadStrided< + layout::PitchLinearShape, kThreadsPerN, ThreadShape::kK>>::type; + + using IteratorB = mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, ElementB, LayoutB, 0, IteratorPolicyB>; + + using IteratorPolicyC = typename platform::conditional< + platform::is_same::value, + mctlass::transform::PitchLinearTilePolicyStripminedThreadContiguous< + layout::PitchLinearShape, kThreadsPerN, ThreadShape::kN>, + mctlass::transform::PitchLinearTilePolicyStripminedThreadStrided< + layout::PitchLinearShape, kThreadsPerN, ThreadShape::kM>>::type; + + using IteratorC = mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, ElementC, LayoutC, 0, IteratorPolicyC>; + + using MmaSimtOp = typename mctlass::gemm::thread::Mma< + mctlass::gemm::GemmShape, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC>; + + using Operator = MmaSimtOp; + + // Assertions for correctness + static_assert((Shape::kM == 1), "M=1 is required for GEMV"); + + static_assert((ThreadShape::kM == 1), "M=1 is required for GEMV"); + + static_assert(Shape::kK % ThreadShape::kK == 0, "Shape::K must be a multiple of ThreadShape::K"); + + static_assert(((ThreadShape::kK == 1) || + (ThreadShape::kK == 2) || + (ThreadShape::kK == 4) || + (ThreadShape::kK == 8) || + (ThreadShape::kK == 16) || + (ThreadShape::kK == 32) + ), + "ThreadShape::K must be a 1, 2, 4, 8, 16 or 32"); +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma.h new file mode 100644 index 0000000..0b4a1cd --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma.h @@ -0,0 +1,823 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/permute.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator_2dthreadtile.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/gemm/threadblock/default_mma_core_wmma.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// Gather operand A by using an index array + bool GatherA = false, + /// Gather operand B by using an index array + bool GatherB = false, + /// Permute operand A + typename PermuteALayout = layout::NoPermute, + /// Permute operand B + typename PermuteBLayout = layout::NoPermute + > +struct DefaultMma; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output (OperatorClass Simt) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operand + typename LayoutC, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Operation performed by GEMM + typename Operator, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout + > +struct DefaultMma { + + static_assert(platform::is_same::value + || platform::is_same>::value, + "simt epilogue must be row major"); + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, LayoutC, + arch::OpClassSimt, 2, Operator>; + + // Define iterators over tiles from the A operand + using IteratorA = + mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, typename MmaCore::IteratorThreadMapA, kAlignmentA, + GatherA, PermuteALayout>; + + // Define iterators over tiles from the B operand + using IteratorB = + mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, typename MmaCore::IteratorThreadMapB, kAlignmentB, + GatherB, PermuteBLayout>; + + // Define the threadblock-scoped pipelined matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaPipelined< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, + LayoutC, typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output (OperatorClass TensorOp) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Operation performed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout + > +struct DefaultMma { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, + arch::OpClassTensorOp, 2, Operator>; + + // Define iterators over tiles from the A operand + using IteratorA = + mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, typename MmaCore::IteratorThreadMapA, kAlignmentA, + GatherA, PermuteALayout>; + + // Define iterators over tiles from the B operand + using IteratorB = + mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, typename MmaCore::IteratorThreadMapB, kAlignmentB, + GatherB, PermuteBLayout>; + + // Define the threadblock-scoped pipelined matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaPipelined< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, + layout::RowMajor, typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// +/// Specialization for row-major output (OperatorClass TensorOp) +template < + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Operation performed by GEMM + typename Operator, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout + > +struct DefaultMma { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, float, LayoutA, float, + LayoutB, float, layout::RowMajor, arch::OpClassTensorOp, 2, + arch::OpMultiplyAddFastF16>; + + // Define iterators over tiles from the A operand + using IteratorA = + mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, + float, LayoutA, 1, typename MmaCore::IteratorThreadMapA, kAlignmentA, + GatherA, PermuteALayout>; + + // Define iterators over tiles from the B operand + using IteratorB = + mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, + float, LayoutB, 0, typename MmaCore::IteratorThreadMapB, kAlignmentB, + GatherB, PermuteBLayout>; + + // Define the threadblock-scoped pipelined matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaPipelined< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, float, + layout::RowMajor, typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for column-major-interleaved output +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Operation performed by GEMM + typename Operator, + /// Number of Interleaved K + int InterleavedK> +struct DefaultMma, OperatorClass, + ArchTag, ThreadblockShape, WarpShape, InstructionShape, 2, + Operator, true, SharedMemoryClearOption::kNone, false, false, + layout::NoPermute, layout::NoPermute> { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, + layout::ColumnMajorInterleaved, OperatorClass, 2, Operator, + true>; + + static_assert(kAlignmentA == 128 / sizeof_bits::value, + "Alignment must match thread data map's vector length"); + + static_assert(kAlignmentB ==128 / sizeof_bits::value, + "Alignment must match thread data map's vector length"); + + // Define iterators over tiles from the A operand + using IteratorA = mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, ElementA, + LayoutA, 1, typename MmaCore::IteratorThreadMapA>; + + // Define iterators over tiles from the B operand + using IteratorB = mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, ElementB, + LayoutB, 0, typename MmaCore::IteratorThreadMapB>; + + // Define the threadblock-scoped pipelined matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaPipelined< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, + layout::ColumnMajorInterleaved, + typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operand + typename LayoutC, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout + > +struct DefaultMma { + + static_assert(platform::is_same::value + || platform::is_same>::value, + "simt epilogue must be row major"); + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, LayoutC, arch::OpClassSimt, + Stages, Operator>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, AccessTypeA, GatherA, PermuteALayout>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, AccessTypeB, GatherB, PermuteBLayout>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, LayoutC, + typename MmaCore::MmaPolicy, Stages>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output (OperatorClass TensorOp) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operand + typename LayoutC, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear, + /// Gather operand A by using an index array + bool GatherA, + /// Gather operand B by using an index array + bool GatherB, + /// Permute operand A + typename PermuteALayout, + /// Permute operand B + typename PermuteBLayout + > +struct DefaultMma { + + static_assert(platform::is_same::value + || platform::is_same>::value, + "simt epilogue must be row major"); + + static mctlass::arch::CacheOperation::Kind const CacheOpA = + ((sizeof_bits::value * kAlignmentA) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * kAlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, LayoutC, arch::OpClassTensorOp, + Stages, Operator, false, CacheOpA, CacheOpB>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, AccessTypeA, GatherA, PermuteALayout>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, AccessTypeB, GatherB, PermuteBLayout>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, LayoutC, + typename MmaCore::MmaPolicy, Stages, SharedMemoryClear>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for column-major-interleaved output +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Operation performed by GEMM + typename Operator, + /// Number of Interleaved K + int InterleavedK> +struct DefaultMma, OperatorClass, + ArchTag, ThreadblockShape, WarpShape, InstructionShape, + Stages, Operator, true, SharedMemoryClearOption::kNone, + false, false, layout::NoPermute, layout::NoPermute> { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, + layout::ColumnMajorInterleaved, OperatorClass, Stages, + Operator, true>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for SIMT IDP4A Kernels +template < + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Operation performed by GEMM + typename Operator, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape> +struct DefaultMma, 2, + Operator, false, SharedMemoryClearOption::kNone, + false, false, layout::NoPermute, layout::NoPermute> { + using InstructionShape = GemmShape<1, 1, 4>; + using ElementA = int8_t; + using ElementB = int8_t; + using OperatorClass = arch::OpClassSimt; + + static const bool transposeA = platform::is_same< LayoutA, layout::ColumnMajor >::value; + static const bool transposeB = platform::is_same< LayoutB, layout::RowMajor >::value; + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, + OperatorClass, 2, Operator>; + + // Define iterators over tiles from the A operand + using IteratorA = + mctlass::transform::threadblock::PredicatedTileIterator2dThreadTile< + mctlass::MatrixShape, + ElementA, LayoutA, 1, typename MmaCore::IteratorThreadMapA, transposeA>; + + // Define iterators over tiles from the B operand + using IteratorB = + mctlass::transform::threadblock::PredicatedTileIterator2dThreadTile< + mctlass::MatrixShape, + ElementB, LayoutB, 0, typename MmaCore::IteratorThreadMapB, transposeB>; + + // Define the threadblock-scoped pipelined matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaPipelined< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, + layout::RowMajor, typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +/// Specialization for Wmma TensorOp operator with 2 staged pipeline +template < + ///< Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Operation performed by GEMM + typename Operator> +struct DefaultMma { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, LayoutC, + arch::OpClassWmmaTensorOp, 2, Operator>; + + // Define iterators over tiles from the A operand + using IteratorA = + mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, typename MmaCore::IteratorThreadMapA, kAlignmentA>; + + // Define iterators over tiles from the B operand + using IteratorB = + mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, typename MmaCore::IteratorThreadMapB, kAlignmentB>; + + // Define the threadblock-scoped pipelined matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaPipelined< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, + LayoutC, typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for Wmma TensorOp operator with 1 staged pipeline +template < + ///< Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Operation performed by GEMM + typename Operator> +struct DefaultMma { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, LayoutC, + arch::OpClassWmmaTensorOp, 1, Operator>; + + // Define iterators over tiles from the A operand + using IteratorA = + mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, typename MmaCore::IteratorThreadMapA, kAlignmentA>; + + // Define iterators over tiles from the B operand + using IteratorB = + mctlass::transform::threadblock::PredicatedTileIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, typename MmaCore::IteratorThreadMapB, kAlignmentB>; + + // Define the threadblock-scoped singlestage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaSingleStage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + IteratorB, typename MmaCore::SmemIteratorB, ElementAccumulator, + LayoutC, typename MmaCore::MmaPolicy>; +}; + +//////////////////////////////////////////////////////////////////////////////// +#endif //MCTLASS_ARCH_WMMA_ENABLED + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core.h new file mode 100644 index 0000000..360b5c7 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core.h @@ -0,0 +1,116 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic properties needed by CTA-level GEMMs assuming expectations about data + layout of the global memory fragments, data types, and internal tile sizes. + + Partial specializations for threadblock::Mma operations targeting TensorOp instructions. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/warp/mma.h" +#include "mctlass/gemm/threadblock/mma_pipelined.h" +#include "mctlass/gemm/threadblock/mma_singlestage.h" +#include "mctlass/arch/cache_operation.h" +#include "mctlass/arch/mma.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template defininng default matrix multiply operators inferred from threadblock tile size, +/// global memory data layout, and target math instruction. +template < + /// Shape of threadblock-scoped matrix multiply operator + typename Shape, + /// Shape of warp-level matrix multiply operator + typename WarpShape, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape, + /// Element data type of A operand + typename ElementA, + /// Layout of operand A + typename LayoutA, + /// Element data type of B operand + typename ElementB, + /// Layout of operand B + typename LayoutB, + /// Data type of accumulator + typename ElementC, + /// Layout of accumulator + typename LayoutC, + /// Indicates type of math operator (arch::OpClassSimt or arch::OpClassTensorOp) + typename OperatorClass, + /// Number of stages + int Stages = 2, + /// Operation performed by MMA + typename Operator = typename platform::conditional< + (platform::is_same::value) && + (platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value), + mctlass::arch::OpMultiplyAddSaturate, + mctlass::arch::OpMultiplyAdd>::type, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA = + mctlass::arch::CacheOperation::Global, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB = + mctlass::arch::CacheOperation::Global, + /// per-element transformation for elements of A + ComplexTransform TransformA = ComplexTransform::kNone, + /// per-element transformation for elements of B + ComplexTransform TransformB = ComplexTransform::kNone, + bool IsComplex = false // (is_complex::value || is_complex::value) +> +struct DefaultMmaCore; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_simt.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_simt.h new file mode 100644 index 0000000..054fb0d --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_simt.h @@ -0,0 +1,1723 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic properties needed by CTA-level GEMMs assuming expectations about data + layout of the global memory fragments, data types, and internal tile sizes. + + Partial specializations for threadblock::Mma operations targeting simt instructions. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/fast_math.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + + +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h" +#include "mctlass/transform/threadblock/regular_tile_iterator_pitch_linear_2dthreadtile.h" + +#include "mctlass/gemm/warp/mma_simt_policy.h" +#include "mctlass/gemm/warp/mma_simt.h" +#include "mctlass/gemm/threadblock/default_mma_core.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +namespace detail { + +// convert a WarpShape which is the whole tile of elements into warp num threads. +// The goal is for each thread's tile of elements to be as square as possible +// for performance (4x4 will be faster than 2x8). +template +constexpr int simt_get_warp_threads_m() { + return (WarpShape::kM > WarpShape::kN) ? 8 : 4; +} + +/// Computes padding in shared memory to perform efficient transpose without bank conflicts. +constexpr int simt_transpose_padding(int threads, int crosswise, int size_in_bits) { + return (size_in_bits >= 32 ? + threads / crosswise / (size_in_bits / 32) : + threads / crosswise * (32 / size_in_bits) + ); +} + +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: row-major +/// Operator: simt class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, ElementA_, + layout::ColumnMajor, ElementB_, layout::RowMajor, + ElementC_, LayoutC_, arch::OpClassSimt, 2, Operator_ + > { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 1>; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + static int const PartitionsK = Shape::kK / WarpShape::kK; + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + PartitionsK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + static int const kElementsPerAccess = 1; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + IteratorThreadMapA + >; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = detail::simt_get_warp_threads_m(); + static const int WarpNumThreadsN = kWarpSize / WarpNumThreadsM; + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy /// Policy describing warp-level MmaSimtOp (concept: MmaSimtOp policy) + >; /// Used for partial specialization + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: column-major +/// Operator: simt class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, ElementA_, + layout::RowMajor, ElementB_, layout::ColumnMajor, + ElementC_, LayoutC_, arch::OpClassSimt, 2, Operator_ + > { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 1>; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + static int const PartitionsK = Shape::kK / WarpShape::kK; + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + PartitionsK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + static int const kElementsPerAccess = 1; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapA = transform::TransposePitchLinearThreadMapSimt; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + SmemThreadMapA // was IteratorThreadMapA + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapB = transform::TransposePitchLinearThreadMapSimt; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + SmemThreadMapB // was IteratorThreadMapA + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = detail::simt_get_warp_threads_m(); + static const int WarpNumThreadsN = kWarpSize / WarpNumThreadsM; + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + + static int const kPaddingM = detail::simt_transpose_padding(kWarpSize, Shape::kK, sizeof_bits::value); + static int const kPaddingN = detail::simt_transpose_padding(kWarpSize, Shape::kK, sizeof_bits::value); + + static_assert(!(kPaddingM % LaneM) && !(kPaddingN % LaneN), + "Padding must be divisible by Lane"); + + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy /// Policy describing warp-level MmaSimtOp (concept: MmaSimtOp policy) + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape, // skew for A matrix to avoid SMEM bank conflicts + MatrixShape<0, kPaddingN>, // skew for B matrix to avoid SMEM bank conflicts + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: row-major +/// Operator: simt class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, ElementA_, + layout::RowMajor, ElementB_, layout::RowMajor, ElementC_, + LayoutC_, arch::OpClassSimt, 2, Operator_ + > { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 1>; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + static int const PartitionsK = Shape::kK / WarpShape::kK; + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + PartitionsK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + static int const kElementsPerAccess = 1; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapA = transform::TransposePitchLinearThreadMapSimt; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + SmemThreadMapA + >; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = detail::simt_get_warp_threads_m(); + static const int WarpNumThreadsN = kWarpSize / WarpNumThreadsM; + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + + static int const kPaddingM = detail::simt_transpose_padding(kWarpSize, Shape::kK, sizeof_bits::value); + + static_assert(!(kPaddingM % LaneM), + "Padding must be divisible by Lane"); + + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy /// Policy describing warp-level MmaSimtOp (concept: MmaSimtOp policy) + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape, // skew for A matrix to avoid SMEM bank conflicts + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: column-major +/// Operator: simt class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, ElementA_, + layout::ColumnMajor, ElementB_, layout::ColumnMajor, + ElementC_, LayoutC_, arch::OpClassSimt, 2, Operator_ + > { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 1>; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + static int const PartitionsK = Shape::kK / WarpShape::kK; + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + PartitionsK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + static int const kElementsPerAccess = 1; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + IteratorThreadMapA + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapB = transform::TransposePitchLinearThreadMapSimt; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + SmemThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = detail::simt_get_warp_threads_m(); + static const int WarpNumThreadsN = kWarpSize / WarpNumThreadsM; + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + + static int const kPaddingN = detail::simt_transpose_padding(kWarpSize, Shape::kK, sizeof_bits::value); + + static_assert(!(kPaddingN % LaneN), + "Padding must be divisible by Lane"); + + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy /// Policy describing warp-level MmaSimtOp (concept: MmaSimtOp policy) + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape<0, 0>, + MatrixShape<0, kPaddingN>, // skew for B matrix to avoid SMEM bank conflicts + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: row-major +/// Operator: simt class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, ElementA_, + layout::AffineRank2ColumnMajor, ElementB_, layout::AffineRank2RowMajor, + ElementC_, LayoutC_, arch::OpClassSimt, 2, Operator_ + > { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 1>; + using ElementA = ElementA_; + using LayoutA = layout::AffineRank2ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::AffineRank2RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + + /// Default Operator + using Operator = Operator_; + + using Base = DefaultMmaCore; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename Base::SmemLayoutA; + using SmemLayoutB = typename Base::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename Base::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename Base::SmemIteratorA; + + /// Policy of iterator B + using IteratorThreadMapB = typename Base::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename Base::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + /// Policy used to define MmaPipelined + using MmaPolicy = typename Base::MmaPolicy; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: column-major +/// Operator: simt class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, ElementA_, + layout::AffineRank2RowMajor, ElementB_, layout::AffineRank2ColumnMajor, + ElementC_, LayoutC_, arch::OpClassSimt, 2, Operator_ + > { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 1>; + using ElementA = ElementA_; + using LayoutA = layout::AffineRank2RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::AffineRank2ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + + /// Default Operator + using Operator = Operator_; + + using Base = DefaultMmaCore; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename Base::SmemLayoutA; + using SmemLayoutB = typename Base::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename Base::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename Base::SmemIteratorA; + + /// Policy of iterator B + using IteratorThreadMapB = typename Base::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename Base::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + /// Policy used to define MmaPipelined + using MmaPolicy = typename Base::MmaPolicy; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: row-major +/// Operator: simt class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, ElementA_, + layout::AffineRank2RowMajor, ElementB_, layout::AffineRank2RowMajor, ElementC_, + LayoutC_, arch::OpClassSimt, 2, Operator_ + > { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 1>; + using ElementA = ElementA_; + using LayoutA = layout::AffineRank2RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::AffineRank2RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + + /// Default Operator + using Operator = Operator_; + + using Base = DefaultMmaCore; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename Base::SmemLayoutA; + using SmemLayoutB = typename Base::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename Base::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename Base::SmemIteratorA; + + /// Policy of iterator B + using IteratorThreadMapB = typename Base::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename Base::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + /// Policy used to define MmaPipelined + using MmaPolicy = typename Base::MmaPolicy; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: column-major +/// Operator: simt class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, ElementA_, + layout::AffineRank2ColumnMajor, ElementB_, layout::AffineRank2ColumnMajor, + ElementC_, LayoutC_, arch::OpClassSimt, 2, Operator_ + > { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 1>; + using ElementA = ElementA_; + using LayoutA = layout::AffineRank2ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::AffineRank2ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + + /// Default Operator + using Operator = Operator_; + + using Base = DefaultMmaCore; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename Base::SmemLayoutA; + using SmemLayoutB = typename Base::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename Base::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename Base::SmemIteratorA; + + /// Policy of iterator B + using IteratorThreadMapB = typename Base::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename Base::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + /// Policy used to define MmaPipelined + using MmaPolicy = typename Base::MmaPolicy; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: row-major +/// Operator: simt class, for dp4a +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, int8_t, + layout::ColumnMajor, int8_t, layout::RowMajor, ElementC_, + LayoutC_, arch::OpClassSimt, 2, Operator_ + > { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 4>; + using ElementA = int8_t; + using LayoutA = layout::ColumnMajor; + using ElementB = int8_t; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + static int const PartitionsK = Shape::kK / WarpShape::kK; + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + PartitionsK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorInterleaved<4>; + using SmemLayoutB = layout::RowMajorInterleaved<4>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinear2DThreadTileStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<4, 4> + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator2dThreadTile< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + IteratorThreadMapA + >; + + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinear2DThreadTileStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<4, 4> + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator2dThreadTile< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = detail::simt_get_warp_threads_m(); + static const int WarpNumThreadsN = kWarpSize / WarpNumThreadsM; + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(4, ThreadTileM); + static const int LaneN = mctlass::const_min(4, ThreadTileN); + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 4>; + + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::ColumnMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy, /// Policy describing warp-level MmaSimtOp (concept: MmaSimtOp policy) + PartitionsK /// Number of partitions along K dimension + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Partial specialization: +// +/// +/// A: Row-major +/// B: Column-major +/// Operator: simt class, for dp4a +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, int8_t, + layout::RowMajor, int8_t, layout::ColumnMajor, ElementC_, + LayoutC_, arch::OpClassSimt, 2, Operator_ + > { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 4>; + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + static int const PartitionsK = Shape::kK / WarpShape::kK; + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + PartitionsK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorInterleaved<4>; + using SmemLayoutB = layout::RowMajorInterleaved<4>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinear2DThreadTileStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<4, 4> + >; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapA = transform::TransposePitchLinearThreadMap2DThreadTile; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator2dThreadTile< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + SmemThreadMapA + >; + + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinear2DThreadTileStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<4, 4> + >; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapB = transform::TransposePitchLinearThreadMap2DThreadTile; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator2dThreadTile< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + SmemThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = detail::simt_get_warp_threads_m(); + static const int WarpNumThreadsN = kWarpSize / WarpNumThreadsM; + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(4, ThreadTileM); + static const int LaneN = mctlass::const_min(4, ThreadTileN); + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 4>; + + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::ColumnMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy, /// Policy describing warp-level MmaSimtOp (concept: MmaSimtOp policy) + PartitionsK /// Number of partitions along K dimension + >; + + static int const kPaddingM = detail::simt_transpose_padding(kWarpSize, Shape::kK, sizeof_bits::value); + static int const kPaddingN = detail::simt_transpose_padding(kWarpSize, Shape::kK, sizeof_bits::value); + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape, + MatrixShape<0, kPaddingN>, + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Partial specialization: +// +/// +/// A: Row-major +/// B: Row-major +/// Operator: simt class, for dp4a +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, int8_t, + layout::RowMajor, int8_t, layout::RowMajor, ElementC_, + LayoutC_, arch::OpClassSimt, 2, Operator_ + > { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 4>; + using ElementA = int8_t; + using LayoutA = layout::RowMajor; + using ElementB = int8_t; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + static int const PartitionsK = Shape::kK / WarpShape::kK; + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + PartitionsK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorInterleaved<4>; + using SmemLayoutB = layout::RowMajorInterleaved<4>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinear2DThreadTileStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<4, 4> + >; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapA = transform::TransposePitchLinearThreadMap2DThreadTile; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator2dThreadTile< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + SmemThreadMapA + >; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinear2DThreadTileStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<4, 4> + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator2dThreadTile< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = detail::simt_get_warp_threads_m(); + static const int WarpNumThreadsN = kWarpSize / WarpNumThreadsM; + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(4, ThreadTileM); + static const int LaneN = mctlass::const_min(4, ThreadTileN); + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 4>; + + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::ColumnMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy, /// Policy describing warp-level MmaSimtOp (concept: MmaSimtOp policy) + PartitionsK /// Number of partitions along K dimension + >; + + static int const kPaddingM = detail::simt_transpose_padding(kWarpSize, Shape::kK, sizeof_bits::value); + static int const kPaddingN = detail::simt_transpose_padding(kWarpSize, Shape::kK, sizeof_bits::value); + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Partial specialization: +// +/// +/// A: Column-major +/// B: Column-major +/// Operator: simt class, for dp4a +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, int8_t, + layout::ColumnMajor, int8_t, layout::ColumnMajor, ElementC_, + LayoutC_, arch::OpClassSimt, 2, Operator_ + > { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 4>; + using ElementA = int8_t; + using LayoutA = layout::ColumnMajor; + using ElementB = int8_t; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + static int const PartitionsK = Shape::kK / WarpShape::kK; + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + PartitionsK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorInterleaved<4>; + using SmemLayoutB = layout::RowMajorInterleaved<4>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinear2DThreadTileStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<4, 4> + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator2dThreadTile< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + IteratorThreadMapA + >; + + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinear2DThreadTileStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<4, 4> + >; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapB = transform::TransposePitchLinearThreadMap2DThreadTile; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator2dThreadTile< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + SmemThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = detail::simt_get_warp_threads_m(); + static const int WarpNumThreadsN = kWarpSize / WarpNumThreadsM; + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(4, ThreadTileM); + static const int LaneN = mctlass::const_min(4, ThreadTileN); + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 4>; + + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::ColumnMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy, /// Policy describing warp-level MmaSimtOp (concept: MmaSimtOp policy) + PartitionsK /// Number of partitions along K dimension + >; + + static int const kPaddingM = detail::simt_transpose_padding(kWarpSize, Shape::kK, sizeof_bits::value); + static int const kPaddingN = detail::simt_transpose_padding(kWarpSize, Shape::kK, sizeof_bits::value); + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape<0, 0>, + MatrixShape<0, kPaddingN>, + WarpCount::kK + >; +}; + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sm70.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sm70.h new file mode 100644 index 0000000..d149fd5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sm70.h @@ -0,0 +1,682 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic properties needed by CTA-level GEMMs assuming expectations about data + layout of the global memory fragments, data types, and internal tile sizes. + + Partial specializations for threadblock::Mma operations targeting TensorOp instructions. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + + +#include "mctlass/layout/tensor_op_multiplicand_sm70.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_iterator_tensor_op_sm70.h" + +#include "mctlass/gemm/warp/mma_tensor_op_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, ElementA_, + layout::ColumnMajor, ElementB_, layout::RowMajor, + ElementC_, LayoutC_, arch::OpClassTensorOp, 2, Operator_ + > { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<8, 8, 4>; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassTensorOp; + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + Shape::kK / WarpShape::kK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + // + // Shared memory layouts + // + + using SmemLayoutA = + layout::ColumnMajorVoltaTensorOpMultiplicandCongruous< + sizeof_bits::value>; + + // Shared memory layout + using SmemLayoutB = + layout::RowMajorVoltaTensorOpMultiplicandBCongruous< + sizeof_bits::value>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + IteratorThreadMapA + >; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + mctlass::gemm::GemmShape<16, 16, 4>, + 32, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + mctlass::layout::RowMajor, + mctlass::arch::OpMultiplyAdd + >, + mctlass::MatrixShape<1, 1> + >; + + using MmaTensorOp = mctlass::gemm::warp::MmaVoltaTensorOp< + WarpShape, + ElementA, + SmemLayoutA, + ElementB, + SmemLayoutB, + ElementC, + LayoutC, + Policy + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +/// Partial specialization: +/// +/// A: row-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, ElementA_, + layout::RowMajor, ElementB_, layout::ColumnMajor, + ElementC_, LayoutC_, arch::OpClassTensorOp, 2, Operator_ + > { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<8, 8, 4>; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassTensorOp; + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + Shape::kK / WarpShape::kK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorVoltaTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // Shared memory layout + using SmemLayoutB = layout::ColumnMajorVoltaTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<4, 8>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 0, + IteratorThreadMapA + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<4, 8>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 1, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + mctlass::gemm::GemmShape<16, 16, 4>, + 32, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + mctlass::layout::RowMajor, + mctlass::arch::OpMultiplyAdd + >, + mctlass::MatrixShape<1, 1> + >; + + using MmaTensorOp = mctlass::gemm::warp::MmaVoltaTensorOp< + WarpShape, + ElementA, + SmemLayoutA, + ElementB, + SmemLayoutB, + ElementC, + LayoutC, + Policy + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, ElementA_, + layout::RowMajor, ElementB_, layout::RowMajor, ElementC_, + LayoutC_, arch::OpClassTensorOp, 2, Operator_ + > { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<8, 8, 4>; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassTensorOp; + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + Shape::kK / WarpShape::kK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorVoltaTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // Shared memory layout + using SmemLayoutB = layout::RowMajorVoltaTensorOpMultiplicandBCongruous< + sizeof_bits::value>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<4, 8>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 0, + IteratorThreadMapA + >; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + mctlass::gemm::GemmShape<16, 16, 4>, + 32, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + mctlass::layout::RowMajor, + mctlass::arch::OpMultiplyAdd + >, + mctlass::MatrixShape<1, 1> + >; + + using MmaTensorOp = mctlass::gemm::warp::MmaVoltaTensorOp< + WarpShape, + ElementA, + SmemLayoutA, + ElementB, + SmemLayoutB, + ElementC, + LayoutC, + Policy + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore, ElementA_, + layout::ColumnMajor, ElementB_, layout::ColumnMajor, + ElementC_, LayoutC_, arch::OpClassTensorOp, 2, Operator_ + > { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<8, 8, 4>; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassTensorOp; + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + Shape::kK / WarpShape::kK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorVoltaTensorOpMultiplicandCongruous< + sizeof_bits::value>; + + // Shared memory layout + using SmemLayoutB = layout::ColumnMajorVoltaTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + IteratorThreadMapA + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<4, 8>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 1, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + mctlass::gemm::GemmShape<16, 16, 4>, + 32, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + mctlass::layout::RowMajor, + mctlass::arch::OpMultiplyAdd + >, + mctlass::MatrixShape<1, 1> + >; + + using MmaTensorOp = mctlass::gemm::warp::MmaVoltaTensorOp< + WarpShape, + ElementA, + SmemLayoutA, + ElementB, + SmemLayoutB, + ElementC, + LayoutC, + Policy + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sm75.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sm75.h new file mode 100644 index 0000000..6c9576b --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sm75.h @@ -0,0 +1,1284 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic properties needed by CTA-level GEMMs assuming expectations about data + layout of the global memory fragments, data types, and internal tile sizes. + + Partial specializations for threadblock::Mma operations targeting TensorOp instructions. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/platform/platform.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_iterator_tensor_op.h" + +#include "mctlass/gemm/warp/default_mma_tensor_op.h" +#include "mctlass/gemm/threadblock/default_mma_core.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassTensorOp; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + Shape::kK / WarpShape::kK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // + // Shared memory layouts + // + + using SmemLayoutA = + layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(ElementA))>; + + // Shared memory layout + using SmemLayoutB = layout::RowMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(ElementB))>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + IteratorThreadMapA + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by MMA + typename Operator_> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassTensorOp; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + Shape::kK / WarpShape::kK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize_0 = warp::WarpSize::value; + static int const kWarpSize_1 = kWarpSize_0 / 2; + static bool const m16n8k4_shape = (InstructionShape::kM == 16 && InstructionShape::kN == 8 && InstructionShape::kK == 4) ? true : false; + static bool const tfloat32_type = std::is_same::value ? true : false; + static bool const flag = (WarpShape::kM < 32 || WarpShape::kN < 32) ? true : false; + static int const kWarpSize = (m16n8k4_shape && tfloat32_type && flag) ? kWarpSize_1 : kWarpSize_0; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousA = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedA = + kWarpSize / kWarpThreadArrangementContiguousA; + + static int const kWarpThreadArrangementContiguousB = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedB = + kWarpSize / kWarpThreadArrangementContiguousB; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // Shared memory layout + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 0, + IteratorThreadMapA + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 1, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by MMA + typename Operator_> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassTensorOp; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + Shape::kK / WarpShape::kK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousA = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedA = + kWarpSize / kWarpThreadArrangementContiguousA; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // Shared memory layout + using SmemLayoutB = layout::RowMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(ElementB))>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 0, + IteratorThreadMapA + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by MMA + typename Operator_> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassTensorOp; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousB = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedB = + kWarpSize / kWarpThreadArrangementContiguousB; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(ElementA))>; + + // Shared memory layout + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// +/// Below is for arch::OpMultiplyAddFastF16 + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = float; + using LayoutA = layout::ColumnMajor; + using ElementB = float; + using LayoutB = layout::RowMajor; + using ElementC = float; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassTensorOp; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + Shape::kK / WarpShape::kK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 256; + + /// Default Operator + using Operator = arch::OpMultiplyAdd; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(half_t))>; + + // Shared memory layout + using SmemLayoutB = + layout::RowMajorTensorOpMultiplicandCongruous::value, + int(128 / sizeof(half_t))>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + half_t, + SmemLayoutA, + 1, + IteratorThreadMapA + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + half_t, + SmemLayoutB, + 0, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, half_t, SmemLayoutA, half_t, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = float; + using LayoutA = layout::RowMajor; + using ElementB = float; + using LayoutB = layout::ColumnMajor; + using ElementC = float; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassTensorOp; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + Shape::kK / WarpShape::kK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 256; + + /// Default Operator + using Operator = arch::OpMultiplyAdd; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousA = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedA = + kWarpSize / kWarpThreadArrangementContiguousA; + + static int const kWarpThreadArrangementContiguousB = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedB = + kWarpSize / kWarpThreadArrangementContiguousB; + + // + // Shared memory layouts + // + + using SmemLayoutA = + layout::RowMajorTensorOpMultiplicandCrosswise::value, + Shape::kK>; + + // Shared memory layout + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + half_t, + SmemLayoutA, + 0, + IteratorThreadMapA + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + half_t, + SmemLayoutB, + 1, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, half_t, SmemLayoutA, half_t, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = float; + using LayoutA = layout::RowMajor; + using ElementB = float; + using LayoutB = layout::RowMajor; + using ElementC = float; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassTensorOp; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + Shape::kK / WarpShape::kK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 256; + + /// Default Operator + using Operator = arch::OpMultiplyAdd; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousA = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedA = + kWarpSize / kWarpThreadArrangementContiguousA; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // Shared memory layout + using SmemLayoutB = layout::RowMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(half_t))>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + half_t, + SmemLayoutA, + 0, + IteratorThreadMapA + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + half_t, + SmemLayoutB, + 0, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, half_t, SmemLayoutA, half_t, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = float; + using LayoutA = layout::ColumnMajor; + using ElementB = float; + using LayoutB = layout::ColumnMajor; + using ElementC = float; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassTensorOp; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 256; + + /// Default Operator + using Operator = arch::OpMultiplyAdd; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousB = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedB = + kWarpSize / kWarpThreadArrangementContiguousB; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(half_t))>; + + // Shared memory layout + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, half_t, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, half_t, SmemLayoutB, 1, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, half_t, SmemLayoutA, half_t, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, MatrixShape<0, 0>, + WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major-interleave +/// B: row-major-interleave +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +/// +/// Column/RowMajorInterleved(m, n) is mapped to Column/RowMajor(m +/// x InterleavedK, n / InterleavedK) so that Column/RowMajor global iterators +/// can be reused. The shared store iterator is the same as the crosswise shared +/// store iterator. So, the only thing we need to do is to swap the coordinates +/// (contiguous <=> strided) used by the global iterator and the shared store +/// iterator. +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by MMA + typename Operator_, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor, + /// Number of interleaved k + int InterleavedK> +struct DefaultMmaCore, ElementB_, + layout::RowMajorInterleaved, ElementC_, + LayoutC_, arch::OpClassTensorOp, 2, Operator_, + AccumulatorsInRowMajor> { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajorInterleaved; + using ElementB = ElementB_; + using LayoutB = layout::RowMajorInterleaved; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassTensorOp; + static int const kInterleavedK = InterleavedK; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kElementsPerAccess = + kAccessSizeInBits / sizeof_bits::value; + + static int const kWarpThreadArrangementContiguous = + kInterleavedK / kElementsPerAccess; + + static int const kWarpThreadArrangementStrided = + kWarpSize / kWarpThreadArrangementContiguous; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, kInterleavedK>; + + // Shared memory layout + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, kInterleavedK>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, layout::PitchLinearShape<32, 1>, kElementsPerAccess>; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapA = transform::TransposePitchLinearThreadMap< + IteratorThreadMapA, + layout::PitchLinearShape>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + SmemThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, layout::PitchLinearShape<32, 1>, kElementsPerAccess>; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapB = transform::TransposePitchLinearThreadMap< + IteratorThreadMapB, + layout::PitchLinearShape>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + SmemThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK, AccumulatorsInRowMajor>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sm80.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sm80.h new file mode 100644 index 0000000..54a10da --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sm80.h @@ -0,0 +1,2916 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Defines basic properties needed by CTA-level GEMMs assuming + expectations about data layout of the global memory fragments, data types, + and internal tile sizes. + + Partial specializations for threadblock::Mma operations targeting TensorOp + instructions. + + SM80 Multi stage kernel expects stage number to be larger or equal to 3 + to use asyncronous copy. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/mctlass.h" + +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" +#include "mctlass/layout/tensor_op_multiplicand_sm80.h" + +#include "mctlass/gemm/warp/mma_simt_policy.h" +#include "mctlass/gemm/warp/mma_simt.h" +#include "mctlass/gemm/warp/default_mma_tensor_op.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" + +#include "mctlass/gemm/threadblock/default_mma_core.h" +#include "mctlass/gemm/threadblock/default_multistage_mma_complex_core.h" +#include "mctlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h" + +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op_sm80.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear.h" +#include "mctlass/gemm/threadblock/mma_multistage.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for double-precision +/// +/// A: column-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = double; + using LayoutA = layout::ColumnMajor; + using ElementB = double; + using LayoutB = layout::ColumnMajor; + using ElementC = double; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 64; + + /// Default Operator + using Operator = Operator_; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorTensorOpMultiplicandCongruous64b; + + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicand64bCrosswise; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpStripedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +/// Partial specialization for double-precision +/// +/// A: column-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = double; + using LayoutA = layout::ColumnMajor; + using ElementB = double; + using LayoutB = layout::RowMajor; + using ElementC = double; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 64; + + /// Default Operator + using Operator = Operator_; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorTensorOpMultiplicandCongruous64b; + + // Shared memory layout + using SmemLayoutB = layout::RowMajorTensorOpMultiplicandCongruous64b; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpStripedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpStripedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for double-precision +/// +/// A: row-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = double; + using LayoutA = layout::RowMajor; + using ElementB = double; + using LayoutB = layout::ColumnMajor; + using ElementC = double; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 64; + + /// Default Operator + using Operator = Operator_; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicand64bCrosswise; + + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicand64bCrosswise; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// +/// +/// Partial specialization for double-precision +/// +/// A: row-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = double; + using LayoutA = layout::RowMajor; + using ElementB = double; + using LayoutB = layout::RowMajor; + using ElementC = double; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 64; + + /// Default Operator + using Operator = Operator_; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicand64bCrosswise; + + using SmemLayoutB = layout::RowMajorTensorOpMultiplicandCongruous64b; + + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpStripedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for double-precision +/// +/// A: column-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = double; + using LayoutA = layout::AffineRank2ColumnMajor; + using ElementB = double; + using LayoutB = layout::AffineRank2ColumnMajor; + using ElementC = double; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Default Operator + using Operator = Operator_; + + using Base = DefaultMmaCore; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename Base::SmemLayoutA; + using SmemLayoutB = typename Base::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename Base::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename Base::SmemIteratorA; + + /// Policy of iterator B + using IteratorThreadMapB = typename Base::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename Base::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + /// Policy used to define MmaPipelined + using MmaPolicy = typename Base::MmaPolicy; +}; + +/// Partial specialization for double-precision +/// +/// A: column-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = double; + using LayoutA = layout::AffineRank2ColumnMajor; + using ElementB = double; + using LayoutB = layout::AffineRank2RowMajor; + using ElementC = double; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Default Operator + using Operator = Operator_; + + using Base = DefaultMmaCore; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename Base::SmemLayoutA; + using SmemLayoutB = typename Base::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename Base::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename Base::SmemIteratorA; + + /// Policy of iterator B + using IteratorThreadMapB = typename Base::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename Base::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + /// Policy used to define MmaPipelined + using MmaPolicy = typename Base::MmaPolicy; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for double-precision +/// +/// A: row-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = double; + using LayoutA = layout::AffineRank2RowMajor; + using ElementB = double; + using LayoutB = layout::AffineRank2ColumnMajor; + using ElementC = double; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Default Operator + using Operator = Operator_; + + using Base = DefaultMmaCore; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename Base::SmemLayoutA; + using SmemLayoutB = typename Base::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename Base::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename Base::SmemIteratorA; + + /// Policy of iterator B + using IteratorThreadMapB = typename Base::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename Base::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + /// Policy used to define MmaPipelined + using MmaPolicy = typename Base::MmaPolicy; +}; + +//////////////////////////////////////////////////////////////////////////////// +/// +/// Partial specialization for double-precision +/// +/// A: row-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = double; + using LayoutA = layout::AffineRank2RowMajor; + using ElementB = double; + using LayoutB = layout::AffineRank2RowMajor; + using ElementC = double; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Default Operator + using Operator = Operator_; + + using Base = DefaultMmaCore; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename Base::SmemLayoutA; + using SmemLayoutB = typename Base::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename Base::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename Base::SmemIteratorA; + + /// Policy of iterator B + using IteratorThreadMapB = typename Base::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename Base::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + /// Policy used to define MmaPipelined + using MmaPolicy = typename Base::MmaPolicy; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for float-precision +/// +/// ElementA: complex +/// ElementB: complex +/// ElementC: complex +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Layout for A operand + typename LayoutA_, + /// Layout for B operand + typename LayoutB_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// per-element transformation for elements of A + ComplexTransform TransformA_, + /// per-element transformation for elements of B + ComplexTransform TransformB_ + > +struct DefaultMmaCore< + Shape_, WarpShape_, GemmShape<16, 8, 8>, + complex, LayoutA_, + complex, LayoutB_, + complex, LayoutC_, + arch::OpClassTensorOp, + Stages, + Operator_, + false, + CacheOpA, + CacheOpB, + TransformA_, TransformB_, true> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<16, 8, 8>; + using ElementA = complex; + using LayoutA = LayoutA_; + using ElementB = complex; + using LayoutB = LayoutB_; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + static const ComplexTransform TransformA = TransformA_; + static const ComplexTransform TransformB = TransformB_; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + static_assert( + platform::is_same::value || + platform::is_same::value || + platform::is_same::value, + "The operator tag must indicate complex multiplication."); + + // + // Underlying template + // + + using MmaComplexCore = DefaultMultistageMmaComplexCore< + Shape, WarpShape, InstructionShape, + ElementA, LayoutA, + ElementB, LayoutB, + ElementC, LayoutC, + arch::OpClassTensorOp, + kStages, + TransformA, + TransformB, + Operator, + kCacheOpA, + kCacheOpB + >; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename MmaComplexCore::SmemLayoutA; + + // Shared memory layout + using SmemLayoutB = typename MmaComplexCore::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename MmaComplexCore::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename MmaComplexCore::SmemIteratorA; + + /// ThreadMap of iterator B + using IteratorThreadMapB = typename MmaComplexCore::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename MmaComplexCore::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename MmaComplexCore::MmaTensorOp; + + /// Policy used to define MmaPipelined + using MmaPolicy = typename MmaComplexCore::MmaPolicy; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for double-precision +/// +/// ElementA: complex +/// ElementB: complex +/// ElementC: complex +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout for A operand + typename LayoutA_, + /// Layout for B operand + typename LayoutB_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// per-element transformation for elements of A + ComplexTransform TransformA_, + /// per-element transformation for elements of B + ComplexTransform TransformB_ + > +struct DefaultMmaCore< + Shape_, WarpShape_, InstructionShape_, + complex, LayoutA_, + complex, LayoutB_, + complex, LayoutC_, + arch::OpClassTensorOp, + Stages, + Operator_, + false, + CacheOpA, + CacheOpB, + TransformA_, TransformB_, true> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = complex; + using LayoutA = LayoutA_; + using ElementB = complex; + using LayoutB = LayoutB_; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + static const ComplexTransform TransformA = TransformA_; + static const ComplexTransform TransformB = TransformB_; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 64; + + /// Default Operator + using Operator = Operator_; + + static_assert( + platform::is_same::value || + platform::is_same::value, + "The operator tag must indicate complex multiplication."); + + // + // Underlying template + // + + using MmaComplexCore = DefaultMultistageMmaComplexCore< + Shape, WarpShape, InstructionShape, + ElementA, LayoutA, + ElementB, LayoutB, + ElementC, LayoutC, + arch::OpClassTensorOp, + kStages, + TransformA, + TransformB, + Operator, + kCacheOpA, + kCacheOpB + >; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename MmaComplexCore::SmemLayoutA; + + // Shared memory layout + using SmemLayoutB = typename MmaComplexCore::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename MmaComplexCore::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename MmaComplexCore::SmemIteratorA; + + /// ThreadMap of iterator B + using IteratorThreadMapB = typename MmaComplexCore::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename MmaComplexCore::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename MmaComplexCore::MmaTensorOp; + + /// Policy used to define MmaPipelined + using MmaPolicy = typename MmaComplexCore::MmaPolicy; +}; + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(ElementA))>; + + // Shared memory layout + using SmemLayoutB = layout::RowMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(ElementB))>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousA = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedA = + kWarpSize / kWarpThreadArrangementContiguousA; + + static int const kWarpThreadArrangementContiguousB = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedB = + kWarpSize / kWarpThreadArrangementContiguousB; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // Shared memory layout + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousB = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedB = + kWarpSize / kWarpThreadArrangementContiguousB; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(ElementA))>; + + // Shared memory layout + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousA = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedA = + kWarpSize / kWarpThreadArrangementContiguousA; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK>; + + // Shared memory layout + using SmemLayoutB = layout::RowMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(ElementB))>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major-interleaved +/// B: row-major-interleaved +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +/// +/// Column/RowMajorInterleved(m, n) is mapped to Column/RowMajor(m +/// x InterleavedK, n / InterleavedK) so that Column/RowMajor global iterators +/// can be reused. The shared store iterator is the same as the crosswise shared +/// store iterator. So, the only thing we need to do is to swap the coordinates +/// (contiguous <=> strided) used by the global iterator and the shared store +/// iterator. +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// Number of interleaved K + int InterleavedK> +struct DefaultMmaCore, ElementB_, + layout::RowMajorInterleaved, ElementC_, + LayoutC_, arch::OpClassTensorOp, Stages, Operator_, + AccumulatorsInRowMajor, CacheOpA, CacheOpB> { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajorInterleaved; + using ElementB = ElementB_; + using LayoutB = layout::RowMajorInterleaved; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + static int const kInterleavedK = InterleavedK; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kElementsPerAccess = + kAccessSizeInBits / sizeof_bits::value; + + static int const kWarpThreadArrangementContiguous = + kInterleavedK / kElementsPerAccess; + + static int const kWarpThreadArrangementStrided = + kWarpSize / kWarpThreadArrangementContiguous; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, kInterleavedK>; + + // Shared memory layout + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, kInterleavedK>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, layout::PitchLinearShape<64, 1>, kElementsPerAccess>; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapA = transform::TransposePitchLinearThreadMap< + IteratorThreadMapA, + layout::PitchLinearShape>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + SmemThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, + kThreads, layout::PitchLinearShape<64, 1>, kElementsPerAccess>; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapB = transform::TransposePitchLinearThreadMap< + IteratorThreadMapB, + layout::PitchLinearShape>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + SmemThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK, AccumulatorsInRowMajor>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for SIMT GEMMs using multistage pipeline. +/// +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by Simt + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kElementsPerAccess = 1; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + + // Shared memory layout + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory + // + + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + IteratorThreadMapA>; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Transpose the ThreadMap of iterator B + using SmemThreadMapB = transform::TransposePitchLinearThreadMapSimt; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + SmemThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = 4; + static const int WarpNumThreadsN = 8; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + + static_assert(!((Shape::kK / 32) % LaneN), + "Padding must be divisible by Lane"); + + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + >; /// Used for partial specialization + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape<0, 0>, + MatrixShape<0, Shape::kK / 32>, + WarpCount::kK>; +}; + +/// Partial specialization for SIMT GEMMs using multistage pipeline. +/// +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by Simt + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kElementsPerAccess = 1; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + + // Shared memory layout + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory + // + + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + IteratorThreadMapA>; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = 4; + static const int WarpNumThreadsN = 8; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + >; /// Used for partial specialization + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK>; +}; + +/// Partial specialization for SIMT GEMMs using multistage pipeline. +/// +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by Simt + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kElementsPerAccess = 1; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + + // Shared memory layout + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory + // + + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapA = transform::TransposePitchLinearThreadMapSimt; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + SmemThreadMapA>; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Transpose the ThreadMap of iterator B + using SmemThreadMapB = transform::TransposePitchLinearThreadMapSimt; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + SmemThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = 4; + static const int WarpNumThreadsN = 8; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + + static_assert(!((Shape::kK / 32) % LaneM) && !((Shape::kK / 32) % LaneN), + "Padding must be divisible by Lane"); + + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + >; /// Used for partial specialization + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape, + MatrixShape<0, Shape::kK / 32>, + WarpCount::kK>; +}; + +/// Partial specialization for SIMT GEMMs using multistage pipeline. +/// +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by Simt + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kElementsPerAccess = 1; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + + // Shared memory layout + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapA = transform::TransposePitchLinearThreadMapSimt; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + SmemThreadMapA>; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = 4; + static const int WarpNumThreadsN = 8; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + + static_assert(!((Shape::kK / 32) % LaneM), + "Padding must be divisible by Lane"); + + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + >; /// Used for partial specialization + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape, + MatrixShape<0, 0>, + WarpCount::kK>; +}; + +/// Partial specialization for SIMT GEMMs using multistage pipeline. +/// +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by Simt + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::AffineRank2ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::AffineRank2RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Default Operator + using Operator = Operator_; + + using Base = DefaultMmaCore; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename Base::SmemLayoutA; + using SmemLayoutB = typename Base::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename Base::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename Base::SmemIteratorA; + + /// Policy of iterator B + using IteratorThreadMapB = typename Base::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename Base::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + /// Policy used to define MmaPipelined + using MmaPolicy = typename Base::MmaPolicy; +}; + +/// Partial specialization for SIMT GEMMs using multistage pipeline. +/// +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by Simt + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::AffineRank2RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::AffineRank2ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Default Operator + using Operator = Operator_; + + using Base = DefaultMmaCore; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename Base::SmemLayoutA; + using SmemLayoutB = typename Base::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename Base::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename Base::SmemIteratorA; + + /// Policy of iterator B + using IteratorThreadMapB = typename Base::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename Base::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + /// Policy used to define MmaPipelined + using MmaPolicy = typename Base::MmaPolicy; +}; + +/// Partial specialization for SIMT GEMMs using multistage pipeline. +/// +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by Simt + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::AffineRank2ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::AffineRank2ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Default Operator + using Operator = Operator_; + + using Base = DefaultMmaCore; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename Base::SmemLayoutA; + using SmemLayoutB = typename Base::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename Base::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename Base::SmemIteratorA; + + /// Policy of iterator B + using IteratorThreadMapB = typename Base::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename Base::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + /// Policy used to define MmaPipelined + using MmaPolicy = typename Base::MmaPolicy; + +}; + +/// Partial specialization for SIMT GEMMs using multistage pipeline. +/// +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by Simt + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::AffineRank2RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::AffineRank2RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Default Operator + using Operator = Operator_; + + using Base = DefaultMmaCore; + + // + // Shared memory layouts + // + + using SmemLayoutA = typename Base::SmemLayoutA; + using SmemLayoutB = typename Base::SmemLayoutB; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = typename Base::IteratorThreadMapA; + + /// Shared memory iterator to A operand + using SmemIteratorA = typename Base::SmemIteratorA; + + /// Policy of iterator B + using IteratorThreadMapB = typename Base::IteratorThreadMapB; + + /// Shared memory iterator to B operand + using SmemIteratorB = typename Base::SmemIteratorB; + + // + // Warp-level matrix multiply operator + // + + /// Policy used to define MmaPipelined + using MmaPolicy = typename Base::MmaPolicy; + +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sparse_sm80.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sparse_sm80.h new file mode 100644 index 0000000..285577f --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_sparse_sm80.h @@ -0,0 +1,834 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Defines basic properties needed by CTA-level GEMMs assuming + expectations about data layout of the global memory fragments, data types, + and internal tile sizes. + + Partial specializations for threadblock::Mma operations targeting sparse + TensorOp instructions. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/mctlass.h" + +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" +#include "mctlass/layout/tensor_op_multiplicand_sm80.h" + +#include "mctlass/gemm/warp/mma_simt_policy.h" +#include "mctlass/gemm/warp/mma_simt.h" +#include "mctlass/gemm/warp/default_mma_sparse_tensor_op.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" + +#include "mctlass/gemm/threadblock/default_mma_core.h" + +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op_sm80.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear.h" +#include "mctlass/gemm/threadblock/mma_sparse_multistage.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Template defininng default matrix multiply operators inferred from threadblock tile size, +/// global memory data layout, and target math instruction. +template < + /// Shape of threadblock-scoped matrix multiply operator + typename Shape, + /// Shape of warp-level matrix multiply operator + typename WarpShape, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape, + /// Element data type of A operand + typename ElementA, + /// Layout of operand A + typename LayoutA, + /// Element data type of B operand + typename ElementB, + /// Layout of operand B + typename LayoutB, + /// Data type of accumulator + typename ElementC, + /// Layout of accumulator + typename LayoutC, + /// Indicates type of math operator (arch::OpClassSimt or arch::OpClassTensorOp) + typename OperatorClass, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator = typename platform::conditional< + (platform::is_same::value) && + (platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value), + mctlass::arch::OpMultiplyAddSaturate, + mctlass::arch::OpMultiplyAdd>::type, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false + /// Cache operation of operand A + , mctlass::arch::CacheOperation::Kind CacheOpA = + mctlass::arch::CacheOperation::Global, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB = + mctlass::arch::CacheOperation::Global +> +struct DefaultSparseMmaCore; + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultSparseMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + static int const kSparse = 2; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(ElementA))>; + + // Shared memory layout + using SmemLayoutB = layout::RowMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(ElementB))>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultSparseMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Cache operation of operand E + static mctlass::arch::CacheOperation::Kind const kCacheOpE = + mctlass::arch::CacheOperation::Global; + + static int const kInterleavedE = MmaTensorOp::kInterleaved; + static int const kMetaSizeInBits = MmaTensorOp::kMetaSizeInBits; + static int const kMaxID2 = MmaTensorOp::kMaxID2; + static int const kElementsPerElementE = MmaTensorOp::kElementsPerElementE; + + using ElementE = typename MmaTensorOp::ElementE; + using GmemLayoutE = mctlass::layout::ColumnMajorInterleaved; + + // Shared memory layout. Interleaved layout is mapped to PitchLinear layout. + using SmemLayoutE = typename MmaTensorOp::LayoutE; + + /// ThreadMap of iterator E + static int const kElementsPerAccessE = + kAccessSizeInBits / sizeof_bits::value; + + /// E is tiny. Not all warps are needed. + static int const kThreadsE = + (Shape::kM * Shape::kK / kSparse / kElementsPerElementE / + (kAccessSizeInBits / sizeof_bits::value) > + kThreads) + ? kThreads + : (Shape::kM * Shape::kK / kSparse / kElementsPerElementE / + (kAccessSizeInBits / sizeof_bits::value)); + + using IteratorThreadMapE = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreadsE, kElementsPerAccessE>; + + /// Shared memory iterator to E operand + using SmemIteratorE = transform::threadblock::RegularTileAccessIterator< + MatrixShape, + ElementE, SmemLayoutE, 0, IteratorThreadMapE>; + + /// Policy used to define MmaPipelined + using MmaPolicy = + SparseMmaPolicy, MatrixShape<0, 0>, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultSparseMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + static int const kSparse = 2; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousA = + Shape::kK / kSparse / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedA = + kWarpSize / kWarpThreadArrangementContiguousA; + + // crosswise cannot be larger than 1024 bit. + static int const kCrosswiseB = + (Shape::kK > (1024 / sizeof_bits::value)) + ? (1024 / sizeof_bits::value) + : Shape::kK; + + static int const kWarpThreadArrangementContiguousB = + kCrosswiseB / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedB = + kWarpSize / kWarpThreadArrangementContiguousB; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK / kSparse>; + + // Shared memory layout + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, kCrosswiseB>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultSparseMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Cache operation of operand E + static mctlass::arch::CacheOperation::Kind const kCacheOpE = + mctlass::arch::CacheOperation::Global; + + static int const kInterleavedE = MmaTensorOp::kInterleaved; + static int const kMetaSizeInBits = MmaTensorOp::kMetaSizeInBits; + static int const kMaxID2 = MmaTensorOp::kMaxID2; + static int const kElementsPerElementE = MmaTensorOp::kElementsPerElementE; + + using ElementE = typename MmaTensorOp::ElementE; + using GmemLayoutE = mctlass::layout::ColumnMajorInterleaved; + + // Shared memory layout. Interleaved layout is mapped to PitchLinear layout. + using SmemLayoutE = typename MmaTensorOp::LayoutE; + + /// ThreadMap of iterator E + static int const kElementsPerAccessE = + kAccessSizeInBits / sizeof_bits::value; + + /// E is tiny. Not all warps are needed. + static int const kThreadsE = + (Shape::kM * Shape::kK / kSparse / kElementsPerElementE / + (kAccessSizeInBits / sizeof_bits::value) > + kThreads) + ? kThreads + : (Shape::kM * Shape::kK / kSparse / kElementsPerElementE / + (kAccessSizeInBits / sizeof_bits::value)); + + using IteratorThreadMapE = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreadsE, kElementsPerAccessE>; + + + /// Shared memory iterator to E operand + using SmemIteratorE = transform::threadblock::RegularTileAccessIterator< + MatrixShape, + ElementE, SmemLayoutE, 0, IteratorThreadMapE>; + + /// Policy used to define MmaPipelined + using MmaPolicy = + SparseMmaPolicy, MatrixShape<0, 0>, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultSparseMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + static int const kSparse = 2; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + // crosswise cannot be larger than 1024 bit. + static int const kCrosswiseB = + (Shape::kK > (1024 / sizeof_bits::value)) + ? (1024 / sizeof_bits::value) + : Shape::kK; + + static int const kWarpThreadArrangementContiguousB = + kCrosswiseB / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedB = + kWarpSize / kWarpThreadArrangementContiguousB; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(ElementA))>; + + // Shared memory layout + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, kCrosswiseB>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultSparseMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Cache operation of operand E + static mctlass::arch::CacheOperation::Kind const kCacheOpE = + mctlass::arch::CacheOperation::Global; + + static int const kInterleavedE = MmaTensorOp::kInterleaved; + static int const kMetaSizeInBits = MmaTensorOp::kMetaSizeInBits; + static int const kMaxID2 = MmaTensorOp::kMaxID2; + static int const kElementsPerElementE = MmaTensorOp::kElementsPerElementE; + + using ElementE = typename MmaTensorOp::ElementE; + using GmemLayoutE = mctlass::layout::ColumnMajorInterleaved; + + // Shared memory layout. Interleaved layout is mapped to PitchLinear layout. + using SmemLayoutE = typename MmaTensorOp::LayoutE; + + /// ThreadMap of iterator E + static int const kElementsPerAccessE = + kAccessSizeInBits / sizeof_bits::value; + + /// E is tiny. Not all warps are needed. + static int const kThreadsE = + (Shape::kM * Shape::kK / kSparse / kElementsPerElementE / + (kAccessSizeInBits / sizeof_bits::value) > + kThreads) + ? kThreads + : (Shape::kM * Shape::kK / kSparse / kElementsPerElementE / + (kAccessSizeInBits / sizeof_bits::value)); + + using IteratorThreadMapE = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreadsE, kElementsPerAccessE>; + + /// Shared memory iterator to E operand + using SmemIteratorE = transform::threadblock::RegularTileAccessIterator< + MatrixShape, + ElementE, SmemLayoutE, 0, IteratorThreadMapE>; + + /// Policy used to define MmaPipelined + using MmaPolicy = + SparseMmaPolicy, MatrixShape<0, 0>, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultSparseMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + static int const kSparse = 2; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousA = + Shape::kK / kSparse / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedA = + kWarpSize / kWarpThreadArrangementContiguousA; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kK / kSparse>; + + // Shared memory layout + using SmemLayoutB = layout::RowMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(ElementB))>; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultSparseMmaTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, WarpCount::kK>::Type; + + /// Cache operation of operand E + static mctlass::arch::CacheOperation::Kind const kCacheOpE = + mctlass::arch::CacheOperation::Global; + + static int const kInterleavedE = MmaTensorOp::kInterleaved; + static int const kMetaSizeInBits = MmaTensorOp::kMetaSizeInBits; + static int const kMaxID2 = MmaTensorOp::kMaxID2; + static int const kElementsPerElementE = MmaTensorOp::kElementsPerElementE; + + using ElementE = typename MmaTensorOp::ElementE; + using GmemLayoutE = mctlass::layout::ColumnMajorInterleaved; + + // Shared memory layout. Interleaved layout is mapped to PitchLinear layout. + using SmemLayoutE = typename MmaTensorOp::LayoutE; + + /// ThreadMap of iterator E + static int const kElementsPerAccessE = + kAccessSizeInBits / sizeof_bits::value; + + /// E is tiny. Not all warps are needed. + static int const kThreadsE = + (Shape::kM * Shape::kK / kSparse / kElementsPerElementE / + (kAccessSizeInBits / sizeof_bits::value) > + kThreads) + ? kThreads + : (Shape::kM * Shape::kK / kSparse / kElementsPerElementE / + (kAccessSizeInBits / sizeof_bits::value)); + + using IteratorThreadMapE = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreadsE, kElementsPerAccessE>; + + /// Shared memory iterator to E operand + using SmemIteratorE = transform::threadblock::RegularTileAccessIterator< + MatrixShape, + ElementE, SmemLayoutE, 0, IteratorThreadMapE>; + + /// Policy used to define MmaPipelined + using MmaPolicy = + SparseMmaPolicy, MatrixShape<0, 0>, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_with_access_size.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_with_access_size.h new file mode 100644 index 0000000..d3c0e77 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_with_access_size.h @@ -0,0 +1,328 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic properties needed by CTA-level GEMMs assuming expectations about data + layout of the global memory fragments, data types, and internal tile sizes. + + Partial specializations for threadblock::Mma operations targeting simt instructions. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/warp/mma.h" +#include "mctlass/gemm/threadblock/mma_pipelined.h" +#include "mctlass/gemm/threadblock/mma_singlestage.h" +#include "mctlass/arch/cache_operation.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +template < + /// Shape of threadblock-scoped matrix multiply operator + typename Shape, + /// Shape of warp-level matrix multiply operator + typename WarpShape, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape, + /// Element data type of A operand + typename ElementA, + /// Layout of operand A + typename LayoutA, + /// Element data type of B operand + typename ElementB, + /// Layout of operand B + typename LayoutB, + /// Data type of accumulator + typename ElementC, + /// Layout of accumulator + typename LayoutC, + /// Indicates type of math operator (arch::OpClassSimt or arch::OpClassTensorOp) + typename OperatorClass, + /// Size of a threadblock-scoped access + int kAccessSizeInBits = -1, // -1 denoting the default + /// Number of stages + int Stages = 2, + /// Operation performed by MMA + typename Operator = typename platform::conditional< + (platform::is_same::value) && + (platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value), + mctlass::arch::OpMultiplyAddSaturate, + mctlass::arch::OpMultiplyAdd>::type, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA = + mctlass::arch::CacheOperation::Global, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB = + mctlass::arch::CacheOperation::Global, + /// per-element transformation for elements of A + ComplexTransform TransformA = ComplexTransform::kNone, + /// per-element transformation for elements of B + ComplexTransform TransformB = ComplexTransform::kNone, + bool IsComplex = false // (is_complex::value || is_complex::value) +> +struct DefaultMmaCoreWithAccessSize; + +template < + /// Shape of threadblock-scoped matrix multiply operator + typename Shape, + /// Shape of warp-level matrix multiply operator + typename WarpShape, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape, + /// Element data type of A operand + typename ElementA, + /// Layout of operand A + typename LayoutA, + /// Element data type of B operand + typename ElementB, + /// Layout of operand B + typename LayoutB, + /// Data type of accumulator + typename ElementC, + /// Layout of accumulator + typename LayoutC, + /// Indicates type of math operator (arch::OpClassSimt or arch::OpClassTensorOp) + typename OperatorClass, + /// Number of stages + int Stages, + /// Operation performed by MMA + typename Operator, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// per-element transformation for elements of A + ComplexTransform TransformA, + /// per-element transformation for elements of B + ComplexTransform TransformB, + bool IsComplex +> +struct DefaultMmaCoreWithAccessSize< + Shape, WarpShape, InstructionShape, + ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, + OperatorClass, -1, Stages, Operator, AccumulatorsInRowMajor, + CacheOpA, CacheOpB, TransformA, TransformB, IsComplex +> : DefaultMmaCore< + Shape, WarpShape, InstructionShape, + ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, + OperatorClass, Stages, Operator, AccumulatorsInRowMajor, + CacheOpA, CacheOpB, TransformA, TransformB, IsComplex +> {}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: row-major +/// Operator: simt class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Size of a threadblock-scoped access (a value of -1 indicates the default) + int kAccessSizeInBits_, + /// Operation performed by GEMM + typename Operator_> +struct DefaultMmaCoreWithAccessSize>::type, ElementA_, + layout::ColumnMajor, ElementB_, layout::RowMajor, + ElementC_, LayoutC_, arch::OpClassSimt, kAccessSizeInBits_, 2, Operator_ + > { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 1>; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassSimt; + static int const PartitionsK = Shape::kK / WarpShape::kK; + + /// Default Operator + using Operator = Operator_; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + PartitionsK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + static int const kElementsPerAccessDefault = 1; + static_assert(kAccessSizeInBits_ == -1 || + sizeof_bits::value == sizeof_bits::value || + kAccessSizeInBits_ / sizeof_bits::value == kElementsPerAccessDefault, + "Non-default value for kAccessSizeInBits_ is only allowed if size(elementA) == sizeof(elementB)"); + static int const kElementsPerAccess = (kAccessSizeInBits_ != -1) ? kAccessSizeInBits_ / sizeof_bits::value : kElementsPerAccessDefault; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + IteratorThreadMapA + >; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = detail::simt_get_warp_threads_m(); + static const int WarpNumThreadsN = kWarpSize / WarpNumThreadsM; + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy /// Policy describing warp-level MmaSimtOp (concept: MmaSimtOp policy) + >; /// Used for partial specialization + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape<0, 0>, + MatrixShape<0, 0>, + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +} // namespace threadblock +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_with_reduction.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_with_reduction.h new file mode 100644 index 0000000..9a96a3e --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_with_reduction.h @@ -0,0 +1,167 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Defines basic properties needed by CTA-level GEMMs assuming + expectations about data layout of the global memory fragments, data types, + and internal tile sizes. + + Partial specializations for threadblock::Mma operations targeting TensorOp + instructions. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/mctlass.h" + +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" +#include "mctlass/layout/tensor_op_multiplicand_sm80.h" + +#include "mctlass/gemm/warp/default_mma_with_reduction_tensor_op.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" + +#include "mctlass/gemm/threadblock/default_mma_core.h" + +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op_sm80.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear.h" +#include "mctlass/gemm/threadblock/mma_with_reduction_multistage.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Template defininng default matrix multiply operators inferred from threadblock tile size, +/// global memory data layout, and target math instruction. +template < + /// Shape of threadblock-scoped matrix multiply operator + typename Shape_, + /// Shape of warp-level matrix multiply operator + typename WarpShape, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape, + /// Element data type of A operand + typename ElementA, + /// Layout of operand A + typename LayoutA, + /// Element data type of B operand + typename ElementB, + /// Layout of operand B + typename LayoutB, + /// Data type of accumulator + typename ElementC, + /// Layout of accumulator + typename LayoutC, + /// Indicates type of math operator (arch::OpClassSimt or arch::OpClassTensorOp) + typename OperatorClass, + /// Reduce operand A or B along K dimension + bool ReduceKForA_, + /// Number of stages + int Stages = 2, + /// Operation performed by MMA + typename Operator = typename platform::conditional< + (platform::is_same::value) && + (platform::is_same::value || + platform::is_same::value || + platform::is_same::value || + platform::is_same::value), + mctlass::arch::OpMultiplyAddSaturate, + mctlass::arch::OpMultiplyAdd>::type, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA = + mctlass::arch::CacheOperation::Global, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB = + mctlass::arch::CacheOperation::Global, + /// per-element transformation for elements of A + ComplexTransform TransformA = ComplexTransform::kNone, + /// per-element transformation for elements of B + ComplexTransform TransformB = ComplexTransform::kNone, + bool IsComplex = false// (is_complex::value || is_complex::value) +> +struct DefaultMmaWithReductionCore { + using Base = DefaultMmaCore; + using Shape = Shape_; + using IteratorThreadMapA = typename Base::IteratorThreadMapA; + using IteratorThreadMapB = typename Base::IteratorThreadMapB; + using SmemIteratorA = typename Base::SmemIteratorA; + using SmemIteratorB = typename Base::SmemIteratorB; + using SmemLayoutA = typename Base::SmemLayoutA; + using SmemLayoutB = typename Base::SmemLayoutB; + using WarpCount = typename Base::WarpCount; + + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaWithReductionTensorOp< + WarpShape, InstructionShape, ElementA, SmemLayoutA, ElementB, SmemLayoutB, + ElementC, LayoutC, Operator, ReduceKForA_, WarpCount::kK>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_wmma.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_wmma.h new file mode 100644 index 0000000..60894e6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_core_wmma.h @@ -0,0 +1,712 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic properties needed by CTA-level GEMMs assuming expectations about data + layout of the global memory fragments, data types, and internal tile sizes. + + Partial specializations for threadblock::Mma operations targeting TensorOp instructions. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/fast_math.h" +#include "mctlass/arch/wmma.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h" + +#include "mctlass/gemm/warp/mma_tensor_op_wmma.h" + +#include "mctlass/gemm/warp/mma_tensor_op_policy.h" +#include "mctlass/gemm/threadblock/default_mma_core.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: row-major +/// Operator: wmma tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + ///< Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_, + /// Number of stages + int Stages> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassWmmaTensorOp; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + Shape::kK / WarpShape::kK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // + // Shared memory layouts + // + // NOTE: shared memory layout for wmma is same as the operands' layout in the global memory + using SmemLayoutA = LayoutA; + using SmemLayoutB = LayoutB; + + // Pad shared memory to avoid bank conflicts + static int const kPaddingA = 128 / sizeof_bits::value; + static int const kPaddingB = 128 / sizeof_bits::value; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + IteratorThreadMapA + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Wmma< + InstructionShape, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + Operator + >, + mctlass::MatrixShape<1, 1> + >; + + using MmaTensorOp = mctlass::gemm::warp::MmaTensorOpWmma< + WarpShape, + ElementA, + SmemLayoutA, + ElementB, + SmemLayoutB, + ElementC, + LayoutC, + Policy + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape, + MatrixShape<0, kPaddingB>, + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: column-major +/// Operator: wmma tensorop class +/// +/// This uses the default warp-level operator given tile sizes +template < + ///< Shape of threadblock-scoped matrix multiply operator + ///< (concept:GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) [allowed + /// wmma instruction shapes, e.g., 16x16x16, 32x8x16, 8x32x16,...] + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by GEMM + typename Operator_, + /// Number of stages + int Stages> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassWmmaTensorOp; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + Shape::kK / WarpShape::kK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads per threadblock + static int const kThreads = WarpCount::kCount * kWarpSize; + + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousA = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedA = + kWarpSize / kWarpThreadArrangementContiguousA; + + static int const kWarpThreadArrangementContiguousB = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedB = + kWarpSize / kWarpThreadArrangementContiguousB; + + // + // Shared memory layouts + // + + // shared memory layout for wmma is same as the operands' layout in global memory + using SmemLayoutA = LayoutA; + using SmemLayoutB = LayoutB; + + // Pad shared memory to avoid bank conflicts + static int const kPaddingA = 128 / sizeof_bits::value; + static int const kPaddingB = 128 / sizeof_bits::value; + + // + // Iterators to write to shared memory + // + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + IteratorThreadMapA + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + IteratorThreadMapB // SmemThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Wmma< + InstructionShape, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + Operator + >, + mctlass::MatrixShape<1, 1> + >; + + using MmaTensorOp = mctlass::gemm::warp::MmaTensorOpWmma< + WarpShape, + ElementA, + SmemLayoutA, + ElementB, + SmemLayoutB, + ElementC, + LayoutC, + Policy + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape<0, kPaddingA>, + MatrixShape, + WarpCount::kK + >; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: row-major +/// B: row-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by MMA + typename Operator_, + /// Number of stages + int Stages> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::RowMajor; + using ElementB = ElementB_; + using LayoutB = layout::RowMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassWmmaTensorOp; + + /// Number of warps present + using WarpCount = GemmShape< + Shape::kM / WarpShape::kM, + Shape::kN / WarpShape::kN, + Shape::kK / WarpShape::kK + >; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && + !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size." + ); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousA = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedA = + kWarpSize / kWarpThreadArrangementContiguousA; + + // + // Shared memory layouts + // + + // shared memory layout for wmma is same as the operands' layout in global memory + using SmemLayoutA = LayoutA; + using SmemLayoutB = LayoutB; + + // Pad shared memory to avoid bank conflicts + static int const kPaddingA = 128 / sizeof_bits::value; + static int const kPaddingB = 128 / sizeof_bits::value; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kAccessSizeInBits / sizeof_bits::value + >; + + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementA, + SmemLayoutA, + 1, + IteratorThreadMapA + >; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, + ElementB, + SmemLayoutB, + 0, + IteratorThreadMapB + >; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Wmma< + InstructionShape, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + Operator + >, + mctlass::MatrixShape<1, 1> + >; + + using MmaTensorOp = mctlass::gemm::warp::MmaTensorOpWmma< + WarpShape, + ElementA, + SmemLayoutA, + ElementB, + SmemLayoutB, + ElementC, + LayoutC, + Policy + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape<0, kPaddingA>, + MatrixShape<0, kPaddingB>, + WarpCount::kK + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization: +/// +/// A: column-major +/// B: column-major +/// Operator: tensor op class +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A operand + typename ElementA_, + /// Data type of B operand + typename ElementB_, + /// Data type of accumulator + typename ElementC_, + /// Layout of accumulator + typename LayoutC_, + /// Operation performed by MMA + typename Operator_, + /// Number of stages + int Stages> +struct DefaultMmaCore { + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = ElementA_; + using LayoutA = layout::ColumnMajor; + using ElementB = ElementB_; + using LayoutB = layout::ColumnMajor; + using ElementC = ElementC_; + using LayoutC = LayoutC_; + using OperatorClass = arch::OpClassWmmaTensorOp; + + /// Number of warps present + using WarpCount = + GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped access + static int const kAccessSizeInBits = 128; + + /// Default Operator + using Operator = Operator_; + + // Warp thread arrangement + static int const kWarpThreadArrangementContiguousB = + Shape::kK / (kAccessSizeInBits / sizeof_bits::value); + + static int const kWarpThreadArrangementStridedB = + kWarpSize / kWarpThreadArrangementContiguousB; + + // + // Shared memory layouts + // + + // shared memory layout for wmma is same as the operands' layout in global memory + using SmemLayoutA = LayoutA; + using SmemLayoutB = LayoutB; + + // Pad shared memory to avoid bank conflicts + static int const kPaddingA = 128 / sizeof_bits::value; + static int const kPaddingB = 128 / sizeof_bits::value; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kAccessSizeInBits / sizeof_bits::value + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Wmma< + InstructionShape, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + Operator + >, + mctlass::MatrixShape<1, 1> + >; + + using MmaTensorOp = mctlass::gemm::warp::MmaTensorOpWmma< + WarpShape, + ElementA, + SmemLayoutA, + ElementB, + SmemLayoutB, + ElementC, + LayoutC, + Policy + >; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaTensorOp, + MatrixShape, + MatrixShape, + WarpCount::kK + >; +}; + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +#endif // defined(MCTLASS_ARCH_WMMA_ENABLED) + diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_layernorm_mainloop_fusion.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_layernorm_mainloop_fusion.h new file mode 100644 index 0000000..a00fb53 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_layernorm_mainloop_fusion.h @@ -0,0 +1,178 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/gemm/threadblock/default_mma_core.h" +#include "mctlass/gemm/threadblock/mma_layernorm_mainloop_fusion_multistage.h" +#include "mctlass/transform/threadblock/predicated_scale_bias_vector_iterator.h" +#include "mctlass/transform/threadblock/predicated_scale_bias_vector_access_iterator.h" +#include "mctlass/transform/threadblock/regular_scale_bias_vector_access_iterator.h" +#include "mctlass/gemm/warp/scale_bias_tile_iterator.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for Scale/Bias vectors + typename ElementScaleBias, + /// Layout type for Scale/Bias vectors + typename LayoutScaleBias, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Use zfill or predicate for SM80 out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone + > +struct DefaultMmaLayernormMainloopFusion { + + static mctlass::arch::CacheOperation::Kind const CacheOpA = + ((sizeof_bits::value * kAlignmentA) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * kAlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpGammaBeta = CacheOpA; + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, Operator, false, CacheOpA, CacheOpB>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, AccessTypeB>; + + /// Define iterators over tiles from scale/bias vectors + using IteratorVarMean = + mctlass::transform::threadblock::PredicatedScaleBiasVectorIterator< + mctlass::MatrixShape<1, WarpShape::kN>, + ElementScaleBias, + LayoutScaleBias>; + + /// Define iterators over tiles from scale/bias vectors + using IteratorGammaBeta = + mctlass::transform::threadblock::PredicatedScaleBiasVectorAccessIterator< + mctlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + using SmemIteratorGammaBeta = + mctlass::transform::threadblock::RegularScaleBiasVectorAccessIterator< + mctlass::MatrixShape<1, ThreadblockShape::kK>, ElementScaleBias, + LayoutScaleBias>; + + static int const kThreadCount = 32; + + // Warp-level iterators to load scale and bias vectors + using WarpIteratorGammaBeta = mctlass::gemm::warp::ScaleBiasTileIterator< + MatrixShape, ElementScaleBias, + LayoutScaleBias, MatrixShape, + typename MmaCore::MmaTensorOp::IteratorA::Base::Policy, kThreadCount, + MmaCore::WarpCount::kK>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaLayernormMainloopFusionMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, IteratorVarMean, IteratorGammaBeta, SmemIteratorGammaBeta, + CacheOpGammaBeta, + ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, WarpIteratorGammaBeta, Stages, SharedMemoryClear>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_planar_complex_multistage.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_planar_complex_multistage.h new file mode 100644 index 0000000..a1a2110 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_planar_complex_multistage.h @@ -0,0 +1,136 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Template for a multistage GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/arch/arch.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/mma_planar_complex_multistage.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex transformation on operand B + ComplexTransform TransformB = ComplexTransform::kNone, + /// Math operator tag (e.g. arch::OpMultiplyAdd) + typename Operator = arch::OpMultiplyAdd +> +struct DefaultMmaPlanarComplexMultistage { + + // Construct a planar complex variant from the real-valued variant + using RealMmaMultistage = typename DefaultMma< + ElementA_, + LayoutA_, + kAlignmentA, + ElementB_, + LayoutB_, + kAlignmentB, + ElementAccumulator_, + LayoutC_, + OperatorClass_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + Stages, + Operator + >::ThreadblockMma; + + using ThreadblockMma = MmaPlanarComplexMultistage< + ThreadblockShape_, + typename RealMmaMultistage::IteratorA, + typename RealMmaMultistage::SmemIteratorA, + mctlass::arch::CacheOperation::Global, + typename RealMmaMultistage::IteratorB, + typename RealMmaMultistage::SmemIteratorB, + mctlass::arch::CacheOperation::Global, + ElementAccumulator_, + LayoutC_, + typename RealMmaMultistage::Policy, + Stages, + TransformA, + TransformB + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_planar_complex_pipelined.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_planar_complex_pipelined.h new file mode 100644 index 0000000..bcbe5b4 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_planar_complex_pipelined.h @@ -0,0 +1,130 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +#include "mctlass/gemm/warp/mma_planar_complex.h" +#include "mctlass/gemm/threadblock/default_mma.h" +#include "mctlass/gemm/threadblock/mma_planar_complex_pipelined.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex transformation on operand B + ComplexTransform TransformB = ComplexTransform::kNone, + /// Math operator tag (e.g. arch::OpMultiplyAdd) + typename Operator = arch::OpMultiplyAdd +> +struct DefaultMmaPlanarComplexPipelined { + + // Construct a planar complex variant from the real-valued variant + using RealMma = typename DefaultMma< + ElementA_, + LayoutA_, + kAlignmentA, + ElementB_, + LayoutB_, + kAlignmentB, + ElementAccumulator_, + LayoutC_, + OperatorClass_, + ArchTag_, + ThreadblockShape_, + WarpShape_, + InstructionShape_, + Stages, + Operator + >::ThreadblockMma; + + using ThreadblockMma = MmaPlanarComplexPipelined< + ThreadblockShape_, + typename RealMma::IteratorA, + typename RealMma::SmemIteratorA, + typename RealMma::IteratorB, + typename RealMma::SmemIteratorB, + ElementAccumulator_, + LayoutC_, + typename RealMma::Policy, + Stages, + TransformA, + TransformB + >; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_softmax_mainloop_fusion.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_softmax_mainloop_fusion.h new file mode 100644 index 0000000..faeb8d2 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_softmax_mainloop_fusion.h @@ -0,0 +1,160 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined softmax-GEMM kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/gemm/threadblock/default_mma_core.h" +#include "mctlass/gemm/threadblock/mma_softmax_mainloop_fusion_multistage.h" +#include "mctlass/transform/threadblock/predicated_scale_bias_vector_iterator.h" +#include "mctlass/transform/threadblock/predicated_scale_bias_vector_access_iterator.h" +#include "mctlass/transform/threadblock/regular_scale_bias_vector_access_iterator.h" +#include "mctlass/gemm/warp/scale_bias_tile_iterator.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for Scale/Bias vectors + typename ElementScaleBias, + /// Layout type for Scale/Bias vectors + typename LayoutScaleBias, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Operator class tag + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Whether problem has been transformed. This determines to which operand + /// the softmax is applied. + bool InternalTranspose, + /// Operation perfomed by GEMM + typename Operator, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Use zfill or predicate for SM80 out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone + > +struct DefaultMmaSoftmaxMainloopFusion { + + static mctlass::arch::CacheOperation::Kind const CacheOpA = + ((sizeof_bits::value * kAlignmentA) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * kAlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpGammaBeta = CacheOpA; + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, Operator, false, CacheOpA, CacheOpB>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, AccessTypeB>; + + /// Define iterators over tiles from scale/bias vectors + using IteratorNormSum = + mctlass::transform::threadblock::PredicatedScaleBiasVectorIterator< + mctlass::MatrixShape<1, WarpShape::kN>, + ElementScaleBias, + LayoutScaleBias>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaSoftmaxMainloopFusionMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, IteratorNormSum, + ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages, InternalTranspose, SharedMemoryClear>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_with_reduction.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_with_reduction.h new file mode 100644 index 0000000..6b3e01a --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_mma_with_reduction.h @@ -0,0 +1,141 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator_2dthreadtile.h" +#include "mctlass/gemm/threadblock/default_mma_core_with_reduction.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Layout type for C and D matrix operands + typename LayoutC, + /// Operator class tag + typename OperatorClass, + /// + bool ReduceKForA_, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Use zfill or predicate for SM80 out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone + > +struct DefaultMmaWithReduction { + + static mctlass::arch::CacheOperation::Kind const CacheOpA = + ((sizeof_bits::value * kAlignmentA) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * kAlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaWithReductionCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + ReduceKForA_, Stages, Operator, false, CacheOpA, CacheOpB>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaWithReductionMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages, SharedMemoryClear>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_mma_complex.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_mma_complex.h new file mode 100644 index 0000000..ffb2675 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_mma_complex.h @@ -0,0 +1,159 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Template for a multistage GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/arch/arch.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/numeric_types.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" +#include "mctlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex transformation on operand B + ComplexTransform TransformB = ComplexTransform::kNone, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator = arch::OpMultiplyAddComplex, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false> +struct DefaultMultistageMmaComplex; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator> +struct DefaultMultistageMmaComplex { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplexCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, TransformA, TransformB, Operator>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages>; +}; + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_mma_complex_core.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_mma_complex_core.h new file mode 100644 index 0000000..49f847a --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_mma_complex_core.h @@ -0,0 +1,119 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic properties needed by CTA-level GEMMs assuming + expectations about data layout of the global memory fragments, data types, + and internal tile sizes. + + Partial specializations for threadblock::Mma operations targeting TensorOp + instructions. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/complex.h" + +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" +#include "mctlass/layout/tensor_op_multiplicand_sm80.h" + +#include "mctlass/gemm/warp/mma_simt_policy.h" +#include "mctlass/gemm/warp/mma_simt.h" +#include "mctlass/gemm/warp/default_mma_tensor_op.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" + +#include "mctlass/gemm/threadblock/default_mma_core.h" + +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" +#include "mctlass/transform/pitch_linear_thread_map.h" + +#include "mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op_sm80.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Template defininng default matrix multiply operators inferred from +/// threadblock tile size, global memory data layout, and target math +/// instruction. +template < + /// Shape of threadblock-scoped matrix multiply operator + typename Shape, + /// Shape of warp-level matrix multiply operator + typename WarpShape, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape, + /// Element data type of A operand + typename ElementA, + /// Layout of operand A + typename LayoutA, + /// Element data type of B operand + typename ElementB, + /// Layout of operand B + typename LayoutB, + /// Data type of accumulator + typename ElementC, + /// Layout of accumulator + typename LayoutC, + /// Indicates type of math operator (arch::OpClassSimt or arch::OpClassTensorOp) + typename OperatorClass, + /// Number of stages + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator = arch::OpMultiplyAddComplex, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA = + mctlass::arch::CacheOperation::Global, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB = + mctlass::arch::CacheOperation::Global> +struct DefaultMultistageMmaComplexCore; + + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h new file mode 100644 index 0000000..57fb7d6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_mma_complex_core_sm80.h @@ -0,0 +1,1808 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic properties needed by CTA-level GEMMs assuming + expectations about data layout of the global memory fragments, data types, + and internal tile sizes. + + Partial specializations for threadblock::Mma operations targeting TensorOp + instructions. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/mctlass.h" + +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" +#include "mctlass/layout/tensor_op_multiplicand_sm80.h" + +#include "mctlass/gemm/warp/mma_simt_policy.h" +#include "mctlass/gemm/warp/mma_simt.h" +#include "mctlass/gemm/warp/default_mma_complex_tensor_op.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" + +#include "mctlass/gemm/threadblock/default_multistage_mma_complex_core.h" + +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" +#include "mctlass/transform/pitch_linear_thread_map.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op_sm80.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear.h" +#include "mctlass/gemm/threadblock/mma_multistage.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex double-precision +/// +/// A: column-major +/// B: row-major +/// Operator: arch::OpMultiplyAddComplex or arch::OpMultiplyGaussianComplex +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMultistageMmaComplexCore< + Shape_, WarpShape_, InstructionShape_, + complex, layout::ColumnMajor, + complex, layout::RowMajor, + complex, LayoutC_, + arch::OpClassTensorOp, + Stages, + TransformA, TransformB, + Operator_, + CacheOpA, CacheOpB> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = complex; + using LayoutA = layout::ColumnMajor; + using ElementB = complex; + using LayoutB = layout::RowMajor; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + using Operator = Operator_; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize / 2; + + /// Size of a threadblock-scoped 128 + static int const kAccessSizeInBits = 128; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorTensorOpMultiplicandCongruous128b; + + using SmemLayoutB = layout::RowMajorTensorOpMultiplicandCongruous128b; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaComplexTensorOp< + WarpShape, InstructionShape, + ElementA, SmemLayoutA, + ElementB, SmemLayoutB, + ElementC, LayoutC, + kTransformA, kTransformB, + Operator>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + + +/// Partial specialization for complex double-precision +/// +/// A: column-major +/// B: row-major +/// Operator: arch::OpMultiplyAddComplex or arch::OpMultiplyGaussianComplex +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMultistageMmaComplexCore< + Shape_, WarpShape_, InstructionShape_, + complex, layout::ColumnMajor, + complex, layout::ColumnMajor, + complex, LayoutC_, + arch::OpClassTensorOp, + Stages, + TransformA, TransformB, + Operator_, + CacheOpA, CacheOpB> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = complex; + using LayoutA = layout::ColumnMajor; + using ElementB = complex; + using LayoutB = layout::ColumnMajor; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + using Operator = Operator_; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize / 2; + + /// Size of a threadblock-scoped 128 + static int const kAccessSizeInBits = 128; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorTensorOpMultiplicandCongruous128b; + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicandCrosswise128x4; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaComplexTensorOp< + WarpShape, InstructionShape, + ElementA, SmemLayoutA, + ElementB, SmemLayoutB, + ElementC, LayoutC, + kTransformA, kTransformB, + Operator>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex double-precision +/// +/// A: row-major +/// B: column-major +/// Operator: arch::OpMultiplyAddComplex or arch::OpMultiplyGaussianComplex +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMultistageMmaComplexCore< + Shape_, WarpShape_, InstructionShape_, + complex, layout::RowMajor, + complex, layout::ColumnMajor, + complex, LayoutC_, + arch::OpClassTensorOp, + Stages, + TransformA, TransformB, + Operator_, + CacheOpA, CacheOpB> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = complex; + using LayoutA = layout::RowMajor; + using ElementB = complex; + using LayoutB = layout::ColumnMajor; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + using Operator = Operator_; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize / 2; + + /// Size of a threadblock-scoped 128 + static int const kAccessSizeInBits = 128; + + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicandCrosswise128x4; + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicandCrosswise128x4; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaComplexTensorOp< + WarpShape, InstructionShape, + ElementA, SmemLayoutA, + ElementB, SmemLayoutB, + ElementC, LayoutC, + kTransformA, kTransformB, + Operator>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + + +/// Partial specialization for complex double-precision +/// +/// A: row-major +/// B: row-major +/// Operator: arch::OpMultiplyAddComplex or arch::OpMultiplyGaussianComplex +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMultistageMmaComplexCore< + Shape_, WarpShape_, InstructionShape_, + complex, layout::RowMajor, + complex, layout::RowMajor, + complex, LayoutC_, + arch::OpClassTensorOp, + Stages, + TransformA, TransformB, + Operator_, + CacheOpA, CacheOpB> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = InstructionShape_; + using ElementA = complex; + using LayoutA = layout::RowMajor; + using ElementB = complex; + using LayoutB = layout::RowMajor; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + using Operator = Operator_; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize / 2; + + /// Size of a threadblock-scoped 128 + static int const kAccessSizeInBits = 128; + + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicandCrosswise128x4; + using SmemLayoutB = layout::RowMajorTensorOpMultiplicandCongruous128b; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<8, 4>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaComplexTensorOp< + WarpShape, InstructionShape, + ElementA, SmemLayoutA, + ElementB, SmemLayoutB, + ElementC, LayoutC, + kTransformA, kTransformB, + Operator>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex floating-point +/// +/// A: column-major +/// B: column-major +/// Operator: arch::OpMultiplyAddComplex +/// Math Instruction: mma.sync.aligned.m16n8k8.f32.tf32.tf32.f32 +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex) + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMultistageMmaComplexCore< + Shape_, WarpShape_, GemmShape<16, 8, 8>, + complex, layout::ColumnMajor, + complex, layout::ColumnMajor, + complex, LayoutC_, + arch::OpClassTensorOp, + Stages, + TransformA, TransformB, + Operator_, + CacheOpA, CacheOpB> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<16, 8, 8>; + using ElementA = complex; + using LayoutA = layout::ColumnMajor; + using ElementB = complex; + using LayoutB = layout::ColumnMajor; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + using Operator = Operator_; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped + static int const kAccessSizeInBits = 64; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorTensorOpMultiplicandCongruous64b; + + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicand64bCrosswise; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpStripedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaComplexTensorOp< + WarpShape, InstructionShape, + ElementA, SmemLayoutA, + ElementB, SmemLayoutB, + ElementC, LayoutC, + kTransformA, kTransformB, + Operator>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + + +/// Partial specialization for complex floating-point +/// +/// A: column-major +/// B: row-major +/// Operator: arch::OpMultiplyAddComplex +/// Math Instruction: mma.sync.aligned.m16n8k8.f32.tf32.tf32.f32 +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex) + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMultistageMmaComplexCore< + Shape_, WarpShape_, GemmShape<16, 8, 8>, + complex, layout::ColumnMajor, + complex, layout::RowMajor, + complex, LayoutC_, + arch::OpClassTensorOp, + Stages, + TransformA, TransformB, + Operator_, + CacheOpA, CacheOpB> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<16, 8, 8>; + using ElementA = complex; + using LayoutA = layout::ColumnMajor; + using ElementB = complex; + using LayoutB = layout::RowMajor; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + using Operator = Operator_; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped + static int const kAccessSizeInBits = 64; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajorTensorOpMultiplicandCongruous64b; + + using SmemLayoutB = layout::RowMajorTensorOpMultiplicandCongruous64b; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpStripedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpStripedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaComplexTensorOp< + WarpShape, InstructionShape, + ElementA, SmemLayoutA, + ElementB, SmemLayoutB, + ElementC, LayoutC, + kTransformA, kTransformB, + Operator>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex floating-point +/// +/// A: row-major +/// B: column-major +/// Operator: arch::OpMultiplyAddComplex +/// Math Instruction: mma.sync.aligned.m16n8k8.f32.tf32.tf32.f32 +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex) + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMultistageMmaComplexCore< + Shape_, WarpShape_, GemmShape<16, 8, 8>, + complex, layout::RowMajor, + complex, layout::ColumnMajor, + complex, LayoutC_, + arch::OpClassTensorOp, + Stages, + TransformA, TransformB, + Operator_, + CacheOpA, CacheOpB> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<16, 8, 8>; + using ElementA = complex; + using LayoutA = layout::RowMajor; + using ElementB = complex; + using LayoutB = layout::ColumnMajor; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + using Operator = Operator_; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped + static int const kAccessSizeInBits = 64; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicand64bCrosswise; + + using SmemLayoutB = layout::ColumnMajorTensorOpMultiplicand64bCrosswise; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaComplexTensorOp< + WarpShape, InstructionShape, + ElementA, SmemLayoutA, + ElementB, SmemLayoutB, + ElementC, LayoutC, + kTransformA, kTransformB, + Operator>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex floating-point +/// +/// A: row-major +/// B: row-major +/// Operator: arch::OpMultiplyAddComplex +/// Math Instruction: mma.sync.aligned.m16n8k8.f32.tf32.tf32.f32 +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex) + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMultistageMmaComplexCore< + Shape_, WarpShape_, GemmShape<16, 8, 8>, + complex, layout::RowMajor, + complex, layout::RowMajor, + complex, LayoutC_, + arch::OpClassTensorOp, + Stages, + TransformA, TransformB, + Operator_, + CacheOpA, CacheOpB> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<16, 8, 8>; + using ElementA = complex; + using LayoutA = layout::RowMajor; + using ElementB = complex; + using LayoutB = layout::RowMajor; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + using Operator = Operator_; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of a threadblock-scoped + static int const kAccessSizeInBits = 64; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::RowMajorTensorOpMultiplicand64bCrosswise; + + using SmemLayoutB = layout::RowMajorTensorOpMultiplicandCongruous64b; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearWarpRakedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 1, + IteratorThreadMapA>; + + /// ThreadMap of iterator B + using IteratorThreadMapB = transform::PitchLinearWarpStripedThreadMap< + layout::PitchLinearShape, kThreads, + layout::PitchLinearShape<16, 2>, + kAccessSizeInBits / sizeof_bits::value>; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 0, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level tensor op + using MmaTensorOp = typename mctlass::gemm::warp::DefaultMmaComplexTensorOp< + WarpShape, InstructionShape, + ElementA, SmemLayoutA, + ElementB, SmemLayoutB, + ElementC, LayoutC, + kTransformA, kTransformB, + Operator>::Type; + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy, + MatrixShape<0, 0>, WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex SIMT operation +/// +/// A: column-major +/// B: column-major +/// Operator: arch::OpMultiplyAddComplex or arch::OpMultiplyGaussianComplex +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + typename RealA, + typename RealB, + typename RealC, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMultistageMmaComplexCore< + Shape_, WarpShape_, GemmShape<1, 1, 1>, + complex, layout::ColumnMajor, + complex, layout::ColumnMajor, + complex, LayoutC_, + arch::OpClassSimt, + Stages, + TransformA, TransformB, + Operator_, + CacheOpA, CacheOpB> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 1>; + using ElementA = complex; + using LayoutA = layout::ColumnMajor; + using ElementB = complex; + using LayoutB = layout::ColumnMajor; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + using Operator = Operator_; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of access + static int const kAccessSizeInBits = sizeof_bits::value; + + /// No vectorized accesses + static int const kElementsPerAccess = 1; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + IteratorThreadMapA>; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Transpose the ThreadMap of iterator B + using SmemThreadMapB = transform::TransposePitchLinearThreadMapSimt; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + SmemThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = 4; + static const int WarpNumThreadsN = 8; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy, /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + 1, /// 1 partition along K dimension + kTransformA, /// Transform for A + kTransformB /// Transform for B + >; /// Used for partial specialization + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape<0, 0>, + MatrixShape<0, Shape::kK / 32>, + WarpCount::kK>; +}; + +/// Partial specialization for complex SIMT operation +/// +/// A: column-major +/// B: row-major +/// Operator: arch::OpMultiplyAddComplex or arch::OpMultiplyGaussianComplex +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + typename RealA, + typename RealB, + typename RealC, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMultistageMmaComplexCore< + Shape_, WarpShape_, GemmShape<1, 1, 1>, + complex, layout::ColumnMajor, + complex, layout::RowMajor, + complex, LayoutC_, + arch::OpClassSimt, + Stages, + TransformA, TransformB, + Operator_, + CacheOpA, CacheOpB> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 1>; + using ElementA = complex; + using LayoutA = layout::ColumnMajor; + using ElementB = complex; + using LayoutB = layout::RowMajor; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + using Operator = Operator_; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of access + static int const kAccessSizeInBits = sizeof_bits::value; + + /// No vectorized accesses + static int const kElementsPerAccess = 1; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + IteratorThreadMapA>; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = 4; + static const int WarpNumThreadsN = 8; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy, /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + 1, /// 1 partition along K dimension + kTransformA, /// Transform for A + kTransformB /// Transform for B + >; /// Used for partial specialization + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape<0, 0>, + MatrixShape<0, 0>, // or Shape::kK / 32 + WarpCount::kK>; +}; + +/// Partial specialization for complex SIMT operation +/// +/// A: row-major +/// B: column-major +/// Operator: arch::OpMultiplyAddComplex or arch::OpMultiplyGaussianComplex +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + typename RealA, + typename RealB, + typename RealC, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMultistageMmaComplexCore< + Shape_, WarpShape_, GemmShape<1, 1, 1>, + complex, layout::RowMajor, + complex, layout::ColumnMajor, + complex, LayoutC_, + arch::OpClassSimt, + Stages, + TransformA, TransformB, + Operator_, + CacheOpA, CacheOpB> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 1>; + using ElementA = complex; + using LayoutA = layout::RowMajor; + using ElementB = complex; + using LayoutB = layout::ColumnMajor; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + using Operator = Operator_; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of access + static int const kAccessSizeInBits = sizeof_bits::value; + + /// No vectorized accesses + static int const kElementsPerAccess = 1; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapA = transform::TransposePitchLinearThreadMapSimt; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + SmemThreadMapA>; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Transpose the ThreadMap of iterator B + using SmemThreadMapB = transform::TransposePitchLinearThreadMapSimt; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + SmemThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = 4; + static const int WarpNumThreadsN = 8; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy, /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + 1, /// 1 partition along K dimension + kTransformA, /// Transform for A + kTransformB /// Transform for B + >; /// Used for partial specialization + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape, + MatrixShape<0, Shape::kK / 32>, + WarpCount::kK>; +}; + +/// Partial specialization for complex SIMT operation +/// +/// A: row-major +/// B: row-major +/// Operator: arch::OpMultiplyAddComplex or arch::OpMultiplyGaussianComplex +/// +/// This uses the default warp-level operator given tile sizes +template < + /// Shape of threadblock-scoped matrix multiply operator (concept: + /// GemmShape) + typename Shape_, + /// Shape of warp-level matrix multiply operator (concept: GemmShape) + typename WarpShape_, + typename RealA, + typename RealB, + typename RealC, + /// Layout of accumulator + typename LayoutC_, + /// Number of stages + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator_, + /// Cache operation of operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Cache operation of operand B + mctlass::arch::CacheOperation::Kind CacheOpB> +struct DefaultMultistageMmaComplexCore< + Shape_, WarpShape_, GemmShape<1, 1, 1>, + complex, layout::RowMajor, + complex, layout::RowMajor, + complex, LayoutC_, + arch::OpClassSimt, + Stages, + TransformA, TransformB, + Operator_, + CacheOpA, CacheOpB> { + + using Shape = Shape_; + using WarpShape = WarpShape_; + using InstructionShape = GemmShape<1, 1, 1>; + using ElementA = complex; + using LayoutA = layout::RowMajor; + using ElementB = complex; + using LayoutB = layout::RowMajor; + using ElementC = complex; + using LayoutC = LayoutC_; + static int const kStages = Stages; + static ComplexTransform const kTransformA = TransformA; + static ComplexTransform const kTransformB = TransformB; + using Operator = Operator_; + static mctlass::arch::CacheOperation::Kind const kCacheOpA = mctlass::arch::CacheOperation::Always; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = mctlass::arch::CacheOperation::Always; + + /// Number of warps present + using WarpCount = GemmShape; + + // Divisility requirements + static_assert( + !(Shape::kM % WarpShape::kM) && !(Shape::kN % WarpShape::kN), + "Threadblock-scoped GEMM should be divisible by warp-scoped GEMM size."); + + static_assert(WarpCount::kCount > 1, + "This specialization requires at least two warps."); + + /// Number of threads per warp + static int const kWarpSize = warp::WarpSize::value; + + /// Number of threads total + static int const kThreads = WarpCount::kCount * kWarpSize; + + /// Size of access + static int const kAccessSizeInBits = sizeof_bits::value; + + /// No vectorized accesses + static int const kElementsPerAccess = 1; + + // + // Shared memory layouts + // + + using SmemLayoutA = layout::ColumnMajor; + + using SmemLayoutB = layout::RowMajor; + + // + // Iterators to write to shared memory + // + + /// ThreadMap of iterator A + using IteratorThreadMapA = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Transpose the ThreadMap of iterator A + using SmemThreadMapA = transform::TransposePitchLinearThreadMapSimt; + + /// Shared memory iterator to A operand + using SmemIteratorA = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementA, SmemLayoutA, 0, + SmemThreadMapA>; + + /// Policy of iterator B + using IteratorThreadMapB = transform::PitchLinearStripminedThreadMap< + layout::PitchLinearShape, + kThreads, + kElementsPerAccess + >; + + /// Shared memory iterator to B operand + using SmemIteratorB = transform::threadblock::RegularTileAccessIterator< + MatrixShape, ElementB, SmemLayoutB, 1, + IteratorThreadMapB>; + + // + // Warp-level matrix multiply operator + // + + // Define the warp-level op + static const int WarpNumThreadsM = 4; + static const int WarpNumThreadsN = 8; + static_assert(!(WarpShape::kM % WarpNumThreadsM) && !(WarpShape::kN % WarpNumThreadsN), + "WarpShape must be divisible by ThreadTile shape."); + static const int ThreadTileM = WarpShape::kM / WarpNumThreadsM; + static const int ThreadTileN = WarpShape::kN / WarpNumThreadsN; + static const int LaneLayout = ThreadTileM > 4 && ThreadTileN > 4 ? 2 : 1; + static const int numElementsA = 128 / sizeof_bits::value; + static const int numElementsB = 128 / sizeof_bits::value; + static const int LaneM = mctlass::const_min(numElementsA, ThreadTileM); + static const int LaneN = mctlass::const_min(numElementsB, ThreadTileN); + // these should have max of thread tile also + using LaneMmaShape = mctlass::gemm::GemmShape< + LaneM, + LaneN, + 1>; + using Policy = mctlass::gemm::warp::MmaSimtPolicy< + mctlass::MatrixShape, // WarpShape + mctlass::layout::RowMajorInterleaved, // LaneLayout + LaneMmaShape + >; + + using MmaWarpSimt = mctlass::gemm::warp::MmaSimt< + WarpShape, /// Size of the Gemm problem - concept: gemm::GemmShape<> 128, 128, 8 + ElementA, /// Data type of A elements + SmemLayoutA, /// Layout of A matrix (concept: MatrixLayout) + ElementB, /// Data type of B elements + SmemLayoutB, /// Layout of B matrix (concept: MatrixLayout) + ElementC, /// Element type of C matrix + LayoutC, /// Layout of C matrix (concept: MatrixLayout) + Policy, /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + 1, /// 1 partition along K dimension + kTransformA, /// Transform for A + kTransformB /// Transform for B + >; /// Used for partial specialization + + /// Policy used to define MmaPipelined + using MmaPolicy = MmaPolicy< + MmaWarpSimt, + MatrixShape, + MatrixShape<0, 0>, // or Shape::kK / 32 + WarpCount::kK>; +}; + +//////////////////////////////////////////////////////////////////////////////// + + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_trmm_complex.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_trmm_complex.h new file mode 100644 index 0000000..ab29722 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_multistage_trmm_complex.h @@ -0,0 +1,556 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Template for a multistage GEMM kernel. Does not compute batching or support split-K. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/arch/arch.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/numeric_types.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator_triangular_matrix.h" +#include "mctlass/gemm/threadblock/mma_blas3_multistage.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Diag Type for the triangular matrix + DiagType kDiagType, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex transformation on operand B + ComplexTransform TransformB = ComplexTransform::kNone, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator = arch::OpMultiplyAddComplex, + /// Blas3 computation mode + BlasMode BlasMode_ = BlasMode::kTriangular, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false> +struct DefaultMultistageTrmmComplex; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Diag Type for the triangular matrix + DiagType kDiagType, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator> +struct DefaultMultistageTrmmComplex { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplexCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, TransformA, TransformB, Operator>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, + kSideMode, kFillMode, kDiagType, + AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, + kSideMode, FillMode::kFull, DiagType::kInvalid, + AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages, SharedMemoryClearOption::kZfill>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output and right-side mode +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Diag Type for the triangular matrix + DiagType kDiagType, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator> +struct DefaultMultistageTrmmComplex { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplexCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, TransformA, TransformB, Operator>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, + SideMode::kRight, FillMode::kFull, DiagType::kInvalid, + AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, + SideMode::kRight, kFillMode, kDiagType, + AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages, SharedMemoryClearOption::kZfill>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output with unit diagonal +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator> +struct DefaultMultistageTrmmComplex { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplexCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, TransformA, TransformB, Operator>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, + kSideMode, kFillMode, DiagType::kUnit, + AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, + kSideMode, FillMode::kFull, DiagType::kInvalid, + AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaBlas3Multistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages, SharedMemoryClearOption::kZfill>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output and right-side mode, unit diagonal +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator> +struct DefaultMultistageTrmmComplex { + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplexCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, TransformA, TransformB, Operator>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, + SideMode::kRight, FillMode::kFull, DiagType::kInvalid, + AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, + SideMode::kRight, kFillMode, DiagType::kUnit, + AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaBlas3Multistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages, SharedMemoryClearOption::kZfill>; +}; + + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output (for TRMM where diagonal imag part is ignored - used by HEMM) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator> +struct DefaultMultistageTrmmComplex { + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplexCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, TransformA, TransformB, Operator>; + + // Define iterators over tiles from the A operand + // PredicatedTileAccessIteratorTriangularMatrix only tracks diagonal elements, + // when DiagType is kUnit + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, + kSideMode, kFillMode, DiagType::kUnit, + AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, + kSideMode, FillMode::kFull, DiagType::kInvalid, + AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaBlas3Multistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages, SharedMemoryClearOption::kZfill, + BlasMode::kHermitian>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output and right-side mode (for TRMM where diagonal imag part is ignored - used by HEMM) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename OperatorClass, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Complex transformation on operand A + ComplexTransform TransformA, + /// Complex transformation on operand B + ComplexTransform TransformB, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator> +struct DefaultMultistageTrmmComplex { + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMultistageMmaComplexCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, OperatorClass, + Stages, TransformA, TransformB, Operator>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, + SideMode::kRight, FillMode::kFull, DiagType::kInvalid, + AccessTypeA>; + + // Define iterators over tiles from the B operand + // PredicatedTileAccessIteratorTriangularMatrix only tracks diagonal elements, + // when DiagType is kUnit + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, + SideMode::kRight, kFillMode, DiagType::kUnit, + AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaBlas3Multistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages, SharedMemoryClearOption::kZfill, + BlasMode::kHermitian>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_sparse_mma.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_sparse_mma.h new file mode 100644 index 0000000..ec0985c --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_sparse_mma.h @@ -0,0 +1,196 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/arch.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator_2dthreadtile.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" +#include "mctlass/gemm/threadblock/default_mma_core_sparse_sm80.h" +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/gemm/threadblock/default_mma_core_wmma.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false + > +struct DefaultSparseMma; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output (OperatorClass TensorOp) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator + > +struct DefaultSparseMma { + static mctlass::arch::CacheOperation::Kind const CacheOpA = + ((sizeof_bits::value * kAlignmentA) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * kAlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultSparseMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, Operator, false, CacheOpA, CacheOpB>; + + static int const kSparse = MmaCore::kSparse; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, AccessTypeB>; + + // Define iterators over tiles from the E operand + using ElementE = typename MmaCore::ElementE; + using LayoutE = typename MmaCore::GmemLayoutE; + using ThreadMapE = typename MmaCore::IteratorThreadMapE; + using AccessTypeE = + mctlass::Array::value>; + using IteratorE = + mctlass::transform::threadblock::PredicatedTileAccessIterator< + mctlass::MatrixShape, + ElementE, LayoutE, 1, ThreadMapE, AccessTypeE>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::SparseMmaMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + IteratorE, typename MmaCore::SmemIteratorE, MmaCore::kCacheOpE, + typename MmaCore::MmaPolicy, Stages>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/default_trmm.h b/csrc/mctlass/include/mctlass/gemm/threadblock/default_trmm.h new file mode 100644 index 0000000..30899a3 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/default_trmm.h @@ -0,0 +1,445 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +// +/*! \file + \brief Template for a pipelined GEMM kernel. Does not compute batching or support split-K. +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/arch/arch.h" +#include "mctlass/arch/wmma.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator_triangular_matrix.h" +#include "mctlass/gemm/threadblock/mma_blas3_multistage.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator.h" +#include "mctlass/transform/threadblock/predicated_tile_iterator_2dthreadtile.h" +#include "mctlass/gemm/threadblock/default_mma_core_simt.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm70.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm75.h" +#include "mctlass/gemm/threadblock/default_mma_core_sm80.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) +#include "mctlass/gemm/threadblock/default_mma_core_wmma.h" +#endif //MCTLASS_ARCH_WMMA_ENABLED + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Element type for A matrix operand + typename ElementA_, + /// Layout type for A matrix operand + typename LayoutA_, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB_, + /// Layout type for B matrix operand + typename LayoutB_, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Diag Type for the triangular matrix + DiagType kDiagType, + /// Element type for internal accumulation + typename ElementAccumulator_, + /// Layout type for C and D matrix operands + typename LayoutC_, + /// Operator class tag + typename OperatorClass_, + /// Tag indicating architecture to tune for + typename ArchTag_, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape_, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape_, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape_, + /// Number of stages used in the pipelined mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false + > +struct DefaultTrmm; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output (OperatorClass TensorOp) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Diag Type for the triangular matrix + DiagType kDiagType, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator + > +struct DefaultTrmm { + + static mctlass::arch::CacheOperation::Kind const CacheOpA = + ((sizeof_bits::value * kAlignmentA) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * kAlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, Operator, false, CacheOpA, CacheOpB>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, kSideMode, kFillMode, kDiagType, AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, kSideMode, FillMode::kFull, DiagType::kInvalid, AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages, SharedMemoryClearOption::kZfill>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output, right side mode (OperatorClass TensorOp) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Diag Type for the triangular matrix + DiagType kDiagType, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator + > +struct DefaultTrmm { + + static mctlass::arch::CacheOperation::Kind const CacheOpA = + ((sizeof_bits::value * kAlignmentA) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * kAlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, Operator, false, CacheOpA, CacheOpB>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, SideMode::kRight, FillMode::kFull, DiagType::kInvalid, AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, SideMode::kRight, kFillMode, kDiagType, AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaMultistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages, SharedMemoryClearOption::kZfill>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output with unit diagonal (OperatorClass TensorOp) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Side Mode for the kernel + SideMode kSideMode, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator + > +struct DefaultTrmm { + + static mctlass::arch::CacheOperation::Kind const CacheOpA = + ((sizeof_bits::value * kAlignmentA) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * kAlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, Operator, false, CacheOpA, CacheOpB>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, kSideMode, kFillMode, DiagType::kUnit, AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, kSideMode, FillMode::kFull, DiagType::kInvalid, AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaBlas3Multistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages, SharedMemoryClearOption::kZfill>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization for row-major output, right side mode, unit diagonal (OperatorClass TensorOp) +template < + /// Element type for A matrix operand + typename ElementA, + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Element type for B matrix operand + typename ElementB, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Fill Mode for the triangular matrix + FillMode kFillMode, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Number of stages used in the multistage mainloop + int Stages, + /// Operation perfomed by GEMM + typename Operator + > +struct DefaultTrmm { + + static mctlass::arch::CacheOperation::Kind const CacheOpA = + ((sizeof_bits::value * kAlignmentA) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + static mctlass::arch::CacheOperation::Kind const CacheOpB = + ((sizeof_bits::value * kAlignmentB) == 128) + ? mctlass::arch::CacheOperation::Global + : mctlass::arch::CacheOperation::Always; + + // Define the MmaCore components + using MmaCore = typename mctlass::gemm::threadblock::DefaultMmaCore< + ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA, + ElementB, LayoutB, ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, + Stages, Operator, false, CacheOpA, CacheOpB>; + + // Define iterators over tiles from the A operand + using ThreadMapA = typename MmaCore::IteratorThreadMapA; + using AccessTypeA = mctlass::Array; + + using IteratorA = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementA, LayoutA, 1, ThreadMapA, SideMode::kRight, FillMode::kFull, DiagType::kInvalid, AccessTypeA>; + + // Define iterators over tiles from the B operand + using ThreadMapB = typename MmaCore::IteratorThreadMapB; + using AccessTypeB = mctlass::Array; + + using IteratorB = + mctlass::transform::threadblock::PredicatedTileAccessIteratorTriangularMatrix< + mctlass::MatrixShape, + ElementB, LayoutB, 0, ThreadMapB, SideMode::kRight, kFillMode, DiagType::kUnit, AccessTypeB>; + + // Define the threadblock-scoped multistage matrix multiply + using ThreadblockMma = mctlass::gemm::threadblock::MmaBlas3Multistage< + typename MmaCore::Shape, IteratorA, typename MmaCore::SmemIteratorA, + MmaCore::kCacheOpA, IteratorB, typename MmaCore::SmemIteratorB, + MmaCore::kCacheOpB, ElementAccumulator, layout::RowMajor, + typename MmaCore::MmaPolicy, Stages, SharedMemoryClearOption::kZfill>; +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/ell_mma_multistage.h b/csrc/mctlass/include/mctlass/gemm/threadblock/ell_mma_multistage.h new file mode 100644 index 0000000..d45e2bf --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/ell_mma_multistage.h @@ -0,0 +1,648 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a multistage threadblock-scoped Blocked-Ell MMA. +*/ + +#pragma once + + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" + +#include "mctlass/gemm/threadblock/mma_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class EllMmaMultistage : + public MmaBase { +public: + ///< Base class + using Base = MmaBase; + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + ///< Policy describing tuning details + using Policy = Policy_; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + using EllIterator = typename mctlass::transform::threadblock::ell::Iterator; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + /// Internal structure exposed for introspection. + struct Detail { + + static_assert(Base::kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + }; + + private: + + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + EllMmaMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + template + MCTLASS_DEVICE + void copy_tiles_and_advance(IteratorA &iterator_A, IteratorB &iterator_B, EllIterator &ell_iter, + int group_start_A = 0, int group_start_B = 0) { + iterator_A.set_iteration_index(group_start_A * + IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_A.get(); + bool is_valid = iterator_A.valid(); + + if (!is_A_sparse){ + if (is_offset_constant){ + auto ell_offset = ell_iter.get_offset_fast(); + is_valid = is_valid && (ell_offset >= 0); + gmem_ptr += ell_offset * sizeof(IteratorA::Element) / kSrcBytes; + } else { + int k_offset = iterator_A.get_k(); + auto ell_offset = ell_iter.get_offset(k_offset); + is_valid = is_valid && (ell_offset >= 0); + gmem_ptr += (ell_offset * sizeof(IteratorA::Element)) / kSrcBytes; + } + } + + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, is_valid); + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + } + + iterator_B.set_iteration_index(group_start_B * + IteratorB::kAccessesPerVector); + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_B.get(); + bool is_valid = iterator_B.valid(); + + if (is_A_sparse){ + if (is_offset_constant){ + auto ell_offset = ell_iter.get_offset_fast(); + is_valid = is_valid && (ell_offset >= 0); + gmem_ptr += ell_offset * sizeof(IteratorB::Element) / kSrcBytes; + } else { + int k_offset = iterator_B.get_k(); + auto ell_offset = ell_iter.get_offset(k_offset); + is_valid = is_valid && (ell_offset >= 0); + gmem_ptr += ( ell_offset * sizeof(IteratorB::Element)) / kSrcBytes; + } + } + + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, is_valid); + + ++iterator_B; + } + ++this->smem_iterator_B_; + } + } + } + + + /// Perform a threadblock-scoped matrix multiply-accumulate + template + MCTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< initial value of accumulator + FragmentC const &src_accum, + EllIterator &ell_iterator + ) { + // + // Prologue + // + + // Issue several complete stages + MCTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; + ++stage, --gemm_k_iterations) { + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + auto gmem_ptr = iterator_A.get(); + bool is_valid = iterator_A.valid(); + + if (!is_A_sparse){ + if (is_offset_constant){ + auto ell_offset = ell_iterator.get_offset_fast(); + is_valid = is_valid && (ell_offset >= 0); + gmem_ptr += ell_offset * sizeof(IteratorA::Element) / kSrcBytes; + } else { + int k_offset = iterator_A.get_k(); + auto ell_offset = ell_iterator.get_offset(k_offset); + is_valid = is_valid && (ell_offset >= 0); + gmem_ptr += (ell_offset * sizeof(IteratorA::Element)) / kSrcBytes; + } + } + + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, is_valid); + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + + iterator_B.set_iteration_index(0); + this->smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + auto gmem_ptr = iterator_B.get(); + bool is_valid = iterator_B.valid(); + + if (is_A_sparse){ + if (is_offset_constant){ + auto ell_offset = ell_iterator.get_offset_fast(); + is_valid = is_valid && (ell_offset >= 0); + gmem_ptr += ell_offset * sizeof(IteratorB::Element) / kSrcBytes; + } else { + int k_offset = iterator_B.get_k(); + auto ell_offset = ell_iterator.get_offset(k_offset); + is_valid = is_valid && (ell_offset >= 0); + gmem_ptr += ( ell_offset * sizeof(IteratorB::Element)) / kSrcBytes; + } + } + + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, is_valid); + + ++iterator_B; + } + + ++this->smem_iterator_B_; + } + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + ++ell_iterator; + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Defines the boundary of a stage of cp.async. + mctlass::arch::cp_async_fence(); + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // Waits until kStages-2 stages have committed. + mctlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[2]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpTransformedFragmentA warp_transformed_frag_A[2]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + if (is_A_sparse){ + iterator_A.ell_add_mask(ell_iterator.get_blocksize()); + } + else { + iterator_B.ell_add_mask(ell_iterator.get_blocksize()); + } + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0], + warp_loaded_frag_A[0], warp_loaded_frag_B[0]); + + // tf32x3 kernels use staging accumulation. warp_mma uses a temporary + // accumulator and this temporary accumulator is added to the final + // accumulator once in every mainloop iteration. + plus plus_accum; + + FragmentC tmp_accum; + + if (platform::is_same::value + || platform::is_same::value) { + + tmp_accum.clear(); + } + + // + // Mainloop + // + + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) + warp_mma.transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % 2], + warp_loaded_frag_B[warp_mma_k % 2]); + + if (platform::is_same::value + || platform::is_same::value) { + + warp_mma( + tmp_accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + tmp_accum + ); + + if (warp_mma_k == 0) { + accum = plus_accum(accum, tmp_accum); + tmp_accum.clear(); + } + } else { + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum + ); + } + + // Issue global->shared copies for the this stage + if (warp_mma_k < Base::kWarpGemmIterations - 1) { + int group_start_iteration_A, group_start_iteration_B; + + group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA; + group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance( + iterator_A, iterator_B, ell_iterator, group_start_iteration_A, + group_start_iteration_B); + } + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + int group_start_iteration_A, group_start_iteration_B; + group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance( + iterator_A, iterator_B, ell_iterator, group_start_iteration_A, + group_start_iteration_B); + + // Inserts a memory fence between stages of cp.async instructions. + mctlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages have committed. + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + ++ell_iterator; + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + } + + // Do any conversions feeding the first stage at the end of the loop so + // we can start right away on mma instructions + if (warp_mma_k + 1 == Base::kWarpGemmIterations) + warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + } + + } + + if (platform::is_same::value + || platform::is_same::value) { + accum = plus_accum(accum, tmp_accum); + } + + + // Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/ell_mma_pipelined.h b/csrc/mctlass/include/mctlass/gemm/threadblock/ell_mma_pipelined.h new file mode 100644 index 0000000..8f43787 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/ell_mma_pipelined.h @@ -0,0 +1,376 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped Blocked-Ell MMA. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/aligned_buffer.h" +#include "mctlass/numeric_conversion.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/mma_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Transformation applied to A operand + typename TransformA_ = NumericArrayConverter< + typename SmemIteratorA_::Element, + typename IteratorA_::Element, + IteratorA_::Fragment::kElements>, + /// + /// Transformation applied to B operand + typename TransformB_ = NumericArrayConverter< + typename SmemIteratorB_::Element, + typename IteratorB_::Element, + IteratorB_::Fragment::kElements>, + /// Used for partial specialization + typename Enable = bool +> +class EllMmaPipelined : public MmaBase { +public: + + ///< Base class + using Base = MmaBase; + + using Shape = Shape_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using IteratorA = IteratorA_; ///< Iterates over tiles of A operand in global memory + using IteratorB = IteratorB_; ///< Iterates over tiles of B operand in global memory + using ElementC = ElementC_; ///< Data type of accumulator matrix + using LayoutC = LayoutC_; ///< Layout of accumulator matrix + using Policy = Policy_; ///< Policy describing tuning details + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + using TransformA = TransformA_; + using TransformB = TransformB_; + + // + // Dependent types + // + + /// Fragment of operand A loaded from global memory + using FragmentA = typename IteratorA::Fragment; + + /// Fragment of operand B loaded from global memory + using FragmentB = typename IteratorB::Fragment; + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Obtain the arch tag from the warp-level operator + using ArchTag = typename Policy::Operator::ArchTag; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + // staticaly assert kStages for EllMmaPipelined is two (Double-buffered pipeline) + static_assert((Base::kStages==2), "EllMmaPipelined requires kStages set to value 2"); + +private: + + using WarpFragmentA = typename Operator::FragmentA; + using WarpFragmentB = typename Operator::FragmentB; + +protected: + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + using EllIterator = typename mctlass::transform::threadblock::ell::Iterator; + +public: + /// Construct from tensor references + MCTLASS_DEVICE + EllMmaPipelined( + typename Base::SharedStorage &shared_storage, ///< Shared storage needed for internal use by threadblock-scoped GEMM + int thread_idx, ///< ID within the threadblock + int warp_idx, ///< ID of warp + int lane_idx ///< ID of each thread within a warp + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) { + + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + template + MCTLASS_DEVICE + void operator()( + int gemm_k_iterations, ///< number of iterations of the mainloop + FragmentC &accum, ///< destination accumulator tile + IteratorA iterator_A, ///< iterator over A operand in global memory + IteratorB iterator_B, ///< iterator over B operand in global memory + FragmentC const &src_accum, ///< source accumulator tile + EllIterator &ell_iterator, + TransformA transform_A = TransformA(), ///< transformation applied to A fragment + TransformB transform_B = TransformB()) { ///< transformation applied to B fragment + + // + // Prologue + // + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + FragmentA tb_frag_A; + FragmentB tb_frag_B; + + tb_frag_A.clear(); + tb_frag_B.clear(); + + // load sparse matrix + if (is_A_sparse){ + iterator_A.load(tb_frag_A); + } else { + iterator_B.load(tb_frag_B); + } + + // load dense matrix + if (is_offset_constant){ + if (is_A_sparse){ + iterator_B.load_with_ell_index_fast(tb_frag_B, ell_iterator); + } else { + iterator_A.load_with_ell_index_fast(tb_frag_A, ell_iterator); + } + } else { + if (is_A_sparse){ + iterator_B.load_with_ell_index(tb_frag_B, ell_iterator); + } else { + iterator_A.load_with_ell_index(tb_frag_A, ell_iterator); + } + } + + ++iterator_A; + ++iterator_B; + ++ell_iterator; + + this->smem_iterator_A_.store(transform_A(tb_frag_A)); + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math instructions + WarpFragmentA warp_frag_A[2]; + WarpFragmentB warp_frag_B[2]; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + Operator warp_mma; + + int smem_write_stage_idx = 1; + + // Avoid reading out of bounds + iterator_A.clear_mask(gemm_k_iterations <= 1); + iterator_B.clear_mask(gemm_k_iterations <= 1); + + if (is_A_sparse){ + iterator_A.ell_add_mask(ell_iterator.get_blocksize()); + } + else { + iterator_B.ell_add_mask(ell_iterator.get_blocksize()); + } + + // Issue loads during the first warp-level matrix multiply-add *AFTER* issuing + // shared memory loads (which have the tightest latency requirement). + + // + // Mainloop + // + + // Note: The main loop does not support Base::kWarpGemmIterations == 2. + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > 0; --gemm_k_iterations) { + // + // Loop over GEMM K dimension + // + + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if this is the last group + // as the case may be. + + if (warp_mma_k == Base::kWarpGemmIterations - 1) { + + // Write fragments to shared memory + this->smem_iterator_A_.store(transform_A(tb_frag_A)); + + this->smem_iterator_B_.store(transform_B(tb_frag_B)); + + __syncthreads(); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + // Add negative offsets to return iterators to the 'start' of the circular buffer in shared memory + if (smem_write_stage_idx == 1) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + } + else { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations, + 0}); + } + + smem_write_stage_idx ^= 1; + } + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k == 0) { + // load sparse matrix + if (is_A_sparse){ + iterator_A.load(tb_frag_A); + } else { + iterator_B.load(tb_frag_B); + } + + // load dense matrix + if (is_offset_constant){ + if (is_A_sparse){ + iterator_B.load_with_ell_index_fast(tb_frag_B, ell_iterator); + } else { + iterator_A.load_with_ell_index_fast(tb_frag_A, ell_iterator); + } + } else { + if (is_A_sparse){ + iterator_B.load_with_ell_index(tb_frag_B, ell_iterator); + } else { + iterator_A.load_with_ell_index(tb_frag_A, ell_iterator); + } + } + + ++iterator_A; + ++iterator_B; + ++ell_iterator; + + // Avoid reading out of bounds if this was the last loop iteration + iterator_A.clear_mask(gemm_k_iterations <= 2); + iterator_B.clear_mask(gemm_k_iterations <= 2); + } + + warp_mma(accum, warp_frag_A[warp_mma_k % 2], + warp_frag_B[warp_mma_k % 2], accum); + } + } + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/gemv.h b/csrc/mctlass/include/mctlass/gemm/threadblock/gemv.h new file mode 100644 index 0000000..e42d667 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/gemv.h @@ -0,0 +1,147 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Template for a threadblock-scoped GEMV kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/gemm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix-vector product using SIMT math instructions. +template < + class Core_ //< GemvCore +> +class Gemv { +public: + using Shape = typename Core_::Shape; + + /// The MMA operator that computes GEMV + using Operator = typename Core_::Operator; + + /// Iterates over A in global memory + using IteratorA = typename Core_::IteratorA; + + /// Iterates over B in global memory + using IteratorB = typename Core_::IteratorB; + + /// Fragment of operand C loaded from global memory + using IteratorC = typename Core_::IteratorC; + + /// Fragment of operand A loaded from global memory + using FragmentA = typename IteratorA::Fragment; + + /// Fragment of operand B loaded from global memory + using FragmentB = typename IteratorB::Fragment; + + /// Fragment of operand accumulator loaded/stored to global memory + using FragmentC = typename Operator::FragmentC; + + /// Shape of the per-thread GEMV operation + using ThreadShape = typename Core_::ThreadShape; + +public: + MCTLASS_DEVICE + Gemv() { } + + MCTLASS_DEVICE + void operator()( + GemmCoord const &problem_size, ///< problem size of batched GEMV + FragmentC &accum, ///< destination accumulator tile + IteratorA iterator_A, ///< iterator over A operand in global memory + IteratorB iterator_B, ///< iterator over B operand in global memory + FragmentC const &src_accum) { ///< source accumualtor tile + + // + // Prologue + // + + FragmentA frag_A; + FragmentB frag_B; + frag_A.clear(); + frag_B.clear(); + + iterator_A.load(frag_A); + iterator_B.load(frag_B); + ++iterator_A; + ++iterator_B; + + // + // Mainloop + // + Operator thread_mma; + int gemm_k = problem_size.k(); + + if (gemm_k < Shape::kK) + { + iterator_A.clear_mask(); + iterator_B.clear_mask(); + } + + // iterate over K to accumulate result + MCTLASS_GEMM_LOOP + for (; gemm_k > 0; gemm_k -= Shape::kK) { + thread_mma(accum, frag_A, frag_B, accum); + + iterator_A.load(frag_A); + iterator_B.load(frag_B); + ++iterator_A; + ++iterator_B; + + if (gemm_k < Shape::kK) + { + iterator_A.clear_mask(); + iterator_B.clear_mask(); + } + } + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/index_remat.h b/csrc/mctlass/include/mctlass/gemm/threadblock/index_remat.h new file mode 100644 index 0000000..5002167 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/index_remat.h @@ -0,0 +1,107 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Helpers for rematerializing indices/dimensions in the thread hierarchy from special registers +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Helper to rematerialize block Idx. Reduces register liveness. +MCTLASS_DEVICE +int RematerializeThreadIdxX() { + return threadIdx.x; +} + +/// Helper to rematerialize block Idx. Reduces register liveness. +MCTLASS_DEVICE +int RematerializeThreadIdxY() { + return threadIdx.y; +} + +/// Helper to rematerialize block Idx. Reduces register liveness. +MCTLASS_DEVICE +int RematerializeThreadIdxZ() { + return threadIdx.z; +} + +/// Helper to rematerialize block Idx. Reduces register liveness. +MCTLASS_DEVICE +int RematerializeBlockIdxX() { + return blockIdx.x; +} + +/// Helper to rematerialize block Idx. Reduces register liveness. +MCTLASS_DEVICE +int RematerializeBlockIdxY() { + return blockIdx.y; +} + +/// Helper to rematerialize block Idx. Reduces register liveness. +MCTLASS_DEVICE +int RematerializeBlockIdxZ() { + return blockIdx.z; +} + +/// Helper to rematerialize block Dim. Reduces register liveness. +MCTLASS_DEVICE +int RematerializeBlockDimX() { + return blockDim.x; +} + +/// Helper to rematerialize block Dim. Reduces register liveness. +MCTLASS_DEVICE +int RematerializeBlockDimY() { + return blockDim.y; +} + +/// Helper to rematerialize block Dim. Reduces register liveness. +MCTLASS_DEVICE +int RematerializeBlockDimZ() { + return blockDim.z; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + + diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/mma_base.h b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_base.h new file mode 100644 index 0000000..2e21b40 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_base.h @@ -0,0 +1,236 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "mctlass/tensor_ref.h" +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Policy object describing MmaTensorOp +template < + /// Warp-level GEMM operator (concept: gemm::warp::Mma) + typename Operator_, + /// Padding used for A operand in shared memory (concept: MatrixShape) + typename SmemPaddingA_, + /// Padding used for B operand in shared memory (concept: MatrixShape) + typename SmemPaddingB_, + /// Number of partitions of K dimension of GEMM + int PartitionsK = 1> +struct MmaPolicy { + /// Warp-level GEMM operator (concept: gemm::warp::MmaTensorOp or gemm::warp::MmaSimt) + using Operator = Operator_; + + /// Padding used for A operand in shared memory + using SmemPaddingA = SmemPaddingA_; + + /// Padding used for B operand in shared memory + using SmemPaddingB = SmemPaddingB_; + + /// Number of partitions of K dimension + static int const kPartitionsK = PartitionsK; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class MmaBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Policy describing tuning details + using Policy = Policy_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = GemmShape; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = + (WarpGemm::kK / Operator::Policy::MmaShape::kK); + + /// Number of stages + static int const kStages = Stages; + + /// Tensor reference to the A operand + using TensorRefA = TensorRef; + + /// Tensor reference to the B operand + using TensorRefB = TensorRef; + + static_assert(kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + static_assert((kWarpGemmIterations % 2) == 0, + "Inner loop iteration must be an even number."); + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + class SharedStorage { + public: + // + // Type definitions + // + + /// Shape of the A matrix operand in shared memory + using ShapeA = MatrixShape; + + /// Shape of the B matrix operand in shared memory + using ShapeB = + MatrixShape; + + public: + // + // Data members + // + + /// Buffer for A operand + AlignedBuffer operand_A; + + /// Buffer for B operand + AlignedBuffer operand_B; + + public: + + // + // Methods + // + + /// Returns a layout object for the A matrix + MCTLASS_DEVICE + static typename Operator::LayoutA LayoutA() { + return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn}); + } + + /// Returns a layout object for the B matrix + MCTLASS_HOST_DEVICE + static typename Operator::LayoutB LayoutB() { + return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); + } + + /// Returns a TensorRef to the A operand + MCTLASS_HOST_DEVICE + TensorRefA operand_A_ref() { + return TensorRefA{operand_A.data(), LayoutA()}; + } + + /// Returns a TensorRef to the B operand + MCTLASS_HOST_DEVICE + TensorRefB operand_B_ref() { + return TensorRefB{operand_B.data(), LayoutB()}; + } + }; + + protected: + + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A operand from shared memory + typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaBase( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx), + warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) { + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/mma_blas3_multistage.h b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_blas3_multistage.h new file mode 100644 index 0000000..8bf6407 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_blas3_multistage.h @@ -0,0 +1,707 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. + Used by BLAS3 kernels that need to treat diagonal elements of a input iterator as a special case. + +*/ + +#pragma once + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" + +#include "mctlass/gemm/threadblock/mma_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kZfill, + /// Blas3 computation mode + BlasMode BlasMode_ = BlasMode::kTriangular, + /// Used for partial specialization + typename Enable = bool> +class MmaBlas3Multistage : + public MmaBase { +public: + ///< Base class + using Base = MmaBase; + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + ///< Policy describing tuning details + using Policy = Policy_; + ///< Blas Mode + static BlasMode const kBlasMode = BlasMode_; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + /// Internal structure exposed for introspection. + struct Detail { + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + }; + + private: + + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaBlas3Multistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + MCTLASS_DEVICE + void copy_tiles_and_advance(IteratorA &iterator_A, IteratorB &iterator_B, + int group_start_A = 0, int group_start_B = 0) { + iterator_A.set_iteration_index(group_start_A * + IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_A.get(); + bool isvalid = iterator_A.valid(); + + if (isvalid && iterator_A.getOnDiag()) { + // Elements that are on diagonal + if (kBlasMode == BlasMode::kHermitian && mctlass::is_complex::value) { + /* Copy real part from gmem, write zero for imag part in smem */ + /* The following logic to determine kSizeRealBytes is so that compiler doesn't complain when + * compiling for not complex datatype and using half the size for cp_async_zfill */ + int const kSizeRealBytes = (platform::is_same>::value) ? 8 : 4; + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, true); + mctlass::arch::cp_async_diag( + reinterpret_cast (dst_ptr + v) + kSizeRealBytes); + } else { + /* Write one (1) directly to smem*/ + mctlass::arch::cp_async_diag(dst_ptr + v); + } + } else { + // Elements that are not of diagonal + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, isvalid); + } + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + } + + iterator_B.set_iteration_index(group_start_B * + IteratorB::kAccessesPerVector); + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_B.get(); + bool isvalid = iterator_B.valid(); + + if (isvalid && iterator_B.getOnDiag()) { + // Elements that are on diagonal + if (kBlasMode == BlasMode::kHermitian && mctlass::is_complex::value) { + /* Copy real part from gmem, write zero for imag part in smem */ + int const kSizeRealBytes = (platform::is_same>::value) ? 8 : 4; + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, true); + mctlass::arch::cp_async_diag( + reinterpret_cast (dst_ptr + v) + kSizeRealBytes); + } else { + /* Write one (1) directly to smem*/ + mctlass::arch::cp_async_diag(dst_ptr + v); + } + } else { + // Elements that are not of diagonal + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, isvalid); + } + + ++iterator_B; + } + ++this->smem_iterator_B_; + } + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< initial value of accumulator + FragmentC const &src_accum) { + + // + // Prologue + // + + // Issue several complete stages + MCTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; + ++stage, --gemm_k_iterations) { + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + auto gmem_ptr = iterator_A.get(); + bool isvalid = iterator_A.valid(); + + if (isvalid && iterator_A.getOnDiag()) { + // Elements that are on diagonal + if (kBlasMode == BlasMode::kHermitian && mctlass::is_complex::value) { + /* Copy real part from gmem, write zero for imag part in smem */ + int const kSizeRealBytes = (platform::is_same>::value) ? 8 : 4; + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, true); + mctlass::arch::cp_async_diag( + reinterpret_cast (dst_ptr + v) + kSizeRealBytes); + } else { + /* Write one (1) directly to smem*/ + mctlass::arch::cp_async_diag(dst_ptr + v); + } + } else { + // Elements that are not of diagonal + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, isvalid); + } + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + + iterator_B.set_iteration_index(0); + this->smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + auto gmem_ptr = iterator_B.get(); + bool isvalid = iterator_B.valid(); + + if (isvalid && iterator_B.getOnDiag()) { + // Elements that are on diagonal + if (kBlasMode == BlasMode::kHermitian && mctlass::is_complex::value) { + /* Copy real part from gmem, write zero for imag part in smem */ + int const kSizeRealBytes = (platform::is_same>::value) ? 8 : 4; + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, true); + mctlass::arch::cp_async_diag( + reinterpret_cast (dst_ptr + v) + kSizeRealBytes); + } else { + /* Write one (1) directly to smem*/ + mctlass::arch::cp_async_diag(dst_ptr + v); + } + } else { + // Elements that are not of diagonal + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, isvalid); + } + + ++iterator_B; + } + + ++this->smem_iterator_B_; + } + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Defines the boundary of a stage of cp.async. + mctlass::arch::cp_async_fence(); + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // + // Clear the remaining tiles of SMEM. This is a functional requirement for some kernels + // so that all accumulator elements outside the GEMM footprint are zero. + // + + if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage) { + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_); + + typename IteratorA::AccessType zero_A; + zero_A.clear(); + + last_smem_iterator_A.set_iteration_index(0); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + last_smem_iterator_A.get()); + + *dst_ptr = zero_A; + + ++last_smem_iterator_A; + } + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB last_smem_iterator_B(this->smem_iterator_B_); + typename IteratorB::AccessType zero_B; + + zero_B.clear(); + last_smem_iterator_B.set_iteration_index(0); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + last_smem_iterator_B.get()); + + *dst_ptr = zero_B; + + ++last_smem_iterator_B; + } + } + + // Waits until kStages-2 stages have committed. + mctlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[2]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpTransformedFragmentA warp_transformed_frag_A[2]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0], + warp_loaded_frag_A[0], warp_loaded_frag_B[0]); + + // tf32x3 kernels use staging accumulation. warp_mma uses a temporary + // accumulator and this temporary accumulator is added to the final + // accumulator once in every mainloop iteration. + plus plus_accum; + + FragmentC tmp_accum; + + if (platform::is_same::value + || platform::is_same::value) { + + tmp_accum.clear(); + } + + // + // Mainloop + // + + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) + warp_mma.transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % 2], + warp_loaded_frag_B[warp_mma_k % 2]); + + if (platform::is_same::value + || platform::is_same::value) { + + warp_mma( + tmp_accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + tmp_accum + ); + + if (warp_mma_k == 0) { + accum = plus_accum(accum, tmp_accum); + tmp_accum.clear(); + } + } else { + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum + ); + } + + // Issue global->shared copies for the this stage + if (warp_mma_k < Base::kWarpGemmIterations - 1) { + int group_start_iteration_A, group_start_iteration_B; + + group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA; + group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, + group_start_iteration_B); + } + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + int group_start_iteration_A, group_start_iteration_B; + group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, + group_start_iteration_B); + + // Inserts a memory fence between stages of cp.async instructions. + mctlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages have committed. + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + } + + // Do any conversions feeding the first stage at the end of the loop so + // we can start right away on mma instructions + if (warp_mma_k + 1 == Base::kWarpGemmIterations) + warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + } + + } + + if (platform::is_same::value + || platform::is_same::value) { + accum = plus_accum(accum, tmp_accum); + } + + if (SharedMemoryClear == SharedMemoryClearOption::kZfill) { + // commit and drain all pending and predicated LDGSTS pnz from the GEMM mainloop + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + } + + // Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/mma_layernorm_mainloop_fusion_multistage.h b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_layernorm_mainloop_fusion_multistage.h new file mode 100644 index 0000000..8b5f59f --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_layernorm_mainloop_fusion_multistage.h @@ -0,0 +1,863 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. + + It loads two loop invariant vectors, mean and var, in the prologue and + stores them in the register file. In the mainloop, it loads two loop + variant vectors, gamma and beta, by using cp.async. We will call + elementwise operation to apply var, mean, gamma, beta between ldmatrix and + warp mma. +*/ + +#pragma once + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" +#include "mctlass/transform/threadblock/predicated_scale_bias_vector_iterator.h" +#include "mctlass/gemm/threadblock/mma_base.h" +#include "mctlass/gemm/warp/layernorm_scale_bias_transform.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Element type of scale and bias vectors + typename ElementScaleBias_, + /// Layout of scale and bias vectors + typename LayoutScaleBias_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// WarpIterator to load Scale or Bias vector from the shared memory + typename WarpIteratorGammaBeta_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class MmaMainloopFusionBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Element type of scale and bias vectors + using ElementScaleBias = ElementScaleBias_; + + /// Layout of scale and bias vectors + using LayoutScaleBias = LayoutScaleBias_; + + ///< Policy describing tuning details + using Policy = Policy_; + + ///< WarpIterator to load Scale or Bias vector from the shared memory + using WarpIteratorGammaBeta = WarpIteratorGammaBeta_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = mctlass::gemm::GemmShape; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = + (WarpGemm::kK / Operator::Policy::MmaShape::kK); + + /// Number of stages + static int const kStages = Stages; + + /// Tensor reference to the A operand + using TensorRefA = TensorRef; + + /// Tensor reference to the scale and bias vectors + using TensorRefGammaBeta = TensorRef; + + /// Tensor reference to the B operand + using TensorRefB = TensorRef; + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + class SharedStorage { + public: + // + // Type definitions + // + + /// Shape of the A matrix operand in shared memory + using ShapeA = MatrixShape; + + /// Shape of the A scale and bias vectors in shared memory + using ShapeGammaBeta = + MatrixShape<1 + Policy::SmemPaddingA::kRow, + 2 * Shape::kK * kStages + Policy::SmemPaddingA::kColumn>; + + /// Shape of the B matrix operand in shared memory + using ShapeB = + MatrixShape; + + public: + // + // Data members + // + + /// Buffer for A operand + AlignedBuffer operand_A; + + /// Buffer for B operand + AlignedBuffer operand_B; + + /// Buffer for A operand Scale and Bias + AlignedBuffer operand_A_gamma_beta; + + public: + + // + // Methods + // + + /// Returns a layout object for the A matrix + MCTLASS_DEVICE + static typename Operator::LayoutA LayoutA() { + return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn}); + } + + /// Returns a layout object for the B matrix + MCTLASS_HOST_DEVICE + static typename Operator::LayoutB LayoutB() { + return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); + } + + /// Returns a layout object for the A scale and bias vectors + MCTLASS_DEVICE + static LayoutScaleBias LayoutScaleBias() { + return LayoutScaleBias::packed( + {ShapeGammaBeta::kRow, ShapeGammaBeta::kColumn}); + } + + /// Returns a TensorRef to the A operand + MCTLASS_HOST_DEVICE + TensorRefA operand_A_ref() { + return TensorRefA{operand_A.data(), LayoutA()}; + } + + /// Returns a TensorRef to the B operand + MCTLASS_HOST_DEVICE + TensorRefB operand_B_ref() { + return TensorRefB{operand_B.data(), LayoutB()}; + } + + /// Returns a TensorRef to the A operand Scale vector + MCTLASS_HOST_DEVICE + TensorRefGammaBeta operand_A_gamma_beta_ref() { + return TensorRefGammaBeta{operand_A_gamma_beta.data(), LayoutScaleBias()}; + } + }; + + protected: + + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A operand from shared memory + typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of A operand scale and bias vector + /// from shared memory + WarpIteratorGammaBeta warp_tile_iterator_A_gamma_beta_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaMainloopFusionBase( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx), + warp_tile_iterator_A_gamma_beta_( + shared_storage.operand_A_gamma_beta_ref(), lane_idx), + warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) {} +}; + + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// Iterates over vectors of var and mean vector in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorVarMean_, + /// Iterates over vectors of scale and bias vector in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorGammaBeta_, + /// Iterates over vectors of scale and bias vector in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorGammaBeta_, + /// Cache operation for scale/bias operand + mctlass::arch::CacheOperation::Kind CacheOpGammaBeta, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// WarpIterator to load Scale or Bias vector from the shared memory + typename WarpIteratorGammaBeta_, + /// Number of stages, + int Stages, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// Used for partial specialization + typename Enable = bool> +class MmaLayernormMainloopFusionMultistage : + public MmaMainloopFusionBase { +public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Iterates over tiles of the var and mean vectors in global memory + using IteratorVarMean = IteratorVarMean_; + ///< Iterates over tiles of the scale and bias vectors in global memory + using IteratorGammaBeta = IteratorGammaBeta_; + ///< WarpIterator to load Scale or Bias vector from the shared memory + using WarpIteratorGammaBeta = WarpIteratorGammaBeta_; + ///< Policy describing tuning details + using Policy = Policy_; + + ///< Base class + using Base = MmaMainloopFusionBase; + + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + using SmemIteratorGammaBeta = SmemIteratorGammaBeta_; + + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + static mctlass::arch::CacheOperation::Kind const kCacheOpGammaBeta = + CacheOpGammaBeta; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + /// Internal structure exposed for introspection. + struct Detail { + + static_assert(Base::kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + }; + + private: + + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + using WarpLoadedFragmentVarMean = typename IteratorVarMean::Fragment; + using WarpLoadedFragmentGammaBeta = + typename WarpIteratorGammaBeta::Fragment; + + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of A operand scale vector to shared memory + SmemIteratorGammaBeta smem_iterator_A_gamma_beta_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + int warp_idx_m_; + + int warp_idx_n_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaLayernormMainloopFusionMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_A_gamma_beta_(shared_storage.operand_A_gamma_beta_ref(), + thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + warp_idx_m_ = warp_idx_mn % Base::WarpCount::kM; + warp_idx_n_ = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m_, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_A_gamma_beta_.add_tile_offset( + {warp_idx_m_, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n_}); + } + + MCTLASS_DEVICE + void copy_tiles_and_advance(IteratorA &iterator_A, + IteratorGammaBeta &iterator_A_gamma_beta, + IteratorB &iterator_B, + int group_start_A = 0, int group_start_B = 0) { + iterator_A.set_iteration_index(group_start_A * + IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_A.get(); + + if (SharedMemoryClear == SharedMemoryClearOption::kZfill) { + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + } else { + mctlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + } + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + } + + // Async Copy for operand A scale and bias vector. Scale and bias vectors + // are small. One iteration is enough. + if (group_start_A == 0) { + typename IteratorGammaBeta::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_gamma_beta_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorGammaBeta::kElementsPerAccess / 8; + + mctlass::arch::cp_async( + dst_ptr, iterator_A_gamma_beta.get(), iterator_A_gamma_beta.valid()); + } + + iterator_B.set_iteration_index(group_start_B * + IteratorB::kAccessesPerVector); + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_B.get(); + + if (SharedMemoryClear == SharedMemoryClearOption::kZfill) { + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + } else { + mctlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + } + + ++iterator_B; + } + ++this->smem_iterator_B_; + } + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< iterator over B operand in global memory + IteratorVarMean iterator_var_mean, + ///< iterator over scale and bias vectors in global memory + IteratorGammaBeta iterator_A_gamma_beta, + ///< initial value of accumulator + FragmentC const &src_accum) { + + // + // Prologue + // + // Issue several complete stages + + WarpLoadedFragmentVarMean warp_loaded_frag_var_mean; + iterator_var_mean.add_tile_offset({0, warp_idx_m_}); + iterator_var_mean.load(warp_loaded_frag_var_mean); + + MCTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; + ++stage, --gemm_k_iterations) { + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_A_gamma_beta.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + int src_bytes = (iterator_A.valid() ? kSrcBytes : 0); + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + + // Async Copy for operand A scale and bias vectors. Scale and bias + // vectors are small. One iteration is enough. + { + typename IteratorGammaBeta::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_gamma_beta_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorGammaBeta::kElementsPerAccess / 8; + + mctlass::arch::cp_async( + dst_ptr, iterator_A_gamma_beta.get(), iterator_A_gamma_beta.valid()); + } + + iterator_B.set_iteration_index(0); + this->smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + } + + ++this->smem_iterator_B_; + } + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_A_gamma_beta.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_A_gamma_beta_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Defines the boundary of a stage of cp.async. + mctlass::arch::cp_async_fence(); + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // Waits until kStages-2 stages have committed. + mctlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[2]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpLoadedFragmentGammaBeta warp_loaded_frag_A_gamma_beta[2]; + WarpTransformedFragmentA warp_transformed_frag_A[2]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + mctlass::gemm::warp::LayernormScaleBiasTransform + elementwise_transform; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_A_gamma_beta_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_A_gamma_beta_.load( + warp_loaded_frag_A_gamma_beta[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_A_gamma_beta_; + ++this->warp_tile_iterator_B_; + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_A_gamma_beta.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0], + warp_loaded_frag_A[0], warp_loaded_frag_B[0]); + + elementwise_transform(warp_transformed_frag_A[0], + warp_loaded_frag_var_mean, + warp_loaded_frag_A_gamma_beta[0]); + + // + // Mainloop + // + + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_A_gamma_beta_.set_kgroup_index( + (warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_A_gamma_beta_.load( + warp_loaded_frag_A_gamma_beta[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_A_gamma_beta_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) { + warp_mma.transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % 2], + warp_loaded_frag_B[warp_mma_k % 2]); + + elementwise_transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_loaded_frag_var_mean, + warp_loaded_frag_A_gamma_beta[warp_mma_k % 2]); + } + + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum + ); + + // Issue global->shared copies for the this stage + if (warp_mma_k < Base::kWarpGemmIterations - 1) { + int group_start_iteration_A, group_start_iteration_B; + + group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA; + group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance(iterator_A, iterator_A_gamma_beta, iterator_B, + group_start_iteration_A, + group_start_iteration_B); + } + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + int group_start_iteration_A, group_start_iteration_B; + group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance(iterator_A, iterator_A_gamma_beta, iterator_B, + group_start_iteration_A, + group_start_iteration_B); + + // Inserts a memory fence between stages of cp.async instructions. + mctlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages have committed. + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_A_gamma_beta.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_A_gamma_beta_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_A_gamma_beta_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_A_gamma_beta_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_A_gamma_beta.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + } + + // Do any conversions feeding the first stage at the end of the loop so + // we can start right away on mma instructions + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + elementwise_transform( + warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_var_mean, + warp_loaded_frag_A_gamma_beta[(warp_mma_k + 1) % 2]); + } + } + + } + + // commit and drain all pending and predicated LDGSTS pnz from the GEMM mainloop + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/mma_multistage.h b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_multistage.h new file mode 100644 index 0000000..a87a68f --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_multistage.h @@ -0,0 +1,745 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" + +#include "mctlass/gemm/threadblock/mma_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// Used for partial specialization + typename Enable = bool> +class MmaMultistage : + public MmaBase { +public: + ///< Base class + using Base = MmaBase; + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + ///< Policy describing tuning details + using Policy = Policy_; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + /// Internal structure exposed for introspection. + struct Detail { + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + // Optional staged-accumulation (e.g., tf32x3 kernels) for improved numerical + // accuracy, where each mainloop iteration first accumulates into a temporary + // set of freshly-cleared accumulators, which are subsequently added to the + // final accumulator set. + static bool const kStagedAccumulation = arch::UseStagedAccumulation::value; + }; + + private: + + + // Structure encapsulating pipeline state live from one iteration to the next + struct PipeState { + + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + /// Temporary accumulator to facilitate staged-accumulation + FragmentC tmp_accum_; + + /// Pair of A fragments used to overlap shared memory loads and math instructions + WarpLoadedFragmentA warp_loaded_frag_A_[2]; + WarpTransformedFragmentA warp_transformed_frag_A_[2]; + + /// Pair of B fragments used to overlap shared memory loads and math instructions + WarpLoadedFragmentB warp_loaded_frag_B_[2]; + WarpTransformedFragmentB warp_transformed_frag_B_[2]; + }; + + + private: + + // + // Data members + // + + /// Warp-level MMA operator + Operator warp_mma_; + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + /// Shared memory write stage index + int smem_write_stage_idx_; + + /// Shared memory read stage index + int smem_read_stage_idx_; + + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx), + smem_write_stage_idx_(0), + smem_read_stage_idx_(0) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + /// Advance shared memory read-iterators to the next stage + MCTLASS_DEVICE + void advance_smem_read_stage() + { + ++smem_read_stage_idx_; + + if (smem_read_stage_idx_ == Base::kStages) { + // Wrap back around to the 'start' of the circular buffer in shared memory + this->warp_tile_iterator_A_.add_tile_offset({0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset({-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations, 0}); + smem_read_stage_idx_ = 0; + } + } + + /// Advance global memory read-iterators and shared memory write-iterators to the stage + MCTLASS_DEVICE + void advance_smem_write_stage( + IteratorA &iterator_A, + IteratorB &iterator_B) + { + // Advance global iterators + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + // Advance shared iterators + smem_iterator_A_.add_tile_offset({0, 1}); + smem_iterator_B_.add_tile_offset({1, 0}); + + // Increment shared memory write stage index + ++smem_write_stage_idx_; + + if (smem_write_stage_idx_ == Base::kStages) { + // Wrap back around to the 'start' of the circular buffer in shared memory + smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx_ = 0; + } + } + + MCTLASS_DEVICE + void copy_tiles_and_advance(IteratorA &iterator_A, IteratorB &iterator_B, + int group_start_A = 0, int group_start_B = 0) { + iterator_A.set_iteration_index(group_start_A * + IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_A.get(); + + if (SharedMemoryClear == SharedMemoryClearOption::kZfill) { + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + } else { + mctlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + } + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + } + + iterator_B.set_iteration_index(group_start_B * + IteratorB::kAccessesPerVector); + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_B.get(); + + if (SharedMemoryClear == SharedMemoryClearOption::kZfill) { + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + } else { + mctlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + } + + ++iterator_B; + } + ++this->smem_iterator_B_; + } + } + } + + /// GEMM prologue. Bootstrap the global->shared memory pipeline by fetching + /// the global fragments needed by the first kStages-1 threadblock mainloop iterations + MCTLASS_DEVICE + void prologue( + IteratorA &iterator_A, ///< [in|out] iterator over A operand in global memory + IteratorB &iterator_B, ///< [in|out] iterator over B operand in global memory + int &gemm_k_iterations) ///< [in|out] number of threadblock mainloop iterations remaining + { + // Issue several complete stages + MCTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; ++stage, --gemm_k_iterations) { + + // Disable global fetching if done with global fetch iterations + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + int src_bytes = (iterator_A.valid() ? kSrcBytes : 0); + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + + iterator_B.set_iteration_index(0); + this->smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + } + + ++this->smem_iterator_B_; + } + + // Move to the next write stage + advance_smem_write_stage(iterator_A, iterator_B); + + // Defines the boundary of a stage of cp.async. + mctlass::arch::cp_async_fence(); + } + + // Optionally clear the remaining stages of SMEM. This is a functional requirement for + // some kernels so that all accumulator elements outside the GEMM footprint are zero. + if (SharedMemoryClear == SharedMemoryClearOption::kClearLastStage) { + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA last_smem_iterator_A(this->smem_iterator_A_); + typename IteratorA::AccessType zero_A; + + zero_A.clear(); + last_smem_iterator_A.set_iteration_index(0); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + last_smem_iterator_A.get()); + + *dst_ptr = zero_A; + + ++last_smem_iterator_A; + } + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB last_smem_iterator_B(this->smem_iterator_B_); + typename IteratorB::AccessType zero_B; + + zero_B.clear(); + last_smem_iterator_B.set_iteration_index(0); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + last_smem_iterator_B.get()); + + *dst_ptr = zero_B; + + ++last_smem_iterator_B; + } + } + } + + + /// Wait until we have at least one completed global fetch stage + MCTLASS_DEVICE + void gmem_wait() + { + // Wait until we have at least one committed global fetch stage. (#uncommitted = Base::kStages - 1 - #committed) + mctlass::arch::cp_async_wait(); + __syncthreads(); + } + + + /// Perform a threadblock mainloop iteration of matrix multiply-accumulate + MCTLASS_DEVICE + void mac_loop_iter( + PipeState &pipe_state, ///< [in|out] loop-carried pipeline state + FragmentC &accum, ///< [in|out] destination accumulator tile + IteratorA &iterator_A, ///< [in|out] iterator over A operand in global memory + IteratorB &iterator_B, ///< [in|out] iterator over B operand in global memory + int &gemm_k_iterations) ///< [in|out] number of threadblock mainloop iterations remaining + { + // Unroll the warp-level MMA tiles of a threadblock's mainloop iteration + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) { + + // Load the next warp-tile's A fragment from shared memory + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_A_.load(pipe_state.warp_loaded_frag_A_[(warp_mma_k + 1) % 2]); + ++this->warp_tile_iterator_A_; + + // Load the next warp-tile's B fragment from shared memory + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.load(pipe_state.warp_loaded_frag_B_[(warp_mma_k + 1) % 2]); + ++this->warp_tile_iterator_B_; + + // Except for the first warp-tile, all warp-tiles convert their incoming shared memory fragments as necessary + if (warp_mma_k > 0) { + warp_mma_.transform( + pipe_state.warp_transformed_frag_A_[warp_mma_k % 2], + pipe_state.warp_transformed_frag_B_[warp_mma_k % 2], + pipe_state.warp_loaded_frag_A_[warp_mma_k % 2], + pipe_state.warp_loaded_frag_B_[warp_mma_k % 2]); + } + + // Execute the current warp-tile of MMA operations + if (Detail::kStagedAccumulation) { + warp_mma_( + pipe_state.tmp_accum_, + pipe_state.warp_transformed_frag_A_[warp_mma_k % 2], + pipe_state.warp_transformed_frag_B_[warp_mma_k % 2], + pipe_state.tmp_accum_ + ); + + // At present, PTX: "cp.async.commit_group", "cp.async.wait_group", "cp.async.wait_all" in memory_sm80.h is not implemented on maca + // gemm/device test: SM80_Device_Her2k_cf32n_cf32n_l_tensor_op_fast_f32.64x64x16_32x32x16 failed + // therefore, block the segment accumulation of fast branch. + // after PTX is implemented ,the following code should be opened again(2023.03.06). + // if (warp_mma_k == 0) { + // plus plus_accum; + // accum = plus_accum(accum, pipe_state.tmp_accum_); + // pipe_state.tmp_accum_.clear(); + // } + } else { + warp_mma_( + accum, + pipe_state.warp_transformed_frag_A_[warp_mma_k % 2], + pipe_state.warp_transformed_frag_B_[warp_mma_k % 2], + accum + ); + } + + // Except for the last warp-tile, all warp-tiles issue their share of + // global->shared fragment copies + if (warp_mma_k < Base::kWarpGemmIterations - 1) { + + int group_start_iteration_A, group_start_iteration_B; + group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA; + group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance( + iterator_A, + iterator_B, + group_start_iteration_A, + group_start_iteration_B); + } + + // The second-to-last warp-tile also: + // - performs the last warp-tile's share of global->shared fragment copies + // - moves to the next global fetch stage + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + + // Performs the last warp-tile's share of global->shared fragment copies + int group_start_iteration_A = (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + int group_start_iteration_B = (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance( + iterator_A, + iterator_B, + group_start_iteration_A, + group_start_iteration_B); + + // Inserts a memory fence between stages of cp.async instructions. + mctlass::arch::cp_async_fence(); + + // Wait until we have at least one completed global fetch stage + gmem_wait(); + + // Move to the next global fetch stage + advance_smem_write_stage(iterator_A, iterator_B); + advance_smem_read_stage(); + + // Disable global fetching when done with global fetch iterations + --gemm_k_iterations; + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + } + + // The last warp-tile also converts the shared memory fragments used by + // the first warp-tile of the next iteration, if necessary (so we can + // immediately start issuing MMA instructions at the top of the loop ) + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + + warp_mma_.transform( + pipe_state.warp_transformed_frag_A_[(warp_mma_k + 1) % 2], + pipe_state.warp_transformed_frag_B_[(warp_mma_k + 1) % 2], + pipe_state.warp_loaded_frag_A_[(warp_mma_k + 1) % 2], + pipe_state.warp_loaded_frag_B_[(warp_mma_k + 1) % 2]); + } + + } + } + + + /// Perform the specified number of threadblock mainloop iterations of matrix + /// multiply-accumulate. Assumes prologue has been initiated. + MCTLASS_DEVICE + void gemm_iters( + int gemm_k_iterations, ///< number of threadblock mainloop iterations + FragmentC &accum, ///< [in|out] accumulator tile + IteratorA &iterator_A, ///< [in|out] iterator over A operand in global memory + IteratorB &iterator_B) ///< [in|out] iterator over B operand in global memory + { + PipeState pipe_state; + + // Disable global fetching if done with global fetch iterations + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + // Load first warp-tile's A fragment from shared memory + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_A_.load(pipe_state.warp_loaded_frag_A_[0]); + ++this->warp_tile_iterator_A_; + + // Load first warp-tile's B fragment from shared memory + this->warp_tile_iterator_B_.set_kgroup_index(0); + this->warp_tile_iterator_B_.load(pipe_state.warp_loaded_frag_B_[0]); + ++this->warp_tile_iterator_B_; + + // Transform, if necessary, the first warp-tile's shared memory fragments + warp_mma_.transform( + pipe_state.warp_transformed_frag_A_[0], + pipe_state.warp_transformed_frag_B_[0], + pipe_state.warp_loaded_frag_A_[0], + pipe_state.warp_loaded_frag_B_[0]); + + if (Detail::kStagedAccumulation) { + pipe_state.tmp_accum_.clear(); + } + + // Mainloop + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { + mac_loop_iter( + pipe_state, + accum, + iterator_A, + iterator_B, + gemm_k_iterations); + } + + if (Detail::kStagedAccumulation) { + plus plus_accum; + accum = plus_accum(accum, pipe_state.tmp_accum_); + } + + // Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } + + + /// Prepares the class for another prologue. + MCTLASS_DEVICE + void wind_down() + { + // Catch-up the smem-read iterator to the smem-write iterator (so this class can be reused for another tile's prologue) + + // First, increment remaining warp tiles to get to the next full stage. (Ideally we would + // just decrement one tile, but not all iterators implement --() decrement.) + #pragma unroll + for (int warp_mma_k = 1; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) + { + this->warp_tile_iterator_A_.set_kgroup_index(warp_mma_k); + this->warp_tile_iterator_B_.set_kgroup_index(warp_mma_k); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + } + smem_read_stage_idx_++; + + // Then wrap back two full stages (one for the tile advancing we just did, and one to catch the write iterators) + static const int kStageIters = Policy::kPartitionsK * Base::kWarpGemmIterations; + if (smem_read_stage_idx_ > 1) + { + this->warp_tile_iterator_A_.add_tile_offset({0, (-2 * kStageIters)}); + this->warp_tile_iterator_B_.add_tile_offset({(-2 * kStageIters), 0}); + } + else + { + this->warp_tile_iterator_A_.add_tile_offset({0, ((Base::kStages - 2) * kStageIters)}); + this->warp_tile_iterator_B_.add_tile_offset({((Base::kStages - 2) * kStageIters), 0}); + } + smem_read_stage_idx_ = smem_write_stage_idx_; + } + + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< initial value of accumulator + FragmentC const &src_accum) { + + // Prologue (start fetching iterations of global fragments into shared memory) + prologue(iterator_A, iterator_B, gemm_k_iterations); + + // Wait until we have at least one completed global fetch stage + gmem_wait(); + + // Initialize destination accumulators with source accumulators + accum = src_accum; + + // Perform the MAC-iterations + gemm_iters(gemm_k_iterations, accum, iterator_A, iterator_B); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/mma_pipelined.h b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_pipelined.h new file mode 100644 index 0000000..7268973 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_pipelined.h @@ -0,0 +1,439 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/aligned_buffer.h" +#include "mctlass/numeric_conversion.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/mma_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Transformation applied to A operand + typename TransformA_ = NumericArrayConverter< + typename SmemIteratorA_::Element, + typename IteratorA_::Element, + IteratorA_::Fragment::kElements>, + /// + /// Transformation applied to B operand + typename TransformB_ = NumericArrayConverter< + typename SmemIteratorB_::Element, + typename IteratorB_::Element, + IteratorB_::Fragment::kElements>, + /// Used for partial specialization + typename Enable = bool +> +class MmaPipelined : public MmaBase { +public: + + ///< Base class + using Base = MmaBase; + + using Shape = Shape_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using IteratorA = IteratorA_; ///< Iterates over tiles of A operand in global memory + using IteratorB = IteratorB_; ///< Iterates over tiles of B operand in global memory + using ElementC = ElementC_; ///< Data type of accumulator matrix + using LayoutC = LayoutC_; ///< Layout of accumulator matrix + using Policy = Policy_; ///< Policy describing tuning details + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + using TransformA = TransformA_; + using TransformB = TransformB_; + + // + // Dependent types + // + + /// Fragment of operand A loaded from global memory + using FragmentA = typename IteratorA::Fragment; + + /// Fragment of operand B loaded from global memory + using FragmentB = typename IteratorB::Fragment; + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Obtain the arch tag from the warp-level operator + using ArchTag = typename Policy::Operator::ArchTag; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + // staticaly assert kStages for MmaPipelined is two (Double-buffered pipeline) + static_assert((Base::kStages==2), "MmaPipelined requires kStages set to value 2"); + +protected: + + // + // Data members + // + + /// Warp-level MMA operator + Operator warp_mma; + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + ///< transformation applied to A fragment + TransformA transform_A_; + + ///< transformation applied to B fragment + TransformB transform_B_; + + /// Shared memory write stage index + int smem_write_stage_idx; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaPipelined( + typename Base::SharedStorage &shared_storage, ///< Shared storage needed for internal use by threadblock-scoped GEMM + int thread_idx, ///< ID within the threadblock + int warp_idx, ///< ID of warp + int lane_idx, ///< ID of each thread within a warp + TransformA transform_A = TransformA(), ///< transformation applied to A fragment + TransformB transform_B = TransformB() ///< transformation applied to B fragment + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx), + transform_A_(transform_A), + transform_B_(transform_B), + smem_write_stage_idx(0) + { + + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + + /// Advance shared memory write-iterators to the next stage + MCTLASS_DEVICE + void advance_smem_write_stage() + { + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + // Add negative offsets to return iterators to the 'start' of the circular buffer in shared memory + if (smem_write_stage_idx == 1) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + } + + smem_write_stage_idx ^= 1; + } + + /// Advance shared memory read- and write-iterators to the next stage + MCTLASS_DEVICE + void advance_smem_stages() + { + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + // Add negative offsets to return iterators to the 'start' of the circular buffer in shared memory + if (smem_write_stage_idx == 1) { + // wrap write stage + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + } + else + { + // wrap read stage + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations, 0}); + } + + smem_write_stage_idx ^= 1; + } + + + /// GEMM prologue. Bootstrap the global->shared memory pipeline by fetching + /// the global fragments needed by the first kStages-1 threadblock mainloop iterations + MCTLASS_DEVICE + void prologue( + IteratorA &iterator_A, ///< [in|out] iterator over A operand in global memory + IteratorB &iterator_B, ///< [in|out] iterator over B operand in global memory + int &gemm_k_iterations) ///< [in|out] number of threadblock mainloop iterations remaining + { + // The last kblock is loaded in the prolog + + // Load A fragment from global A + FragmentA tb_frag_A; + tb_frag_A.clear(); + iterator_A.load(tb_frag_A); + ++iterator_A; + + // Load B fragment from global B + FragmentB tb_frag_B; + tb_frag_B.clear(); + iterator_B.load(tb_frag_B); + ++iterator_B; + + // Store A and B fragments to shared + this->smem_iterator_A_.store(transform_A_(tb_frag_A)); + this->smem_iterator_B_.store(transform_B_(tb_frag_B)); + + // Advance write stage + advance_smem_write_stage(); + } + + /// Wait until we have at least one completed global fetch stage + MCTLASS_DEVICE + void gmem_wait() + { + __syncthreads(); + } + + + /// Perform the specified number of threadblock mainloop iterations of matrix + /// multiply-accumulate. Assumes prologue has been initiated. + MCTLASS_DEVICE + void gemm_iters( + int gemm_k_iterations, ///< number of threadblock mainloop iterations + FragmentC &accum, ///< [in|out] accumulator tile + IteratorA &iterator_A, ///< [in|out] iterator over A operand in global memory + IteratorB &iterator_B) ///< [in|out] iterator over B operand in global memory + { + using WarpFragmentA = typename Operator::FragmentA; + using WarpFragmentB = typename Operator::FragmentB; + + // Pair of fragments used to overlap shared memory loads and math instructions + WarpFragmentA warp_frag_A[2]; + WarpFragmentB warp_frag_B[2]; + + // Load A fragment from shared A + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_A_.load(warp_frag_A[0]); + ++this->warp_tile_iterator_A_; + + // Load B fragment from shared B + this->warp_tile_iterator_B_.set_kgroup_index(0); + this->warp_tile_iterator_B_.load(warp_frag_B[0]); + ++this->warp_tile_iterator_B_; + + // Pair of fragments used to overlap global memory loads and math instructions; + FragmentA tb_frag_A; + FragmentB tb_frag_B; + + // Avoid reading out of bounds + iterator_A.clear_mask(gemm_k_iterations <= 1); + iterator_B.clear_mask(gemm_k_iterations <= 1); + + // + // Mainloop + // + + // Note: The main loop does not support Base::kWarpGemmIterations == 2. + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > 0; --gemm_k_iterations) { + // + // Loop over GEMM K dimension + // + + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if this is the last group + // as the case may be. + + if (warp_mma_k == Base::kWarpGemmIterations - 1) { + + // Write fragments to shared memory + this->smem_iterator_A_.store(transform_A_(tb_frag_A)); + + this->smem_iterator_B_.store(transform_B_(tb_frag_B)); + + // Wait until we have at least one completed global fetch stage + gmem_wait(); + + // Advance smem read and write stages + advance_smem_stages(); + } + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k == 0) { + + // Load fragment from global A + tb_frag_A.clear(); + iterator_A.load(tb_frag_A); + ++iterator_A; + + // Load fragment from global B + tb_frag_B.clear(); + iterator_B.load(tb_frag_B); + ++iterator_B; + + // Avoid reading out of bounds if this was the last loop iteration + iterator_A.clear_mask(gemm_k_iterations <= 2); + iterator_B.clear_mask(gemm_k_iterations <= 2); + } + + warp_mma( + accum, + warp_frag_A[warp_mma_k % 2], + warp_frag_B[warp_mma_k % 2], + accum); + } + } + + } + + + /// Prepares the class for another prologue. + MCTLASS_DEVICE + void wind_down() + { + // First, increment remaining warp tiles to catch it up with the write stage. + #pragma unroll + for (int warp_mma_k = 1; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) + { + this->warp_tile_iterator_A_.set_kgroup_index(warp_mma_k); + this->warp_tile_iterator_B_.set_kgroup_index(warp_mma_k); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + } + + // If we bumped the read iterators to the end of the circular buffer, wrap them around to + // align them with the write iterators + if (smem_write_stage_idx == 0) + { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations, 0}); + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + int gemm_k_iterations, ///< number of iterations of the mainloop + FragmentC &accum, ///< destination accumulator tile + IteratorA iterator_A, ///< iterator over A operand in global memory + IteratorB iterator_B, ///< iterator over B operand in global memory + FragmentC const &src_accum) ///< source accumulator tile + { + // Prologue + prologue(iterator_A, iterator_B, gemm_k_iterations); + + // Wait until we have at least one completed global fetch stage + gmem_wait(); + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // Perform the MAC-iterations + gemm_iters(gemm_k_iterations, accum, iterator_A, iterator_B); + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/mma_planar_complex_base.h b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_planar_complex_base.h new file mode 100644 index 0000000..0801f56 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_planar_complex_base.h @@ -0,0 +1,208 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class MmaPlanarComplexBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Policy describing tuning details + using Policy = Policy_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = GemmShape; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = + (WarpGemm::kK / Operator::Policy::MmaShape::kK); + + /// Number of stages + static int const kStages = Stages; + + /// Tensor reference to the A operand + using TensorRefA = TensorRef; + + /// Tensor reference to the B operand + using TensorRefB = TensorRef; + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + class SharedStorage { + public: + // + // Type definitions + // + + /// Shape of the A matrix operand in shared memory + using ShapeA = MatrixShape; + + /// Stride to the imaginary part of the A operand + static int const kImaginaryStrideA = ShapeA::kCount; + + /// Shape of the B matrix operand in shared memory + using ShapeB = + MatrixShape; + + /// Stride to the imaginary part of the A operand + static int const kImaginaryStrideB = ShapeB::kCount; + + public: + // + // Data members + // + + /// Buffer for A operand + AlignedBuffer operand_A; + + /// Buffer for B operand + AlignedBuffer operand_B; + + public: + + // + // Methods + // + + /// Returns a layout object for the A matrix + MCTLASS_DEVICE + static typename Operator::LayoutA LayoutA() { + return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn}); + } + + /// Returns a layout object for the B matrix + MCTLASS_HOST_DEVICE + static typename Operator::LayoutB LayoutB() { + return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); + } + + /// Returns a TensorRef to the A operand + MCTLASS_HOST_DEVICE + TensorRefA operand_A_ref() { + return TensorRefA{operand_A.data(), LayoutA()}; + } + + /// Returns a TensorRef to the B operand + MCTLASS_HOST_DEVICE + TensorRefB operand_B_ref() { + return TensorRefB{operand_B.data(), LayoutB()}; + } + }; + + protected: + + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A operand from shared memory + typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaPlanarComplexBase( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx), + warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) { + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/mma_planar_complex_multistage.h b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_planar_complex_multistage.h new file mode 100644 index 0000000..2e83e16 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_planar_complex_multistage.h @@ -0,0 +1,646 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/array_planar_complex.h" +#include "mctlass/functional.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/mma_planar_complex_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Transformation applied to A + ComplexTransform TransformA = ComplexTransform::kNone, + /// Transformation applied to B + ComplexTransform TransformB = ComplexTransform::kNone +> +class MmaPlanarComplexMultistage : + public MmaPlanarComplexBase { +public: + ///< Base class + using Base = MmaPlanarComplexBase; + + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + + ///< Data type of accumulator matrix + using ElementC = ElementC_; + + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + + ///< Policy describing tuning details + using Policy = Policy_; + + ///< Archtecture tag + using ArchTag = arch::Sm80; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + /// Transformation applied to A + static ComplexTransform const kTransformA = TransformA; + + /// Transformation applied to B + static ComplexTransform const kTransformB = TransformB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC = ArrayPlanarComplex< + typename Policy::Operator::FragmentC::Element, + Policy::Operator::FragmentC::kElements + >; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Internal structure exposed for introspection. + struct Detail { + + static_assert(Base::kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + /// Number of cp.async instructions to load one stage of operand A + static int const TBLoadIterationsA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const TBLoadIterationsB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + static int const kAccessesPerGroupA = + (TBLoadIterationsA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + static int const kAccessesPerGroupB = + (TBLoadIterationsB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + }; + + private: + + using WarpFragmentA = typename Operator::FragmentA; + using WarpFragmentB = typename Operator::FragmentB; + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaPlanarComplexMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + +private: + + MCTLASS_DEVICE + void copy_tiles_and_advance( + IteratorA &iterator_A_real, + IteratorA &iterator_A_imag, + + IteratorB &iterator_B_real, + IteratorB &iterator_B_imag, + + int group_start_A = 0, + int group_start_B = 0) { + + iterator_A_real.set_iteration_index(group_start_A * IteratorA::kAccessesPerVector); + iterator_A_imag.set_iteration_index(group_start_A * IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Load for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast(this->smem_iterator_A_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / IteratorA::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + + auto gmem_ptr_real = iterator_A_real.get(); + auto gmem_ptr_imag = iterator_A_imag.get(); + + bool pred_guard = iterator_A_real.valid(); + mctlass::arch::cp_async( + dst_ptr + v, + gmem_ptr_real, + pred_guard); + mctlass::arch::cp_async( + dst_ptr + v + (Base::SharedStorage::kImaginaryStrideA / IteratorA::ThreadMap::kElementsPerAccess), + reinterpret_cast(gmem_ptr_imag), + pred_guard); + + ++iterator_A_real; + ++iterator_A_imag; + } + + ++this->smem_iterator_A_; + } + + iterator_B_real.set_iteration_index(group_start_B * IteratorB::kAccessesPerVector); + iterator_B_imag.set_iteration_index(group_start_B * IteratorB::kAccessesPerVector); + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Load for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast(this->smem_iterator_B_.get()); + + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / IteratorB::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + auto gmem_ptr_real = iterator_B_real.get(); + auto gmem_ptr_imag = iterator_B_imag.get(); + + bool pred_guard = iterator_B_real.valid(); + mctlass::arch::cp_async( + dst_ptr + v, + gmem_ptr_real, + pred_guard); + mctlass::arch::cp_async( + dst_ptr + v + (Base::SharedStorage::kImaginaryStrideB / IteratorB::ThreadMap::kElementsPerAccess), + reinterpret_cast(gmem_ptr_imag), + pred_guard); + + ++iterator_B_real; + ++iterator_B_imag; + } + ++this->smem_iterator_B_; + } + } + + MCTLASS_DEVICE + void warp_mma_planar_complex( + Operator & warp_mma, + FragmentC &accum, + WarpFragmentA const & real_A, + WarpFragmentA const & imag_A, + WarpFragmentB const & real_B, + WarpFragmentB const & imag_B) { + + mctlass::negate> neg_op_B; + + WarpFragmentB neg_real_B = neg_op_B(real_B); + WarpFragmentB neg_imag_B = neg_op_B(imag_B); + + warp_mma(accum.real, real_A, real_B, accum.real); + + if (kTransformB == ComplexTransform::kNone) { + warp_mma(accum.imag, real_A, imag_B, accum.imag); + } + else { + warp_mma(accum.imag, real_A, neg_imag_B, accum.imag); + } + + if (kTransformA == ComplexTransform::kNone) { + warp_mma(accum.imag, imag_A, real_B, accum.imag); + } + else { + warp_mma(accum.imag, imag_A, neg_real_B, accum.imag); + } + + if (kTransformA == ComplexTransform::kNone ^ kTransformB == ComplexTransform::kNone) { + warp_mma(accum.real, imag_A, imag_B, accum.real); + } + else { + warp_mma(accum.real, imag_A, neg_imag_B, accum.real); + } + } + +public: + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A_real, + ///< iterator over A operand in global memory + IteratorA iterator_A_imag, + ///< iterator over B operand in global memory + IteratorB iterator_B_real, + ///< iterator over B operand in global memory + IteratorB iterator_B_imag, + ///< initial value of accumulator + FragmentC const &src_accum) { + + // + // Prologue + // + + // Issue several complete stages + MCTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; + ++stage, --gemm_k_iterations) { + + iterator_A_real.clear_mask(gemm_k_iterations == 0); + iterator_A_imag.clear_mask(gemm_k_iterations == 0); + iterator_B_real.clear_mask(gemm_k_iterations == 0); + iterator_B_imag.clear_mask(gemm_k_iterations == 0); + + iterator_A_real.set_iteration_index(0); + iterator_A_imag.set_iteration_index(0); + + this->smem_iterator_A_.set_iteration_index(0); + + // Load for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::TBLoadIterationsA; ++j) { + + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast(this->smem_iterator_A_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / IteratorA::kAccessesPerVector / 8; + + bool pred_guard = iterator_A_real.valid(); + + auto src_ptr_real = iterator_A_real.get(); + auto src_ptr_imag = iterator_A_imag.get(); + + mctlass::arch::cp_async_zfill( + dst_ptr + v, src_ptr_real, pred_guard); + + mctlass::arch::cp_async_zfill( + dst_ptr + v + + Base::SharedStorage::kImaginaryStrideA / + IteratorA::ThreadMap::kElementsPerAccess, + reinterpret_cast(src_ptr_imag), + pred_guard); + + ++iterator_A_real; + ++iterator_A_imag; + } + + ++this->smem_iterator_A_; + } + + iterator_B_real.set_iteration_index(0); + iterator_B_imag.set_iteration_index(0); + + this->smem_iterator_B_.set_iteration_index(0); + + // Load for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::TBLoadIterationsB; ++j) { + + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast(this->smem_iterator_B_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / IteratorB::kAccessesPerVector / 8; + + bool pred_guard = iterator_B_real.valid(); + + auto src_ptr_real = iterator_B_real.get(); + auto src_ptr_imag = iterator_B_imag.get(); + + mctlass::arch::cp_async_zfill( + dst_ptr + v, src_ptr_real, pred_guard); + + mctlass::arch::cp_async_zfill( + dst_ptr + v + + Base::SharedStorage::kImaginaryStrideB / + IteratorB::ThreadMap::kElementsPerAccess, + reinterpret_cast(src_ptr_imag), + pred_guard); + + ++iterator_B_real; + ++iterator_B_imag; + } + + ++this->smem_iterator_B_; + } + + // Move to the next stage + iterator_A_real.add_tile_offset({0, 1}); + iterator_A_imag.add_tile_offset({0, 1}); + + iterator_B_real.add_tile_offset({1, 0}); + iterator_B_imag.add_tile_offset({1, 0}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Inserts a memory fence between stages of cp.async instructions + mctlass::arch::cp_async_fence(); + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // Blocks until all but kStages-2 cp.async stages have committed. + mctlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + + WarpFragmentA warp_frag_real_A[2]; + WarpFragmentA warp_frag_imag_A[2]; + + WarpFragmentB warp_frag_real_B[2]; + WarpFragmentB warp_frag_imag_B[2]; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_frag_real_A[0]); + this->warp_tile_iterator_A_.load_with_pointer_offset(warp_frag_imag_A[0], Base::SharedStorage::kImaginaryStrideA); + + this->warp_tile_iterator_B_.load(warp_frag_real_B[0]); + this->warp_tile_iterator_B_.load_with_pointer_offset(warp_frag_imag_B[0], Base::SharedStorage::kImaginaryStrideB); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + iterator_A_real.clear_mask(gemm_k_iterations == 0); + iterator_A_imag.clear_mask(gemm_k_iterations == 0); + iterator_B_real.clear_mask(gemm_k_iterations == 0); + iterator_B_imag.clear_mask(gemm_k_iterations == 0); + + // Start issuing the first group of the next stage outside of the mainloop + copy_tiles_and_advance(iterator_A_real, iterator_A_imag, iterator_B_real, iterator_B_imag); + + Operator warp_mma; + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + // + // Mainloop + // + + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_frag_real_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_A_.load_with_pointer_offset(warp_frag_imag_A[(warp_mma_k + 1) % 2], Base::SharedStorage::kImaginaryStrideA); + + this->warp_tile_iterator_B_.load(warp_frag_real_B[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load_with_pointer_offset(warp_frag_imag_B[(warp_mma_k + 1) % 2], Base::SharedStorage::kImaginaryStrideB); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + // Issue global->shared copies for the next stage + int group_start_iteration_A, group_start_iteration_B; + + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + group_start_iteration_A = 0; + group_start_iteration_B = 0; + } + else { + group_start_iteration_A = (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + } + + copy_tiles_and_advance( + iterator_A_real, + iterator_A_imag, + iterator_B_real, + iterator_B_imag, + group_start_iteration_A, + group_start_iteration_B); + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + // Inserts a memory fence between stages of cp.async instructions + mctlass::arch::cp_async_fence(); + + // Blocks until all but kStages-2 cp.async stages have committed. + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A_real.add_tile_offset({0, 1}); + iterator_A_imag.add_tile_offset({0, 1}); + + iterator_B_real.add_tile_offset({1, 0}); + iterator_B_imag.add_tile_offset({1, 0}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + iterator_A_real.clear_mask(gemm_k_iterations == 0); + iterator_A_imag.clear_mask(gemm_k_iterations == 0); + iterator_B_real.clear_mask(gemm_k_iterations == 0); + iterator_B_imag.clear_mask(gemm_k_iterations == 0); + } + + warp_mma_planar_complex( + warp_mma, + accum, + warp_frag_real_A[warp_mma_k % 2], + warp_frag_imag_A[warp_mma_k % 2], + warp_frag_real_B[warp_mma_k % 2], + warp_frag_imag_B[warp_mma_k % 2]); + } + + } + + + // Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/mma_planar_complex_pipelined.h b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_planar_complex_pipelined.h new file mode 100644 index 0000000..b9923b6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_planar_complex_pipelined.h @@ -0,0 +1,424 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/aligned_buffer.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/mma_planar_complex_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Transformation applied to A + ComplexTransform TransformA = ComplexTransform::kNone, + /// Transformation applied to B + ComplexTransform TransformB = ComplexTransform::kNone +> +class MmaPlanarComplexPipelined : + public MmaPlanarComplexBase { +public: + ///< Base class + using Base = MmaPlanarComplexBase; + + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + + ///< Data type of accumulator matrix + using ElementC = ElementC_; + + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + + ///< Policy describing tuning details + using Policy = Policy_; + + using ArchTag = typename Policy::Operator::ArchTag; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + /// Transformation applied to A + static ComplexTransform const kTransformA = TransformA; + + /// Transformation applied to B + static ComplexTransform const kTransformB = TransformB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC = ArrayPlanarComplex< + typename Policy::Operator::FragmentC::Element, + Policy::Operator::FragmentC::kElements + >; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + private: + + using FragmentA = typename IteratorA::Fragment; + using FragmentB = typename IteratorB::Fragment; + using WarpFragmentA = typename Operator::FragmentA; + using WarpFragmentB = typename Operator::FragmentB; + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaPlanarComplexPipelined( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + +private: + + MCTLASS_DEVICE + void warp_mma_planar_complex( + Operator & warp_mma, + FragmentC &accum, + WarpFragmentA const & real_A, + WarpFragmentA const & imag_A, + WarpFragmentB const & real_B, + WarpFragmentB const & imag_B) { + + mctlass::negate> neg_op_B; + + WarpFragmentB neg_real_B = neg_op_B(real_B); + WarpFragmentB neg_imag_B = neg_op_B(imag_B); + + warp_mma(accum.real, real_A, real_B, accum.real); + + if (kTransformB == ComplexTransform::kNone) { + warp_mma(accum.imag, real_A, imag_B, accum.imag); + } + else { + warp_mma(accum.imag, real_A, neg_imag_B, accum.imag); + } + + if (kTransformA == ComplexTransform::kNone) { + warp_mma(accum.imag, imag_A, real_B, accum.imag); + } + else { + warp_mma(accum.imag, imag_A, neg_real_B, accum.imag); + } + + if (kTransformA == ComplexTransform::kNone ^ kTransformB == ComplexTransform::kNone) { + warp_mma(accum.real, imag_A, imag_B, accum.real); + } + else { + warp_mma(accum.real, imag_A, neg_imag_B, accum.real); + } + } + +public: + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A_real, + ///< iterator over A operand in global memory + IteratorA iterator_A_imag, + ///< iterator over B operand in global memory + IteratorB iterator_B_real, + ///< iterator over B operand in global memory + IteratorB iterator_B_imag, + ///< initial value of accumulator + FragmentC const &src_accum) { + + // + // Prologue + // + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + FragmentA tb_frag_A_real; + FragmentA tb_frag_A_imag; + + FragmentB tb_frag_B_real; + FragmentB tb_frag_B_imag; + + tb_frag_A_real.clear(); + tb_frag_A_imag.clear(); + + tb_frag_B_real.clear(); + tb_frag_B_imag.clear(); + + // The last kblock is loaded in the prolog + iterator_A_real.load(tb_frag_A_real); + iterator_A_imag.load(tb_frag_A_imag); + + iterator_B_real.load(tb_frag_B_real); + iterator_B_imag.load(tb_frag_B_imag); + + ++iterator_A_real; + ++iterator_A_imag; + + ++iterator_B_real; + ++iterator_B_imag; + + this->smem_iterator_A_.store(tb_frag_A_real); + this->smem_iterator_A_.store_with_pointer_offset(tb_frag_A_imag, Base::SharedStorage::kImaginaryStrideA); + + this->smem_iterator_B_.store(tb_frag_B_real); + this->smem_iterator_B_.store_with_pointer_offset(tb_frag_B_imag, Base::SharedStorage::kImaginaryStrideB); + + ++this->smem_iterator_A_; + ++this->smem_iterator_B_; + + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math instructions + WarpFragmentA warp_frag_real_A[2]; + WarpFragmentA warp_frag_imag_A[2]; + + WarpFragmentB warp_frag_real_B[2]; + WarpFragmentB warp_frag_imag_B[2]; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_frag_real_A[0]); + this->warp_tile_iterator_A_.load_with_pointer_offset(warp_frag_imag_A[0], Base::SharedStorage::kImaginaryStrideA); + + this->warp_tile_iterator_B_.load(warp_frag_real_B[0]); + this->warp_tile_iterator_B_.load_with_pointer_offset(warp_frag_imag_B[0], Base::SharedStorage::kImaginaryStrideB); + + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + Operator warp_mma; + + int smem_write_stage_idx = 1; + + // Avoid reading out of bounds + iterator_A_real.clear_mask(gemm_k_iterations <= 1); + iterator_A_imag.clear_mask(gemm_k_iterations <= 1); + + iterator_B_real.clear_mask(gemm_k_iterations <= 1); + iterator_B_imag.clear_mask(gemm_k_iterations <= 1); + + // Issue loads during the first warp-level matrix multiply-add *AFTER* issuing + // shared memory loads (which have the tighest latency requirement). + + // + // Mainloop + // + + // Note: The main loop does not support Base::kWarpGemmIterations == 2. + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > 0; --gemm_k_iterations) { + // + // Loop over GEMM K dimension + // + + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if this is the last group + // as the case may be. + + if (warp_mma_k == Base::kWarpGemmIterations - 1) { + + // Write fragments to shared memory + this->smem_iterator_A_.store(tb_frag_A_real); + this->smem_iterator_A_.store_with_pointer_offset(tb_frag_A_imag, Base::SharedStorage::kImaginaryStrideA); + + this->smem_iterator_B_.store(tb_frag_B_real); + this->smem_iterator_B_.store_with_pointer_offset(tb_frag_B_imag, Base::SharedStorage::kImaginaryStrideB); + + __syncthreads(); + + ++this->smem_iterator_B_; + ++this->smem_iterator_A_; + + // Add negative offsets to return iterators to the 'start' of the circular buffer in shared memory + if (smem_write_stage_idx == 1) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + } + else { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * Base::kWarpGemmIterations, + 0}); + } + + smem_write_stage_idx ^= 1; + } + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_frag_real_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_A_.load_with_pointer_offset(warp_frag_imag_A[(warp_mma_k + 1) % 2], Base::SharedStorage::kImaginaryStrideA); + + this->warp_tile_iterator_B_.load(warp_frag_real_B[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load_with_pointer_offset(warp_frag_imag_B[(warp_mma_k + 1) % 2], Base::SharedStorage::kImaginaryStrideB); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k == 0) { + + iterator_A_real.load(tb_frag_A_real); + iterator_A_imag.load(tb_frag_A_imag); + + iterator_B_real.load(tb_frag_B_real); + iterator_B_imag.load(tb_frag_B_imag); + + ++iterator_A_real; + ++iterator_A_imag; + ++iterator_B_real; + ++iterator_B_imag; + + // Avoid reading out of bounds if this was the last loop iteration + iterator_A_real.clear_mask(gemm_k_iterations <= 2); + iterator_A_imag.clear_mask(gemm_k_iterations <= 2); + iterator_B_real.clear_mask(gemm_k_iterations <= 2); + iterator_B_imag.clear_mask(gemm_k_iterations <= 2); + } + + warp_mma_planar_complex( + warp_mma, + accum, + warp_frag_real_A[warp_mma_k % 2], + warp_frag_imag_A[warp_mma_k % 2], + warp_frag_real_B[warp_mma_k % 2], + warp_frag_imag_B[warp_mma_k % 2]); + } + } + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/mma_singlestage.h b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_singlestage.h new file mode 100644 index 0000000..85aadc6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_singlestage.h @@ -0,0 +1,265 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/aligned_buffer.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/threadblock/mma_base.h" + + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Used for partial specialization + typename Enable = bool +> +class MmaSingleStage : public MmaBase { +public: + + ///< Base class + using Base = MmaBase; + + using Shape = Shape_; ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using IteratorA = IteratorA_; ///< Iterates over tiles of A operand in global memory + using IteratorB = IteratorB_; ///< Iterates over tiles of B operand in global memory + using ElementC = ElementC_; ///< Data type of accumulator matrix + using LayoutC = LayoutC_; ///< Layout of accumulator matrix + using Policy = Policy_; ///< Policy describing tuning details + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + // + // Dependent types + // + + /// Fragment of operand A loaded from global memory + using FragmentA = typename IteratorA::Fragment; + + /// Fragment of operand B loaded from global memory + using FragmentB = typename IteratorB::Fragment; + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + using ArchTag = arch::Sm70; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + // staticaly assert kStages for MmaSingleStage is 1 (single stage mma pipeline) + static_assert((Base::kStages==1), "MmaSingleStage requires kStages set to value 1"); +private: + + using WarpFragmentA = typename Operator::FragmentA; + using WarpFragmentB = typename Operator::FragmentB; + +protected: + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaSingleStage( + typename Base::SharedStorage &shared_storage, ///< Shared storage needed for internal use by threadblock-scoped GEMM + int thread_idx, ///< ID within the threadblock + int warp_idx, ///< ID of warp + int lane_idx ///< ID of each thread within a warp + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) { + + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset({warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset({Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + int gemm_k_iterations, ///< number of iterations of the mainloop + FragmentC &accum, ///< destination accumulator tile + IteratorA iterator_A, ///< iterator over A operand in global memory + IteratorB iterator_B, ///< iterator over B operand in global memory + FragmentC const &src_accum) { ///< source accumualtor tile + + // + // Prologue + // + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + FragmentA tb_frag_A; + FragmentB tb_frag_B; + + tb_frag_A.clear(); + tb_frag_B.clear(); + + // The last kblock is loaded in the prolog + iterator_A.load(tb_frag_A); + iterator_B.load(tb_frag_B); + + ++iterator_A; + ++iterator_B; + + // Pair of fragments used to overlap shared memory loads and math instructions + WarpFragmentA warp_frag_A; + WarpFragmentB warp_frag_B; + + Operator warp_mma; + + // Avoid reading out of bounds + iterator_A.clear_mask(gemm_k_iterations <= 1); + iterator_B.clear_mask(gemm_k_iterations <= 1); + + // + // Mainloop + // + + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > 0; --gemm_k_iterations) { + this->smem_iterator_A_.store(tb_frag_A); + this->smem_iterator_B_.store(tb_frag_B); + + __syncthreads(); + + // + // Loop over GEMM K dimension + // + + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if this is the last group + // as the case may be. + + this->warp_tile_iterator_A_.set_kgroup_index(warp_mma_k % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index(warp_mma_k % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_frag_A); + this->warp_tile_iterator_B_.load(warp_frag_B); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + warp_mma(accum, warp_frag_A, warp_frag_B, accum); + } + + // Add negative offsets to return smem load iterators to the 'start' of the shared memory + this->warp_tile_iterator_A_.add_tile_offset({0, -Policy::kPartitionsK * Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset({-Policy::kPartitionsK * Base::kWarpGemmIterations, 0}); + + __syncthreads(); + + iterator_A.load(tb_frag_A); + iterator_B.load(tb_frag_B); + + ++iterator_A; + ++iterator_B; + + // Avoid reading out of bounds if this was the last loop iteration + iterator_A.clear_mask(gemm_k_iterations <= 2); + iterator_B.clear_mask(gemm_k_iterations <= 2); + } + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/mma_softmax_mainloop_fusion_multistage.h b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_softmax_mainloop_fusion_multistage.h new file mode 100644 index 0000000..52a32c9 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_softmax_mainloop_fusion_multistage.h @@ -0,0 +1,756 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. + + It loads two loop invariant vectors, norm and sum, in the prologue and + stores them in the register file. We will call elementwise operation to + apply norm and sum between ldmatrix and warp mma. +*/ + +#pragma once + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" +#include "mctlass/transform/threadblock/predicated_scale_bias_vector_iterator.h" +#include "mctlass/gemm/threadblock/mma_base.h" +#include "mctlass/gemm/warp/softmax_scale_bias_transform.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class MmaMainloopFusionBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Policy describing tuning details + using Policy = Policy_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = mctlass::gemm::GemmShape; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = + (WarpGemm::kK / Operator::Policy::MmaShape::kK); + + /// Number of stages + static int const kStages = Stages; + + /// Tensor reference to the A operand + using TensorRefA = TensorRef; + + /// Tensor reference to the B operand + using TensorRefB = TensorRef; + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + class SharedStorage { + public: + // + // Type definitions + // + + /// Shape of the A matrix operand in shared memory + using ShapeA = MatrixShape; + + /// Shape of the B matrix operand in shared memory + using ShapeB = + MatrixShape; + + public: + // + // Data members + // + + /// Buffer for A operand + AlignedBuffer operand_A; + + /// Buffer for B operand + AlignedBuffer operand_B; + + public: + + // + // Methods + // + + /// Returns a layout object for the A matrix + MCTLASS_DEVICE + static typename Operator::LayoutA LayoutA() { + return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn}); + } + + /// Returns a layout object for the B matrix + MCTLASS_HOST_DEVICE + static typename Operator::LayoutB LayoutB() { + return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); + } + + /// Returns a TensorRef to the A operand + MCTLASS_HOST_DEVICE + TensorRefA operand_A_ref() { + return TensorRefA{operand_A.data(), LayoutA()}; + } + + /// Returns a TensorRef to the B operand + MCTLASS_HOST_DEVICE + TensorRefB operand_B_ref() { + return TensorRefB{operand_B.data(), LayoutB()}; + } + }; + + protected: + + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A operand from shared memory + typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaMainloopFusionBase( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx) + : warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx), + warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx) {} +}; + + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// Iterates over vectors of var and mean vector in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorNormSum_, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Whether problem has been transformed. This determines to which operand + /// the softmax is applied. + bool InternalTranspose, + /// Use zfill or predicate for out-of-bound cp.async + SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone, + /// Used for partial specialization + typename Enable = bool> +class MmaSoftmaxMainloopFusionMultistage : + public MmaMainloopFusionBase { +public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Iterates over tiles of the var and mean vectors in global memory + using IteratorNormSum = IteratorNormSum_; + ///< Policy describing tuning details + using Policy = Policy_; + + ///< Base class + using Base = MmaMainloopFusionBase; + + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + /// Internal structure exposed for introspection. + struct Detail { + + static_assert(Base::kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + }; + + private: + + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + using WarpLoadedFragmentNormSum = typename IteratorNormSum::Fragment; + + static bool const kInternalTranspose = InternalTranspose; + + using SoftmaxFragment = typename platform::conditional::type; + + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + int warp_idx_m_; + + int warp_idx_n_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaSoftmaxMainloopFusionMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + warp_idx_m_ = warp_idx_mn % Base::WarpCount::kM; + warp_idx_n_ = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m_, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n_}); + } + + MCTLASS_DEVICE + void copy_tiles_and_advance(IteratorA &iterator_A, + IteratorB &iterator_B, + int group_start_A = 0, int group_start_B = 0) { + iterator_A.set_iteration_index(group_start_A * + IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_A.get(); + + if (SharedMemoryClear == SharedMemoryClearOption::kZfill) { + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + } else { + mctlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + } + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + } + + iterator_B.set_iteration_index(group_start_B * + IteratorB::kAccessesPerVector); + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_B.get(); + + if (SharedMemoryClear == SharedMemoryClearOption::kZfill) { + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + } else { + mctlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + } + + ++iterator_B; + } + ++this->smem_iterator_B_; + } + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< iterator over B operand in global memory + IteratorNormSum iterator_norm_sum, + ///< initial value of accumulator + FragmentC const &src_accum) { + + // + // Prologue + // + // Issue several complete stages + + WarpLoadedFragmentNormSum warp_loaded_frag_norm_sum; + iterator_norm_sum.add_tile_offset({0, warp_idx_m_}); + iterator_norm_sum.load(warp_loaded_frag_norm_sum); + + MCTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; + ++stage, --gemm_k_iterations) { + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + int src_bytes = (iterator_A.valid() ? kSrcBytes : 0); + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + + iterator_B.set_iteration_index(0); + this->smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + } + + ++this->smem_iterator_B_; + } + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Defines the boundary of a stage of cp.async. + mctlass::arch::cp_async_fence(); + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // Waits until kStages-2 stages have committed. + mctlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[2]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpTransformedFragmentA warp_transformed_frag_A[2]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + mctlass::gemm::warp::SoftmaxScaleBiasTransform< + SoftmaxFragment, WarpLoadedFragmentNormSum> elementwise_transform; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + // Start issuing the first group of the next stage outside of the mainloop + copy_tiles_and_advance(iterator_A, iterator_B); + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0], + warp_loaded_frag_A[0], warp_loaded_frag_B[0]); + + if (kInternalTranspose) { + elementwise_transform(warp_transformed_frag_B[0], + warp_loaded_frag_norm_sum); + } else { + elementwise_transform(warp_transformed_frag_A[0], + warp_loaded_frag_norm_sum); + } + + // + // Mainloop + // + + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) { + warp_mma.transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % 2], + warp_loaded_frag_B[warp_mma_k % 2]); + + if (kInternalTranspose) { + elementwise_transform(warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_norm_sum); + } else { + elementwise_transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_loaded_frag_norm_sum); + } + } + + // Issue global->shared copies for the next stage + int group_start_iteration_A, group_start_iteration_B; + + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + group_start_iteration_A = 0; + group_start_iteration_B = 0; + } else { + group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + } + + copy_tiles_and_advance(iterator_A, iterator_B, + group_start_iteration_A, + group_start_iteration_B); + + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum + ); + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + + // Inserts a memory fence between stages of cp.async instructions. + mctlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages have committed. + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + } + + // Do any conversions feeding the first stage at the end of the loop so + // we can start right away on mma instructions + if (warp_mma_k + 1 == Base::kWarpGemmIterations) { + warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + if (kInternalTranspose) { + elementwise_transform(warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_norm_sum); + } else { + elementwise_transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_norm_sum); + } + } + } + + } + + if (SharedMemoryClear == SharedMemoryClearOption::kZfill) { + // commit and drain all pending and predicated LDGSTS pnz from the GEMM mainloop + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + } + + // Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/mma_sparse_base.h b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_sparse_base.h new file mode 100644 index 0000000..624561f --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_sparse_base.h @@ -0,0 +1,273 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Policy object describing MmaTensorOp +template < + /// Warp-level GEMM operator (concept: gemm::warp::Mma) + typename Operator_, + /// Padding used for A operand in shared memory (concept: MatrixShape) + typename SmemPaddingA_, + /// Padding used for B operand in shared memory (concept: MatrixShape) + typename SmemPaddingB_, + /// Padding used for E operand in shared memory (concept: MatrixShape) + typename SmemPaddingE_, + /// Number of partitions of K dimension of GEMM + int PartitionsK = 1> +struct SparseMmaPolicy { + /// Warp-level GEMM operator (concept: gemm::warp::MmaTensorOp or gemm::warp::MmaSimt) + using Operator = Operator_; + + /// Padding used for A operand in shared memory + using SmemPaddingA = SmemPaddingA_; + + /// Padding used for B operand in shared memory + using SmemPaddingB = SmemPaddingB_; + + /// Padding used for B operand in shared memory + using SmemPaddingE = SmemPaddingE_; + + /// Number of partitions of K dimension + static int const kPartitionsK = PartitionsK; +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class SparseMmaBase { + public: + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + + ///< Policy describing tuning details + using Policy = Policy_; + + // + // Dependent types + // + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// Shape describing the overall GEMM computed from shared memory + /// by each warp. + using WarpGemm = typename Policy::Operator::Shape; + + /// Shape describing the number of warps filling the CTA + using WarpCount = GemmShape; + + /// Number of warp-level GEMM oeprations + static int const kWarpGemmIterations = + (WarpGemm::kK / Operator::Policy::MmaShape::kK); + + static_assert(kWarpGemmIterations > 1, + "The pipelined structure requires at least two warp-level " + "GEMM operations."); + + static_assert((kWarpGemmIterations % 2) == 0, + "Inner loop iteration must be an even number."); + + /// Number of stages + static int const kStages = Stages; + + static int const kSparse = Operator::kSparse; + + static int const kElementsPerElementE = Operator::kElementsPerElementE; + + /// Tensor reference to the A operand + using TensorRefA = TensorRef; + + /// Tensor reference to the B operand + using TensorRefB = TensorRef; + + /// Tensor reference to the E operand + using TensorRefE = TensorRef; + + // + // Nested structs + // + + /// Shared storage object needed by threadblock-scoped GEMM + class SharedStorage { + public: + // + // Type definitions + // + + /// Shape of the A matrix operand in shared memory + using ShapeA = MatrixShape; + + /// Shape of the B matrix operand in shared memory + using ShapeB = + MatrixShape; + + /// Shape of the E matrix operand in shared memory + using ShapeE = + MatrixShape; + + public: + // + // Data members + // + + /// Buffer for A operand + AlignedBuffer operand_A; + + /// Buffer for B operand + AlignedBuffer operand_B; + + /// Buffer for E operand + AlignedBuffer operand_E; + + public: + + // + // Methods + // + + /// Returns a layout object for the A matrix + MCTLASS_DEVICE + static typename Operator::LayoutA LayoutA() { + return Operator::LayoutA::packed({ShapeA::kRow, ShapeA::kColumn}); + } + + /// Returns a layout object for the B matrix + MCTLASS_HOST_DEVICE + static typename Operator::LayoutB LayoutB() { + return Operator::LayoutB::packed({ShapeB::kRow, ShapeB::kColumn}); + } + + /// Returns a layout object for the E matrix + MCTLASS_HOST_DEVICE + static typename Operator::LayoutE LayoutE() { + return Operator::LayoutE::packed({ShapeE::kRow, ShapeE::kColumn}); + } + + /// Returns a TensorRef to the A operand + MCTLASS_HOST_DEVICE + TensorRefA operand_A_ref() { + return TensorRefA{operand_A.data(), LayoutA()}; + } + + /// Returns a TensorRef to the B operand + MCTLASS_HOST_DEVICE + TensorRefB operand_B_ref() { + return TensorRefB{operand_B.data(), LayoutB()}; + } + + /// Returns a TensorRef to the E operand + MCTLASS_HOST_DEVICE + TensorRefE operand_E_ref() { + return TensorRefE{operand_E.data(), LayoutE()}; + } + }; + + protected: + + // + // Data members + // + + /// Iterator to load a warp-scoped tile of A operand from shared memory + typename Operator::IteratorA warp_tile_iterator_A_; + + /// Iterator to load a warp-scoped tile of B operand from shared memory + typename Operator::IteratorB warp_tile_iterator_B_; + + /// Iterator to load a warp-scoped tile of E operand from shared memory + typename Operator::IteratorE warp_tile_iterator_E_; + + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + SparseMmaBase( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + warp_tile_iterator_A_(shared_storage.operand_A_ref(), lane_idx), + warp_tile_iterator_B_(shared_storage.operand_B_ref(), lane_idx), + warp_tile_iterator_E_(shared_storage.operand_E_ref(), lane_idx) { + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/mma_sparse_multistage.h b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_sparse_multistage.h new file mode 100644 index 0000000..904240e --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_sparse_multistage.h @@ -0,0 +1,668 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" + +#include "mctlass/gemm/threadblock/mma_sparse_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Iterates over tiles of E operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorE_, + /// Iterates over tiles of E operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorE_, + /// Cache operation for operand E + mctlass::arch::CacheOperation::Kind CacheOpE, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Used for partial specialization + typename Enable = bool> +class SparseMmaMultistage : + public SparseMmaBase { +public: + ///< Base class + using Base = SparseMmaBase; + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Iterates over tiles of E operand in global memory + using IteratorE = IteratorE_; + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + ///< Policy describing tuning details + using Policy = Policy_; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + using SmemIteratorE = SmemIteratorE_; + + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + static mctlass::arch::CacheOperation::Kind const kCacheOpE = CacheOpE; + + static int const kSparse = Policy::Operator::kSparse; + static int const kMetaSizeInBits = Policy::Operator::kMetaSizeInBits; + static int const kMaxID2 = Policy::Operator::kMaxID2; + static int const kElementsPerElementE = + Policy::Operator::kElementsPerElementE; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + /// ElementE + using ElementE = typename IteratorE::Element; + + /// LayoutE + using LayoutE = typename IteratorE::Layout; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + /// Internal structure exposed for introspection. + struct Detail { + + /// Number of async copies to load one stage of operand A + static int const TBLoadIterationsA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of async copies to load one stage of operand B + static int const TBLoadIterationsB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of async copies to load one stage of operand E + static int const TBLoadIterationsE = + IteratorE::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of async copies to load one group of operand A + static int const kAccessesPerGroupA = + (TBLoadIterationsA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// Number of async copies to load one group of operand B + static int const kAccessesPerGroupB = + (TBLoadIterationsB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// Number of async copies to load one group of operand E + static int const kAccessesPerGroupE = + (TBLoadIterationsE + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// E operand is tiny. For the most of time, not all the warps are needed + /// to load it from the global memory. + static int const kValidWarps = IteratorE::ThreadMap::kThreads / 32; + + /// B operand is twice as big as A which brings very high register pressure. + /// We have to sacrifice the double buffer when the warp tile size is big. + static int const kBBufferSize = + ((sizeof(typename Operator::ElementC) == 4) && + ((platform::is_same::value && + platform::is_same::value)) && + (Operator::Shape::kM >= 64 && Operator::Shape::kN >= 64)) + ? 1 + : 2; + }; + + private: + + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + using WarpFragmentE = typename Operator::FragmentE; + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + + /// Iterator to write threadblock-scoped tile of E operand to shared memory + SmemIteratorE smem_iterator_E_; + + /// Warp id + bool is_warp_valid_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + SparseMmaMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx), + smem_iterator_E_(shared_storage.operand_E_ref(), thread_idx) + { + is_warp_valid_ = warp_idx < Detail::kValidWarps; + + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + this->warp_tile_iterator_E_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + } + + MCTLASS_DEVICE + void copy_tiles_and_advance(IteratorA &iterator_A, IteratorB &iterator_B, + IteratorE &iterator_E, int group_start_A = 0, + int group_start_B = 0, int group_start_E = 0) { + iterator_A.set_iteration_index(group_start_A * + IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // async copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + if (group_start_A + j < Detail::TBLoadIterationsA) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_A.get(); + + mctlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + } + + iterator_B.set_iteration_index(group_start_B * + IteratorB::kAccessesPerVector); + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // async copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::TBLoadIterationsB) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_B.get(); + + mctlass::arch::cp_async( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + + ++iterator_B; + } + ++this->smem_iterator_B_; + } + } + + iterator_E.set_iteration_index(group_start_E); + this->smem_iterator_E_.set_iteration_index(group_start_E); + + // async copy for operand E + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupE; ++j) { + if (group_start_E + j < Detail::TBLoadIterationsE) { + typename IteratorE::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_E_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorE::ThreadMap::kElementsPerAccess / 8; + + auto gmem_ptr = iterator_E.get(); + + mctlass::arch::cp_async( + dst_ptr, gmem_ptr, iterator_E.valid() && is_warp_valid_); + + ++iterator_E; + ++this->smem_iterator_E_; + } + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< iterator over E operand in global memory + IteratorE iterator_E, + ///< initial value of accumulator + FragmentC const &src_accum) { + + // + // Prologue + // + + // Issue several complete stages + MCTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; + ++stage, --gemm_k_iterations) { + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + iterator_E.clear_mask(gemm_k_iterations == 0); + + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_index(0); + + // async copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::TBLoadIterationsA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + + iterator_B.set_iteration_index(0); + this->smem_iterator_B_.set_iteration_index(0); + + // async copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::TBLoadIterationsB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + } + + ++this->smem_iterator_B_; + } + + iterator_E.set_iteration_index(0); + this->smem_iterator_E_.set_iteration_index(0); + + // async copy for operand E + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::TBLoadIterationsE; ++j) { + typename IteratorE::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_E_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorE::ThreadMap::kElementsPerAccess / 8; + if (is_warp_valid_) + mctlass::arch::cp_async_zfill( + dst_ptr, iterator_E.get(), iterator_E.valid()); + + ++iterator_E; + + ++this->smem_iterator_E_; + } + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + iterator_E.add_tile_offset({0, 1}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + this->smem_iterator_E_.add_tile_offset({0, 1}); + + // cp.async.commit_group - completes a stage + mctlass::arch::cp_async_fence(); + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + mctlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[2]; + WarpLoadedFragmentB warp_loaded_frag_B[Detail::kBBufferSize]; + WarpTransformedFragmentA warp_transformed_frag_A[2]; + WarpTransformedFragmentB warp_transformed_frag_B[Detail::kBBufferSize]; + WarpFragmentE warp_frag_E[2]; + + Operator warp_mma; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + this->warp_tile_iterator_E_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + this->warp_tile_iterator_E_.load(warp_frag_E[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + ++this->warp_tile_iterator_E_; + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + iterator_E.clear_mask(gemm_k_iterations == 0); + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0], + warp_loaded_frag_A[0], warp_loaded_frag_B[0]); + + // + // Mainloop + // + + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_E_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_E_.load(warp_frag_E[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_E_; + + if (Detail::kBBufferSize == 2) { + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.load( + warp_loaded_frag_B[(warp_mma_k + 1) % Detail::kBBufferSize]); + ++this->warp_tile_iterator_B_; + } + + if (warp_mma_k > 0) + warp_mma.transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % Detail::kBBufferSize], + warp_loaded_frag_A[warp_mma_k % 2], + warp_loaded_frag_B[warp_mma_k % Detail::kBBufferSize]); + + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % Detail::kBBufferSize], accum, + warp_frag_E[warp_mma_k % 2] + ); + + if (Detail::kBBufferSize == 1) { + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + ++this->warp_tile_iterator_B_; + + } + + // Issue global->shared copies for the this stage + if (warp_mma_k < Base::kWarpGemmIterations - 1) { + int group_start_iteration_A, group_start_iteration_B, group_start_iteration_E; + + group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA; + group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB; + group_start_iteration_E = warp_mma_k * Detail::kAccessesPerGroupE; + + copy_tiles_and_advance( + iterator_A, iterator_B, iterator_E, group_start_iteration_A, + group_start_iteration_B, group_start_iteration_E); + } + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + int group_start_iteration_A, group_start_iteration_B, group_start_iteration_E; + group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + group_start_iteration_E = + (warp_mma_k + 1) * Detail::kAccessesPerGroupE; + + copy_tiles_and_advance( + iterator_A, iterator_B, iterator_E, group_start_iteration_A, + group_start_iteration_B, group_start_iteration_E); + + // Inserts a memory fence between stages of cp.async instructions. + mctlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages have committed. + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + iterator_E.add_tile_offset({0, 1}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + this->smem_iterator_E_.add_tile_offset({0, 1}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + this->smem_iterator_E_.add_tile_offset({0, -Base::kStages}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + this->warp_tile_iterator_E_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + iterator_E.clear_mask(gemm_k_iterations == 0); + } + + // Do any conversions feeding the first stage at the end of the loop so + // we can start right away on mma instructions + if (warp_mma_k + 1 == Base::kWarpGemmIterations) + warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + } + + } + + + // Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/mma_with_reduction_multistage.h b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_with_reduction_multistage.h new file mode 100644 index 0000000..d23c7e5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/mma_with_reduction_multistage.h @@ -0,0 +1,539 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template for a double-buffered threadblock-scoped GEMM kernel. +*/ + +#pragma once + +#include "mctlass/aligned_buffer.h" +#include "mctlass/arch/memory.h" +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_types.h" + +#include "mctlass/gemm/threadblock/mma_base.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math +/// instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Iterates over tiles of A operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorA_, + /// Iterates over tiles of A operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorA_, + /// Cache operation for operand A + mctlass::arch::CacheOperation::Kind CacheOpA, + /// Iterates over tiles of B operand in global memory + // (concept: ReadableTileIterator | ForwardTileIterator | + // MaskedTileIterator) + typename IteratorB_, + /// Iterates over tiles of B operand in shared memory + /// (concept: WriteableTileIterator | RandomAccessTileIterator) + typename SmemIteratorB_, + /// Cache operation for operand B + mctlass::arch::CacheOperation::Kind CacheOpB, + /// Data type of accumulator matrix + typename ElementC_, + /// Data type of accumulator matrix + typename LayoutC_, + /// Policy describing tuning details (concept: MmaPolicy) + typename Policy_, + /// Number of stages, + int Stages, + /// Use zfill or predicate for out-of-bound cp.async + /// Used for partial specialization + typename Enable = bool> +class MmaWithReductionMultistage : + public MmaBase { + ///< Base class + using Base = MmaBase; + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + using Shape = Shape_; + ///< Iterates over tiles of A operand in global memory + using IteratorA = IteratorA_; + ///< Iterates over tiles of B operand in global memory + using IteratorB = IteratorB_; + ///< Data type of accumulator matrix + using ElementC = ElementC_; + ///< Layout of accumulator matrix + using LayoutC = LayoutC_; + ///< Policy describing tuning details + using Policy = Policy_; + + using SmemIteratorA = SmemIteratorA_; + using SmemIteratorB = SmemIteratorB_; + + static mctlass::arch::CacheOperation::Kind const kCacheOpA = CacheOpA; + static mctlass::arch::CacheOperation::Kind const kCacheOpB = CacheOpB; + + // + // Dependent types + // + + /// Fragment of accumulator tile + using FragmentC = typename Policy::Operator::FragmentC; + + /// Warp-level Mma + using Operator = typename Policy::Operator; + + using FragmentReduction = typename Operator::FragmentReduction; + + /// Minimum architecture is Sm80 to support cp.async + using ArchTag = arch::Sm80; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Operator::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Operator::kTransformB; + + static int const kReduceKForA = Operator::kReduceKForA; + + /// Internal structure exposed for introspection. + struct Detail { + + /// Number of cp.async instructions to load one stage of operand A + static int const AsyncCopyIterationsPerStageA = + IteratorA::ThreadMap::Iterations::kCount; + + /// Number of cp.async instructions to load one stage of operand B + static int const AsyncCopyIterationsPerStageB = + IteratorB::ThreadMap::Iterations::kCount; + + /// Number of stages + static int const kStages = Stages; + + /// Number of cp.async instructions to load on group of operand A + static int const kAccessesPerGroupA = + (AsyncCopyIterationsPerStageA + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + + /// Number of cp.async instructions to load on group of operand B + static int const kAccessesPerGroupB = + (AsyncCopyIterationsPerStageB + Base::kWarpGemmIterations - 1) / Base::kWarpGemmIterations; + }; + + private: + + using WarpLoadedFragmentA = typename Operator::FragmentA; + using WarpLoadedFragmentB = typename Operator::FragmentB; + using WarpTransformedFragmentA = typename Operator::TransformedFragmentA; + using WarpTransformedFragmentB = typename Operator::TransformedFragmentB; + + private: + + // + // Data members + // + + /// Iterator to write threadblock-scoped tile of A operand to shared memory + SmemIteratorA smem_iterator_A_; + + /// Iterator to write threadblock-scoped tile of B operand to shared memory + SmemIteratorB smem_iterator_B_; + +public: + + /// Construct from tensor references + MCTLASS_DEVICE + MmaWithReductionMultistage( + ///< Shared storage needed for internal use by threadblock-scoped GEMM + typename Base::SharedStorage &shared_storage, + ///< ID within the threadblock + int thread_idx, + ///< ID of warp + int warp_idx, + ///< ID of each thread within a warp + int lane_idx + ): + Base(shared_storage, thread_idx, warp_idx, lane_idx), + smem_iterator_A_(shared_storage.operand_A_ref(), thread_idx), + smem_iterator_B_(shared_storage.operand_B_ref(), thread_idx) + { + // Compute warp location within threadblock tile by mapping the warp_id to + // three coordinates: + // _m: the warp's position within the threadblock along the M dimension + // _n: the warp's position within the threadblock along the N dimension + // _k: the warp's position within the threadblock along the K dimension + + int warp_idx_mn = warp_idx % (Base::WarpCount::kM * Base::WarpCount::kN); + int warp_idx_k = warp_idx / (Base::WarpCount::kM * Base::WarpCount::kN); + + int warp_idx_m = warp_idx_mn % Base::WarpCount::kM; + int warp_idx_n = warp_idx_mn / Base::WarpCount::kM; + + // Add per-warp offsets in units of warp-level tiles + this->warp_tile_iterator_A_.add_tile_offset( + {warp_idx_m, Base::kWarpGemmIterations * warp_idx_k}); + this->warp_tile_iterator_B_.add_tile_offset( + {Base::kWarpGemmIterations * warp_idx_k, warp_idx_n}); + } + + MCTLASS_DEVICE + void copy_tiles_and_advance(IteratorA &iterator_A, IteratorB &iterator_B, + int group_start_A = 0, int group_start_B = 0) { + iterator_A.set_iteration_index(group_start_A * + IteratorA::kAccessesPerVector); + this->smem_iterator_A_.set_iteration_index(group_start_A); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupA; ++j) { + if (group_start_A + j < Detail::AsyncCopyIterationsPerStageA) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_A.get(); + + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_A.valid()); + } else { + mctlass::arch::cp_async( + } + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + } + + iterator_B.set_iteration_index(group_start_B * + IteratorB::kAccessesPerVector); + this->smem_iterator_B_.set_iteration_index(group_start_B); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::kAccessesPerGroupB; ++j) { + if (group_start_B + j < Detail::AsyncCopyIterationsPerStageB) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + int const kSrcBytes = sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + auto gmem_ptr = iterator_B.get(); + + mctlass::arch::cp_async_zfill( + dst_ptr + v, gmem_ptr, iterator_B.valid()); + } else { + mctlass::arch::cp_async( + } + + ++iterator_B; + } + ++this->smem_iterator_B_; + } + } + } + + /// Perform a threadblock-scoped matrix multiply-accumulate + MCTLASS_DEVICE + void operator()( + ///< problem size of GEMM + int gemm_k_iterations, + ///< destination accumulator tile + FragmentC &accum, + ///< iterator over A operand in global memory + IteratorA iterator_A, + ///< iterator over B operand in global memory + IteratorB iterator_B, + ///< initial value of accumulator + FragmentC const &src_accum, + FragmentReduction &gemm_k_reduction_accum) { + + // + // Prologue + // + // Issue several complete stages + + MCTLASS_PRAGMA_UNROLL + for (int stage = 0; stage < Base::kStages - 1; + ++stage, --gemm_k_iterations) { + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + iterator_A.set_iteration_index(0); + this->smem_iterator_A_.set_iteration_index(0); + + // Async Copy for operand A + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageA; ++j) { + typename IteratorA::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_A_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorA::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorA::ThreadMap::kElementsPerAccess / + IteratorA::kAccessesPerVector / 8; + + int src_bytes = (iterator_A.valid() ? kSrcBytes : 0); + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_A.get(), iterator_A.valid()); + + ++iterator_A; + } + + ++this->smem_iterator_A_; + } + + iterator_B.set_iteration_index(0); + this->smem_iterator_B_.set_iteration_index(0); + + // Async Copy for operand B + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < Detail::AsyncCopyIterationsPerStageB; ++j) { + typename IteratorB::AccessType *dst_ptr = + reinterpret_cast( + this->smem_iterator_B_.get()); + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < IteratorB::kAccessesPerVector; ++v) { + int const kSrcBytes = + sizeof_bits::value * + IteratorB::ThreadMap::kElementsPerAccess / + IteratorB::kAccessesPerVector / 8; + + mctlass::arch::cp_async_zfill( + dst_ptr + v, iterator_B.get(), iterator_B.valid()); + + ++iterator_B; + } + + ++this->smem_iterator_B_; + } + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Defines the boundary of a stage of cp.async. + mctlass::arch::cp_async_fence(); + } + + // Perform accumulation in the 'd' output operand + accum = src_accum; + + // Waits until kStages-2 stages have committed. + mctlass::arch::cp_async_wait(); + __syncthreads(); + + // Pair of fragments used to overlap shared memory loads and math + // instructions + WarpLoadedFragmentA warp_loaded_frag_A[2]; + WarpLoadedFragmentB warp_loaded_frag_B[2]; + WarpTransformedFragmentA warp_transformed_frag_A[2]; + WarpTransformedFragmentB warp_transformed_frag_B[2]; + + Operator warp_mma; + + this->warp_tile_iterator_A_.set_kgroup_index(0); + this->warp_tile_iterator_B_.set_kgroup_index(0); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[0]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[0]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + + int smem_write_stage_idx = Base::kStages - 1; + int smem_read_stage_idx = 0; + + warp_mma.transform(warp_transformed_frag_A[0], warp_transformed_frag_B[0], + warp_loaded_frag_A[0], warp_loaded_frag_B[0]); + + // + // Mainloop + // + + MCTLASS_GEMM_LOOP + for (; gemm_k_iterations > (-Base::kStages + 1);) { + // + // Loop over GEMM K dimension + // + + // Computes a warp-level GEMM on data held in shared memory + // Each "warp_mma_k" refers to a warp-level matrix multiply-accumulate + MCTLASS_PRAGMA_UNROLL + for (int warp_mma_k = 0; warp_mma_k < Base::kWarpGemmIterations; + ++warp_mma_k) { + + // Load warp-level tiles from shared memory, wrapping to k offset if + // this is the last group as the case may be. + + this->warp_tile_iterator_A_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + this->warp_tile_iterator_B_.set_kgroup_index((warp_mma_k + 1) % Base::kWarpGemmIterations); + + this->warp_tile_iterator_A_.load(warp_loaded_frag_A[(warp_mma_k + 1) % 2]); + this->warp_tile_iterator_B_.load(warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + + ++this->warp_tile_iterator_A_; + ++this->warp_tile_iterator_B_; + + if (warp_mma_k > 0) + warp_mma.transform(warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + warp_loaded_frag_A[warp_mma_k % 2], + warp_loaded_frag_B[warp_mma_k % 2]); + + warp_mma( + accum, + warp_transformed_frag_A[warp_mma_k % 2], + warp_transformed_frag_B[warp_mma_k % 2], + accum, + gemm_k_reduction_accum + ); + + // Issue global->shared copies for the this stage + if (warp_mma_k < Base::kWarpGemmIterations - 1) { + int group_start_iteration_A, group_start_iteration_B; + + group_start_iteration_A = warp_mma_k * Detail::kAccessesPerGroupA; + group_start_iteration_B = warp_mma_k * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, + group_start_iteration_B); + } + + if (warp_mma_k + 2 == Base::kWarpGemmIterations) { + int group_start_iteration_A, group_start_iteration_B; + group_start_iteration_A = + (warp_mma_k + 1) * Detail::kAccessesPerGroupA; + group_start_iteration_B = + (warp_mma_k + 1) * Detail::kAccessesPerGroupB; + + copy_tiles_and_advance(iterator_A, iterator_B, group_start_iteration_A, + group_start_iteration_B); + + // Inserts a memory fence between stages of cp.async instructions. + mctlass::arch::cp_async_fence(); + + // Waits until kStages-2 stages have committed. + arch::cp_async_wait(); + __syncthreads(); + + // Move to the next stage + iterator_A.add_tile_offset({0, 1}); + iterator_B.add_tile_offset({1, 0}); + + this->smem_iterator_A_.add_tile_offset({0, 1}); + this->smem_iterator_B_.add_tile_offset({1, 0}); + + // Add negative offsets to return iterators to the 'start' of the + // circular buffer in shared memory + if (smem_write_stage_idx == (Base::kStages - 1)) { + this->smem_iterator_A_.add_tile_offset({0, -Base::kStages}); + this->smem_iterator_B_.add_tile_offset({-Base::kStages, 0}); + smem_write_stage_idx = 0; + } else { + ++smem_write_stage_idx; + } + + if (smem_read_stage_idx == (Base::kStages - 1)) { + this->warp_tile_iterator_A_.add_tile_offset( + {0, -Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations}); + this->warp_tile_iterator_B_.add_tile_offset( + {-Base::kStages * Policy::kPartitionsK * + Base::kWarpGemmIterations, + 0}); + smem_read_stage_idx = 0; + } else { + ++smem_read_stage_idx; + } + + --gemm_k_iterations; + iterator_A.clear_mask(gemm_k_iterations == 0); + iterator_B.clear_mask(gemm_k_iterations == 0); + } + + // Do any conversions feeding the first stage at the end of the loop so + // we can start right away on mma instructions + if (warp_mma_k + 1 == Base::kWarpGemmIterations) + warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2], + warp_transformed_frag_B[(warp_mma_k + 1) % 2], + warp_loaded_frag_A[(warp_mma_k + 1) % 2], + warp_loaded_frag_B[(warp_mma_k + 1) % 2]); + } + + } + + // commit and drain all pending and predicated LDGSTS pnz from the GEMM mainloop + mctlass::arch::cp_async_fence(); + mctlass::arch::cp_async_wait<0>(); + __syncthreads(); + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/threadblock_swizzle.h b/csrc/mctlass/include/mctlass/gemm/threadblock/threadblock_swizzle.h new file mode 100644 index 0000000..9122b39 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/threadblock_swizzle.h @@ -0,0 +1,459 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Implements several possible threadblock-swizzling functions mapping blockIdx to + GEMM problems. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/platform/platform.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/gemm/threadblock/index_remat.h" +#include "mctlass/gemm/threadblock/threadblock_swizzle_streamk.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Threadblock swizzling function for GEMMs +template +struct GemmIdentityThreadblockSwizzle { + + MCTLASS_HOST_DEVICE + GemmIdentityThreadblockSwizzle() { } + + /// Returns the shape of the problem in units of logical tiles + /// *Gemm* problem size: gemm(M, N, K) + MCTLASS_HOST_DEVICE + GemmCoord get_tiled_shape( + GemmCoord problem_size, + GemmCoord tile_size, + int split_k_slices) const { + + return GemmCoord( + (problem_size.m() + tile_size.m() - 1) / tile_size.m(), + (problem_size.n() + tile_size.n() - 1) / tile_size.n(), + split_k_slices); + } + + /// Returns the shape of the problem in units of logical tiles + /// *ImplicitGemm* Conv2d problem size: conv_operator(NPQK, NHWC, KRSC) + MCTLASS_HOST_DEVICE + GemmCoord get_tiled_shape( + mctlass::conv::Operator conv_operator, + mctlass::conv::Conv2dProblemSize const &problem_size, + GemmCoord tile_size, + int split_k_slices) const { + + gemm::GemmCoord implicit_gemm_problem_size = + mctlass::conv::implicit_gemm_problem_size(conv_operator, problem_size); + + return get_tiled_shape( + implicit_gemm_problem_size, tile_size, split_k_slices); + } + + /// Returns the shape of the problem in units of logical tiles + /// *ImplicitGemm* Conv3d problem size: conv_operator(NZPQK, NDHWC, KTRSC) + MCTLASS_HOST_DEVICE + GemmCoord get_tiled_shape( + mctlass::conv::Operator conv_operator, + mctlass::conv::Conv3dProblemSize const &problem_size, + GemmCoord tile_size, + int split_k_slices) const { + + gemm::GemmCoord implicit_gemm_problem_size = + mctlass::conv::implicit_gemm_problem_size(conv_operator, problem_size); + + return get_tiled_shape( + implicit_gemm_problem_size, tile_size, split_k_slices); + } + + /// Computes CUDA grid dimensions given a size in units of logical tiles + MCTLASS_HOST_DEVICE + dim3 get_grid_shape(GemmCoord tiled_shape) const { + int tile = 1 << get_log_tile(tiled_shape); + return dim3(tiled_shape.m() * tile, (tiled_shape.n() + tile - 1) / tile, tiled_shape.k()); + } + + /// Calculates optimal swizzle width + MCTLASS_HOST_DEVICE + int get_log_tile(GemmCoord tiled_shape) const { + auto n = tiled_shape.n(); + // Thresholds picked so that it doesn't cause too many no-op CTAs + if (N >= 8 && n >= 6) + return 3; + else if (N >= 4 && n >= 3) + return 2; + else if (N >= 2 && n >= 2) + return 1; + else + return 0; + } + + /// Obtains the threadblock offset (in units of threadblock-scoped tiles) + MCTLASS_DEVICE + GemmCoord get_tile_offset(int log_tile) const { + int block_idx_x = RematerializeBlockIdxX(); + int block_idx_y = RematerializeBlockIdxY(); + int block_idx_z = RematerializeBlockIdxZ(); + + return GemmCoord{(block_idx_x >> log_tile), // + (block_idx_y << log_tile) + ((block_idx_x) & ((1 << (log_tile)) - 1)), + block_idx_z}; + } + + /// Obtains the threadblock offset (in units of threadblock-scoped tiles) + MCTLASS_DEVICE + GemmCoord get_tile_offset(GemmCoord tiled_shape) const { + + int const kTile = N; + int block_idx_x = RematerializeBlockIdxX(); + int block_idx_y = RematerializeBlockIdxY(); + + if ((tiled_shape.m() < kTile) || (tiled_shape.n() < kTile)) + return GemmCoord{block_idx_x, block_idx_y, RematerializeBlockIdxZ()}; + + return GemmCoord{ + (block_idx_x / kTile), + (block_idx_y * kTile) + (block_idx_x % kTile), + RematerializeBlockIdxZ() + }; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Threadblock swizzling function for GEMMs +struct GemmHorizontalThreadblockSwizzle { + + MCTLASS_HOST_DEVICE + GemmHorizontalThreadblockSwizzle() { } + + /// Returns the shape of the problem in units of logical tiles + MCTLASS_HOST_DEVICE + GemmCoord get_tiled_shape( + GemmCoord problem_size, + GemmCoord tile_size, + int split_k_slices) const { + + return GemmCoord( + (problem_size.m() + tile_size.m() - 1) / tile_size.m(), + (problem_size.n() + tile_size.n() - 1) / tile_size.n(), + split_k_slices); + } + + /// Computes CUDA grid dimensions given a size in units of logical tiles + MCTLASS_HOST_DEVICE + dim3 get_grid_shape(GemmCoord tiled_shape) const { + return dim3(tiled_shape.n(), tiled_shape.m(), tiled_shape.k()); + } + + /// Calculates optimal swizzle width + MCTLASS_HOST_DEVICE + int get_log_tile(GemmCoord tiled_shape) const { + return 0; + } + + /// Obtains the threadblock offset (in units of threadblock-scoped tiles) + MCTLASS_DEVICE + GemmCoord get_tile_offset(GemmCoord tiled_shape) const { + return GemmCoord{ + RematerializeBlockIdxY(), + RematerializeBlockIdxX(), + RematerializeBlockIdxZ() + }; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Threadblock swizzling function for batched GEMMs +struct GemmBatchedIdentityThreadblockSwizzle { + + /// Returns the shape of the problem in units of logical tiles + MCTLASS_HOST_DEVICE + GemmCoord get_tiled_shape( + GemmCoord problem_size, + GemmCoord tile_size, + int batch_count) const { + + return GemmCoord( + (problem_size.m() + tile_size.m() - 1) / tile_size.m(), + (problem_size.n() + tile_size.n() - 1) / tile_size.n(), + batch_count % (1 << 16)); + } + + /// Computes CUDA grid dimensions given a size in units of logical tiles + MCTLASS_HOST_DEVICE + dim3 get_grid_shape(GemmCoord tiled_shape) const { + return dim3(tiled_shape.m(), tiled_shape.n(), tiled_shape.k()); + } + + /// Calculates optimal swizzle width + MCTLASS_HOST_DEVICE + int get_log_tile(GemmCoord tiled_shape) const { + return 0; + } + + /// Obtains the threadblock offset (in units of threadblock-scoped tiles) + MCTLASS_DEVICE + GemmCoord get_tile_offset(GemmCoord tiled_shape) const { + return GemmCoord{ + RematerializeBlockIdxX(), + RematerializeBlockIdxY(), + RematerializeBlockIdxZ() + }; + } + + /// Obtains the threadblock offset (in units of threadblock-scoped tiles) + MCTLASS_DEVICE + GemmCoord get_tile_offset(int log_tile) const { + int block_idx_x = RematerializeBlockIdxX(); + int block_idx_y = RematerializeBlockIdxY(); + int block_idx_z = RematerializeBlockIdxZ(); + + return GemmCoord{(block_idx_x >> log_tile), // + (block_idx_y << log_tile) + ((block_idx_x) & ((1 << (log_tile)) - 1)), + block_idx_z}; + } + + /// Gets the batch index + MCTLASS_DEVICE + int get_batch_idx() const { + return RematerializeBlockIdxZ(); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Threadblock swizzling function for split-K GEMMs +template +struct GemmSplitKIdentityThreadblockSwizzle { + + int const kTile = N; + + /// Returns the shape of the problem in units of logical tiles + MCTLASS_HOST_DEVICE + GemmCoord get_tiled_shape( + GemmCoord problem_size, + GemmCoord tile_size, + int partitions) const { + + return GemmCoord( + (problem_size.m() + tile_size.m() - 1) / tile_size.m(), + (problem_size.n() + tile_size.n() - 1) / tile_size.n(), + partitions); + } + + /// Calculates optimal swizzle width + MCTLASS_HOST_DEVICE + int get_log_tile(GemmCoord tiled_shape) const { + auto n = tiled_shape.n(); + // Thresholds picked so that it doesn't cause too many no-op CTAs + if (N >= 8 && n >= 6) + return 3; + else if (N >= 4 && n >= 3) + return 2; + else if (N >= 2 && n >= 2) + return 1; + else + return 0; + } + + /// Computes CUDA grid dimensions given a size in units of logical tiles + MCTLASS_HOST_DEVICE + dim3 get_grid_shape(GemmCoord tiled_shape) const { + int tile = 1 << get_log_tile(tiled_shape); + return dim3(tiled_shape.m() * tile, (tiled_shape.n() + tile - 1) / tile, tiled_shape.k()); + } + + /// Obtains the threadblock offset (in units of threadblock-scoped tiles) + MCTLASS_DEVICE + GemmCoord get_tile_offset(int log_tile) const { + int block_idx_x = RematerializeBlockIdxX(); + int block_idx_y = RematerializeBlockIdxY(); + int block_idx_z = RematerializeBlockIdxZ(); + + return GemmCoord{(block_idx_x >> log_tile), // + (block_idx_y << log_tile) + ((block_idx_x) & ((1 << (log_tile)) - 1)), + block_idx_z}; + } + + /// Obtains the threadblock offset (in units of threadblock-scoped tiles) + MCTLASS_DEVICE + GemmCoord get_tile_offset(GemmCoord tiled_shape) const { + + int const kTile = N; + int block_idx_x = RematerializeBlockIdxX(); + int block_idx_y = RematerializeBlockIdxY(); + + if ((tiled_shape.m() < kTile) || (tiled_shape.n() < kTile)) + return GemmCoord{block_idx_x, block_idx_y, RematerializeBlockIdxZ()}; + + return GemmCoord{ + (block_idx_x / kTile), + (block_idx_y * kTile) + (block_idx_x % kTile), + RematerializeBlockIdxZ() + }; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Threadblock swizzling function for split-K GEMMs +struct GemmSplitKHorizontalThreadblockSwizzle { + + /// Returns the shape of the problem in units of logical tiles + MCTLASS_HOST_DEVICE + GemmCoord get_tiled_shape( + GemmCoord problem_size, + GemmCoord tile_size, + int partitions) const { + + return GemmCoord( + (problem_size.m() + tile_size.m() - 1) / tile_size.m(), + (problem_size.n() + tile_size.n() - 1) / tile_size.n(), + partitions); + } + + /// Computes CUDA grid dimensions given a size in units of logical tiles + MCTLASS_HOST_DEVICE + dim3 get_grid_shape(GemmCoord tiled_shape) const { + return dim3(tiled_shape.n(), tiled_shape.m(), tiled_shape.k()); + } + + /// Calculates optimal swizzle width + MCTLASS_HOST_DEVICE + int get_log_tile(GemmCoord tiled_shape) const { + return 0; + } + + /// Obtains the threadblock offset (in units of threadblock-scoped tiles) + MCTLASS_DEVICE + GemmCoord get_tile_offset(int log_tile) const { + return GemmCoord{ + RematerializeBlockIdxY(), + RematerializeBlockIdxX(), + RematerializeBlockIdxZ() + }; + } + + /// Obtains the threadblock offset (in units of threadblock-scoped tiles) + MCTLASS_DEVICE + GemmCoord get_tile_offset(GemmCoord tiled_shape) const { + return GemmCoord{ + RematerializeBlockIdxY(), + RematerializeBlockIdxX(), + RematerializeBlockIdxZ() + }; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Threadblock swizzling function for batched GEMVs +struct GemvBatchedStridedThreadblockDefaultSwizzle { + + /// Returns the shape of the problem in units of logical tiles + MCTLASS_HOST_DEVICE + BatchedGemmCoord get_tiled_shape( + BatchedGemmCoord problem_size, + BatchedGemmCoord tile_size) const { + + return BatchedGemmCoord( + 1, // M is always 1 + (problem_size.n() + tile_size.n() - 1) / tile_size.n(), + (problem_size.k() + tile_size.k() - 1) / tile_size.k(), + (problem_size.batch() + tile_size.batch() - 1) / tile_size.batch()); + } + + /// Computes CUDA grid dimensions given a size in units of logical tiles + MCTLASS_HOST_DEVICE + dim3 get_grid_shape(BatchedGemmCoord tiled_shape) const { + return dim3(tiled_shape.n(), tiled_shape.batch(), tiled_shape.k()); + } + + /// Calculates optimal swizzle width + MCTLASS_HOST_DEVICE + int get_log_tile(GemmCoord tiled_shape) const { + return 0; + } + + /// Obtains the threadblock offset (in units of threadblock-scoped tiles) + MCTLASS_DEVICE + BatchedGemmCoord get_tile_offset(int log_tile) const { + return BatchedGemmCoord{ + 0, // M is always 1 + RematerializeBlockIdxX(), + RematerializeBlockIdxZ(), + RematerializeBlockIdxY(), + }; + } + + /// Obtains the threadblock offset (in units of threadblock-scoped tiles) + MCTLASS_DEVICE + BatchedGemmCoord get_tile_offset() const { + return BatchedGemmCoord{ + 0, // M is always 1 + RematerializeBlockIdxX(), + RematerializeBlockIdxZ(), + RematerializeBlockIdxY(), + }; + } + + /// Gets the batch tile index + MCTLASS_DEVICE + int get_batch_tile_idx() const { + return RematerializeBlockIdxY(); + } + + /// Gets the absolute batch index + MCTLASS_DEVICE + int get_batch_idx() const { + return RematerializeBlockDimY()*RematerializeBlockIdxY() + RematerializeThreadIdxY(); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/gemm/threadblock/threadblock_swizzle_streamk.h b/csrc/mctlass/include/mctlass/gemm/threadblock/threadblock_swizzle_streamk.h new file mode 100644 index 0000000..f11c301 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/threadblock/threadblock_swizzle_streamk.h @@ -0,0 +1,801 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Implements streamk threadblock mapping blockIdx to GEMM problems. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/platform/platform.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/conv/conv2d_problem_size.h" +#include "mctlass/conv/conv3d_problem_size.h" +#include "mctlass/gemm/threadblock/index_remat.h" + +#if !defined(__MACACC_RTC__) +#include +#include "mctlass/core_io.h" +#include "mctlass/trace.h" +#endif + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Threadblock mapping control for GEMMs +struct ThreadblockSwizzleStreamK { + + /// Advertise StreamkFeature + using StreamkFeature = void; + + + /// Kernel traits + template + struct KernelTraits {}; + + + /// Reduction strategy + enum ReductionStrategy + { + kNone, // Data-parallel strategy (no seams, fixup, etc.) + + kAtomic, // Non-deterministic reduction of SK-block partials using atomic aggregation in L2 + + kMixed, // Deterministic reduction of SK-block partials employing either: + // (a) A separate wave of reduction thread blocks" (for scenarios with lots of + // SK-blocks per SK-tile) + // (b) Turnstile-ordered atomic aggregation in L2 (for scenarios with few + // SK-blocks per SK-tile) + }; + + static ReductionStrategy const kReductionStrategy = kMixed; + + + // + // Heuristics + // + + /// Data-parallel wave-quantization efficiency threshold (above which we go data-parallel) + static float constexpr kDpEfficiencyThreshold = 0.92f; + + /// Minimum number of MAC-iterations per streamk block + static int const kMinItersPerSkBlock = 2; + + /// Height in CTAs of a grid rasterization cohort + static int const kCohortCtasM = 8; + + /// Width in CTAs of a grid rasterization cohort + static int const kCohortCtasN = 4; + + /// Number of CTAs per cohort + static int const kCtasPerCohort = kCohortCtasN * kCohortCtasM; + + /// Cost-equivalent number of SM-iterations for fixup I/O + static int const kFixupStartupIterEquiv = 10; + static int const kFixupPeerIterEquiv = 3; + + + // + // Member state + // + + + /// The 3D value-extents of the GEMM computation volume (m,n,k) + GemmCoord problem_size; + + /// Div/mod accelerators + FastDivmod div_mod_tiled_shape_m; + FastDivmod div_mod_tiled_shape_n; + FastDivmod div_mod_tiled_cohort_shape_n; + FastDivmod div_mod_iters_per_tile; + + /// Whether to perform cohort CTA rasterization + bool cohort_raster; + + // Whether to pad and remap block indices + bool remap_block_indices; + + /// CTA occupancy per SM + int sm_occupancy; + + /// Number of SMs for dispatch heuristics to load-balance using Stream-K CTAs (wave size) + int avail_sms; + + int dp_blocks; /// Number of data-parallel thread blocks in the grid + int dp_first_wave_tiles; /// Number of output tiles each CTA in the first DP wave will produce + + /// Number of reduction blocks in the grid + int reduction_blocks; + + int sk_waves; + int sk_tiles; + int sk_big_blocks_per_region; + int sk_iters_per_region; + + /// Div/mod accelerators + FastDivmod div_mod_sk_iters_per_normal_block; + FastDivmod div_mod_sk_iters_per_big_block; + FastDivmod div_mod_sk_iters_per_region; + FastDivmod div_mod_sk_regions; //!! used in block map + FastDivmod div_mod_sk_blocks_per_region; //!! used in block map + + /// The batch count + int batch_count; + + + // + // Host+device interface + // + + /// Constructor + MCTLASS_HOST_DEVICE + ThreadblockSwizzleStreamK() {} + + /// Returns the GEMM volume in thread block tiles + MCTLASS_HOST_DEVICE + GemmCoord tiled_shape() const + { + return GemmCoord( + static_cast(div_mod_tiled_shape_m), + static_cast(div_mod_tiled_shape_n), + batch_count); + } + + /// Number of iterations per output tile + MCTLASS_HOST_DEVICE + int iters_per_tile() const + { + return static_cast(div_mod_iters_per_tile); + } + + /// Number of iterations for normal SK-blocks + MCTLASS_HOST_DEVICE + int sk_iters_per_normal_block() const + { + return static_cast(div_mod_sk_iters_per_normal_block); + } + + /// Number of SK regions + MCTLASS_HOST_DEVICE + int sk_regions() const + { + return static_cast(div_mod_sk_regions); + } + + /// Number of SK blocks per region (splitting factor) + MCTLASS_HOST_DEVICE + int sk_blocks_per_region() const + { + return static_cast(div_mod_sk_blocks_per_region); + } + + + // + // Host-side interface + // + + /// Debug print + void Print() + { +#ifndef __MACA_ARCH__ + auto tiles = tiled_shape().mn().product(); + std::cout << + "problem_size: (" << problem_size.m() << "," << problem_size.n() << ")" << + ", tiled_shape: (" << tiled_shape().m() << "," << tiled_shape().n() << ")" << + ", tiles: " << tiles << + ", dp_tiles: " << tiles - sk_tiles << + ", sk_tiles: " << sk_tiles << + ", iters_per_tile: " << iters_per_tile() << + ", reduction_blocks: " << reduction_blocks << + ", dp_blocks: " << dp_blocks << + ", dp_waves: " << dp_blocks / avail_sms << + ", dp_first_wave_tiles: " << dp_first_wave_tiles << + ", sk_blocks_per_region: " << sk_blocks_per_region() << + ", sk_regions: " << sk_regions() << + ", sk_waves: " << sk_waves << + ", sk_iters_per_normal_block: " << sk_iters_per_normal_block() << + ", sk_big_blocks_per_region: " << sk_big_blocks_per_region << + ", remap_block_indices: " << remap_block_indices << + ", cohort_raster: " << cohort_raster << + ", sm_occupancy: " << sm_occupancy << + ", avail_sms: " << avail_sms << + ", num_blocks: " << get_num_blocks() << + "\n\n"; +#endif + } + + + // Compute sk_blocks to dispatch for a given number of sk_tiles + static void get_sk_blocks( + int &sk_blocks, /// [out] + int &savings_iters, /// [out] + int sk_tiles, + int iters_per_tile, + int avail_sms, + int max_sk_occupancy, + bool allow_partial_wave) + { + savings_iters = INT_MIN; + sk_blocks = 0; + + if (sk_tiles == 0) { + return; + } + + int sk_iters = sk_tiles * iters_per_tile; + + int dp_equiv_waves = (sk_tiles + avail_sms - 1) / avail_sms; + int dp_equiv_iters = iters_per_tile * dp_equiv_waves; + + int min_sk_blocks = (allow_partial_wave) ? fast_min(avail_sms, sk_tiles + 1) : avail_sms; + int max_sk_blocks = fast_min(avail_sms * max_sk_occupancy, sk_iters / kMinItersPerSkBlock); + + for (int trial_sk_blocks = min_sk_blocks; trial_sk_blocks <= max_sk_blocks; ++trial_sk_blocks) + { + int sk_waves = (trial_sk_blocks + avail_sms - 1) / avail_sms; + int max_sk_iters_per_block = (sk_iters + trial_sk_blocks - 1) / trial_sk_blocks; + int sk_iter_equiv = max_sk_iters_per_block * sk_waves; + + int num_peers = ((trial_sk_blocks + sk_tiles - 1) / sk_tiles) + 1; // add one for alignment skew + + float iter_cost = 0.02f * float(num_peers) * float(sk_iter_equiv); + + if (trial_sk_blocks % sk_tiles == 0) + { + // aligned + num_peers = (trial_sk_blocks / sk_tiles); + + iter_cost = 0.0f; + } + + float peer_cost = 2.0f * float(num_peers); + + float base_cost = 2.0f * float(sk_waves); + + int fixup_iter_equiv = int(base_cost + iter_cost + peer_cost); + + int trial_savings_iters = dp_equiv_iters - sk_iter_equiv - fixup_iter_equiv; + + if (trial_savings_iters >= savings_iters) { + savings_iters = trial_savings_iters; + sk_blocks = trial_sk_blocks; + } + } + } + + + /// Determine the populations of DP and SK blocks to invoke for the given number of output tiles + static void get_blocks( + int &dp_tiles, /// [out] + int &sk_blocks, /// [out] + int output_tiles, + int iters_per_tile, + int avail_sms, + int sm_occupancy) + { + int full_waves = output_tiles / avail_sms; + int full_wave_tiles = full_waves * avail_sms; + int partial_wave_tiles = output_tiles - full_wave_tiles; + + int score = -1; + dp_tiles = output_tiles; + sk_blocks = 0; + + if (partial_wave_tiles == 0) + { + // Perfect quantization + return; + } + + if (full_waves < sm_occupancy) + { + // We're less than full GPU occupancy + + // Form the SK wave from the partial wave to get us up to full GPU occupancy + int max_sk_occupancy = sm_occupancy - full_waves; + + dp_tiles = full_wave_tiles; + + get_sk_blocks( + sk_blocks, + score, + partial_wave_tiles, + iters_per_tile, + avail_sms, + max_sk_occupancy, + true); // we can run with less than a full wave of SK-blocks + + if (score < 0) { + // not profitable + sk_blocks = 0; + dp_tiles = output_tiles; + } + + return; + } + + // We're at (or greater) than GPU occupancy + + if ((sm_occupancy > 1 ) && (full_waves % sm_occupancy == sm_occupancy - 1)) + { + // If occupancy is more than one CTA per SM, form the SK wave from the partial + // wave to get us to full GPU occupancy + int max_sk_occupancy = 1; + + dp_tiles = full_wave_tiles; + + get_sk_blocks( + sk_blocks, + score, + partial_wave_tiles, + iters_per_tile, + avail_sms, + max_sk_occupancy, + true); // we can run with less than a full wave of SK-blocks + + if (score >= 0) { + return; + } + } + + // Form the SK wave by combining the last full wave and the partial wave + // We're less than full GPU occupancy + dp_tiles = full_wave_tiles - avail_sms; + + int max_sk_occupancy = sm_occupancy - ((full_waves - 1) % sm_occupancy); + + get_sk_blocks( + sk_blocks, + score, + partial_wave_tiles + avail_sms, + iters_per_tile, + avail_sms, + max_sk_occupancy, + false); // we cannot run with less than a full wave of SK-blocks + + if (score < 0) { + // not profitable + sk_blocks = 0; + dp_tiles = output_tiles; + } + + } + + /// Constructor: *Gemm* problem size (m, n, k) + template + ThreadblockSwizzleStreamK( + KernelTraits const kernel_traits_, + GemmUniversalMode const mode_, + GemmCoord const problem_size_, + GemmCoord const tile_size_, + int const batch_split_, /// Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor (1 defaults to StreamK, >1 emulates Split-K) + int const sm_occupancy_, + int const device_sms_, + int const avail_sms_) /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling) + : + problem_size(problem_size_), + batch_count((mode_ == GemmUniversalMode::kBatched) ? batch_split_ : 1), + reduction_blocks(0), + dp_blocks(0), + dp_first_wave_tiles(1), // Default: one tile per DP-block in the first wave of DP blocks + sk_tiles(0), + sk_big_blocks_per_region(0), + sk_iters_per_region(0), + sk_waves(0), + sm_occupancy(sm_occupancy_), + remap_block_indices(false), + avail_sms(fast_max(1, avail_sms_)), + cohort_raster(false) + { + int gpu_occupancy = device_sms_ * sm_occupancy; + int iters_per_tile = (problem_size.k() + tile_size_.k() - 1) / tile_size_.k(); + int sk_iters_per_normal_block = 0; + + int sk_regions = 1; // Default: a single region of iteration space (across all SK tiles) + int sk_blocks_per_region = 0; + + GemmCoord tiled_shape( + (problem_size.m() + tile_size_.m() - 1) / tile_size_.m(), + (problem_size.n() + tile_size_.n() - 1) / tile_size_.n(), + batch_count); + + size_t problem_bytes = + (sizeof(typename GemmKernel::ElementC) * problem_size.m() * problem_size.n()) + + (sizeof(typename GemmKernel::ElementA) * problem_size.m() * problem_size.k()) + + (sizeof(typename GemmKernel::ElementB) * problem_size.k() * problem_size.n()); + + size_t problem_flops = size_t(problem_size.m()) * size_t(problem_size.n()) * size_t(problem_size.k()) * 2; + + float flops_per_byte = float(problem_flops) / float(problem_bytes); + + int output_tiles = tiled_shape.m() * tiled_shape.n(); + int waves = (output_tiles + avail_sms - 1) / avail_sms; + float dp_efficiency = float(output_tiles) / float(waves * avail_sms); + + // + // Determine dispatch composition of DP-tiles and SK-blocks + // + + // Start with a DP-only configuration + int dp_tiles = output_tiles; // Number of data-parallel tiles + int sk_blocks = 0; // Number of thread blocks to produce the remaining SK tiles + + // Only kGemm mode allows for SK load balancing + if (mode_ == GemmUniversalMode::kGemm) + { + int split_factor = batch_split_; + if (split_factor > 1) + { + // Split-K override + dp_tiles = 0; + sk_blocks = output_tiles * split_factor; + } + else if ((kReductionStrategy != kNone) && // Load-balancing strategy statically enabled + (avail_sms > 1)) // Plurality of SMs to load balance across + { + // Use heuristics + get_blocks( + dp_tiles, /// [out] + sk_blocks, /// [out] + output_tiles, + iters_per_tile, + avail_sms, + sm_occupancy); + } + } + + sk_tiles = output_tiles - dp_tiles; + + + // Compute SK block iteration details + if (sk_blocks > 0) + { + sk_waves = (sk_blocks + avail_sms - 1) / avail_sms; + + int sk_iters = sk_tiles * iters_per_tile; + sk_blocks = fast_min(sk_blocks, sk_iters); + + sk_iters_per_normal_block = sk_iters / sk_blocks; + int extra_sk_iters = sk_iters - (sk_iters_per_normal_block * sk_blocks); + int sk_big_blocks = extra_sk_iters; + + if ((sk_blocks > sk_tiles) && (sk_blocks % sk_tiles == 0)) + { + // Split-K decomposition + sk_regions = sk_tiles; + } + + sk_blocks_per_region = sk_blocks / sk_regions; + sk_big_blocks_per_region = sk_big_blocks / sk_regions; + sk_iters_per_region = sk_iters / sk_regions; + + // Use a separate reduction wave when all of: + // - Non-atomic reduction stratgy + // - The number of SK waves won't fully occupy the GPU (Otherwise we don't have + // a strong-scaling case for more parallel reduction) + // - More than three peers working on an SK tile. (This occurs when the ratio of + // SK-blocks to SK-tiles > 2, as a single tile may be covered by four SK-blocks, + // e.g.:[partial-block | block | block | partial-block] ). With three or + // less peers, the two non-finishing SK-blocks are not expexted to contend. + if ((kReductionStrategy == kMixed) && + (sk_waves < sm_occupancy) && + (sk_blocks > 2 * sk_tiles)) + { + // Launch a reduction block for every accumulator fragment in each SK-tile + static const int kAccumulatorFragments = GemmKernel::Epilogue::kAccumulatorFragments; + reduction_blocks = sk_tiles * kAccumulatorFragments; + + } + + // When we have a multi-occupancy kernel and at least two waves of active blocks (where + // at least one wave is SK blocks), we need to (1) dispatch at least four waves, and (2) + // remap the block indices so that we can reliably spread the SK blocks evenly across the + // device's first SM occupancy valence. Also see get_num_blocks() and get_block_idx(). + remap_block_indices = ( + (sm_occupancy > 1) && + (device_sms_ == avail_sms) && + (get_num_active_blocks() > avail_sms * 2)); + + // Initialize fast div/mod members related to SK + div_mod_sk_iters_per_normal_block = FastDivmod(sk_iters_per_normal_block); + div_mod_sk_iters_per_big_block = FastDivmod(sk_iters_per_normal_block + 1); + div_mod_sk_iters_per_region = FastDivmod(sk_iters_per_region); + div_mod_sk_regions = FastDivmod(sk_regions); + div_mod_sk_blocks_per_region = FastDivmod(sk_blocks_per_region); + } + + // + // Compute DP blocks + // + + dp_blocks = dp_tiles; + + mctlass::gemm::GemmCoord tiled_cohort_shape( + (tiled_shape.m() + kCohortCtasM - 1) / kCohortCtasM, + (tiled_shape.n() + kCohortCtasN - 1) / kCohortCtasN, + tiled_shape.k()); + int cohort_blocks = (tiled_cohort_shape.m() * tiled_cohort_shape.n()) * kCtasPerCohort; + float cohort_efficiency = float(dp_blocks) / float(cohort_blocks); + + // Check if the SK tiles would be in cohorts that are in-bounds + bool sk_in_range = true; + if (sk_tiles > 0) + { + int last_sk_tile = sk_tiles - 1; + int cohort_tile_idx = last_sk_tile / kCtasPerCohort; + int cohort_grid_m = cohort_tile_idx / tiled_cohort_shape.n(); + int cohort_grid_n = (cohort_grid_m > 0) ? + tiled_cohort_shape.n() - 1 : + cohort_tile_idx % tiled_cohort_shape.n(); + + if ((((cohort_grid_m + 1) * kCohortCtasM) >= tiled_shape.m()) || + (((cohort_grid_n + 1) * kCohortCtasN) >= tiled_shape.n())) + { + sk_in_range = false; + } + + } + + // Decide if we're going to be doing cohort raster + if (sk_in_range && + (dp_blocks >= gpu_occupancy * 2) && + (cohort_efficiency > 0.85f)) + { + cohort_raster = true; + dp_blocks = cohort_blocks; + } + else if (sk_waves > 0) + { + // Update semi-persistence of first DP wave to ensure full grid wavesets + // (Only applies when there's an SK component and we're not doing blocked cohort rasterization) + int dp_tile_waves = (dp_tiles + avail_sms - 1) / avail_sms; + int full_dp_tile_waves = dp_tiles / avail_sms; + int waveset_excess = (sk_waves + dp_tile_waves) % sm_occupancy; + + if (dp_first_wave_tiles + waveset_excess <= full_dp_tile_waves) + { + dp_first_wave_tiles += waveset_excess; + dp_blocks -= (waveset_excess * avail_sms); + } + } + + // Setup fast-div/mod for device-side usage + div_mod_tiled_shape_m = FastDivmod(tiled_shape.m()); + div_mod_tiled_shape_n = FastDivmod(tiled_shape.n()); + div_mod_tiled_cohort_shape_n = FastDivmod(tiled_cohort_shape.n()); + div_mod_iters_per_tile = FastDivmod(iters_per_tile); + + } + + /// Number of blocks performing useful work + int get_num_active_blocks() const + { + return (sk_waves * avail_sms) + dp_blocks + reduction_blocks; + } + + /// Obtains number of threadblocks per GEMM + int get_num_blocks() const + { + int active_blocks = get_num_active_blocks(); + if (remap_block_indices) + { + // Add padding blocks if we are performing remapping in order to dispatch a grid of at least four waves + return fast_max(active_blocks, avail_sms * 4); + } + + return active_blocks; + } + + + /// Obtains grid extents in CTAs + dim3 get_grid_dims() const + { + return dim3(get_num_blocks(), 1, batch_count); + } + + + // + // Device-side interface + // + + /// Obtains number of threadblocks per GEMM + MCTLASS_DEVICE + int device_num_blocks() const + { + return gridDim.x; + } + + /// Obtains tile index for the given sk iteration + MCTLASS_DEVICE + int get_sk_tile_idx(int iter) const + { + int tile_idx = div_mod_iters_per_tile.div(iter); + return tile_idx; + } + + /// Obtains the batch index + MCTLASS_DEVICE + int get_batch_idx() const + { + return RematerializeBlockIdxZ(); + } + + /// Obtains the calling threadblock's tiled coordinates for the given tile index + MCTLASS_DEVICE + GemmCoord get_tile_offset(int tile_idx) const + { + int m, n; + + // row-major raster + div_mod_tiled_shape_n(m, n, tile_idx); + + if (tiled_shape().m() < tiled_shape().n()) + { + // column-major raster + div_mod_tiled_shape_m(n, m, tile_idx); + } + + if (cohort_raster) + { + // tiled cohort raster + int cohort_tile_idx = tile_idx / kCtasPerCohort; + int cohort_grid_m, cohort_grid_n; + div_mod_tiled_cohort_shape_n(cohort_grid_m, cohort_grid_n, cohort_tile_idx); + + int block_idx_cohort = tile_idx % kCtasPerCohort; + int block_cohort_m = block_idx_cohort / kCohortCtasN; + int block_cohort_n = block_idx_cohort % kCohortCtasN; + + m = (cohort_grid_m * kCohortCtasM) + block_cohort_m; + n = (cohort_grid_n * kCohortCtasN) + block_cohort_n; + } + + return GemmCoord(m, n, get_batch_idx()); + } + + /// Obtains the calling threadblock's tiled coordinates for the given tile index (row-major rasterization) + MCTLASS_DEVICE + GemmCoord get_tile_offset_row_major(int tile_idx) const + { + // row-major raster + int m, n; + div_mod_tiled_shape_n(m, n, tile_idx); + return GemmCoord(m, n, get_batch_idx()); + } + + /// Obtains calling threadblock's linear threadblock index + MCTLASS_DEVICE + int get_block_idx() const + { + int block_idx = RematerializeBlockIdxX(); + + // Remap the block indices for the first two waves of thread blocks if + // we have multi-occupancy and the grid constitutes four or more waves + if (remap_block_indices && (block_idx < avail_sms * 2)) + { + int dest_sm = block_idx / 2; + int dest_wave = block_idx % 2; + int remapped_block_idx = dest_sm + (dest_wave * avail_sms); + block_idx = remapped_block_idx; + } + + // Remap block indices to interleave SK regions to limit intra-region waiting + if (block_idx < sk_regions() * sk_blocks_per_region()) + { + int block_in_region; + int region; + div_mod_sk_regions(block_in_region, region, block_idx); + block_idx = (region * sk_blocks_per_region()) + block_in_region; + } + + return block_idx; + } + + + /// Obtains calling linear threadblock index of the first block to work on the given tile + MCTLASS_DEVICE + int get_sk_block_idx(int iter) const + { + int region_idx; + int iter_in_region; + div_mod_sk_iters_per_region(region_idx, iter_in_region, iter); + + int big_block_iters = (sk_big_blocks_per_region * sk_iters_per_normal_block()) + sk_big_blocks_per_region; // number of iterations in the region's big blocks + int normal_block_iters = iter_in_region - big_block_iters; // number of iterations in the region's normal blocks + + int big_block_idx_in_region = div_mod_sk_iters_per_big_block.div(iter_in_region); + int normal_block_idx_in_region = sk_big_blocks_per_region + div_mod_sk_iters_per_normal_block.div(normal_block_iters); + + int block_idx_in_region = (big_block_idx_in_region < sk_big_blocks_per_region) ? + big_block_idx_in_region : + normal_block_idx_in_region; + + int owning_block_idx = (sk_blocks_per_region() * region_idx) + block_idx_in_region; + + return owning_block_idx; + } + + /// Obtains iteration extends for the given SK block index + MCTLASS_DEVICE + void get_iter_extents( + int sk_block_idx, + int &block_iter_begin, + int &block_iter_end) const + { + int region_idx; + int block_idx_in_region; + div_mod_sk_blocks_per_region(region_idx, block_idx_in_region, sk_block_idx); + + block_iter_begin = (region_idx * sk_iters_per_region) + (block_idx_in_region * sk_iters_per_normal_block()); + + // Adjust extents for the first "num_big_blocks" blocks that get one extra iteration + int block_iters = sk_iters_per_normal_block(); + if (block_idx_in_region < sk_big_blocks_per_region) { + // This is a +1 iteration block + block_iter_begin += block_idx_in_region; + block_iters++; + } else { + // This is a regular block + block_iter_begin += sk_big_blocks_per_region; + } + block_iter_end = block_iter_begin + block_iters; + } + + + /// Obtains calling linear threadblock index of the first block to work on the given tile + MCTLASS_DEVICE + int get_first_block_idx(int tile_idx, int block_idx) const + { + if (tile_idx >= sk_tiles) { + // DP tile + return block_idx; + } + + int iter = tile_idx * iters_per_tile(); + return get_sk_block_idx(iter); + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace gemm +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/gemm/warp/default_mma_complex_tensor_op.h b/csrc/mctlass/include/mctlass/gemm/warp/default_mma_complex_tensor_op.h new file mode 100644 index 0000000..fba2c8c --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/default_mma_complex_tensor_op.h @@ -0,0 +1,612 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Default warp-level GEMM operators selected by data type, size, and layouts of operands. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/warp/mma_complex_tensor_op.h" +#include "mctlass/gemm/warp/mma_complex_tensor_op_fast_f32.h" +#include "mctlass/gemm/warp/mma_gaussian_complex_tensor_op.h" +#include "mctlass/layout/tensor_op_multiplicand_sm80.h" + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Complex transform on A operand + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex transform on B operand + ComplexTransform TransformB = ComplexTransform::kNone, + /// Multiply-add operator (arch::OpMultiplyAddComplex, arch::OpMultiplyGaussianComplex) + typename Operator_ = arch::OpMultiplyAddComplex> +struct DefaultMmaComplexTensorOp; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex*complex case +// 4 real-valued mma operations +// A = (ar + j ai), B (br +j bi), D = AB +// D = dr + j di = (ar*br - ai*bi) + j (ar*bi + ai*br) +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Real-valued underlying type of complex-valued A operand + typename RealElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Real-valued underlying type of complex-valued B operand + typename RealElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Real-valued underlying type of complex-valued C operand + typename RealElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB> +struct DefaultMmaComplexTensorOp< + WarpShape_, + InstructionShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + TransformA, + TransformB, + arch::OpMultiplyAddComplex> { + + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + InstructionShape_, + 32, + RealElementA, + mctlass::layout::RowMajor, + RealElementB, + mctlass::layout::ColumnMajor, + RealElementC, + mctlass::layout::RowMajor, + arch::OpMultiplyAdd>, + mctlass::MatrixShape<1, 1> + >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaComplexTensorOp< + WarpShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + Policy, + TransformA, + TransformB>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex*complex case using GaussianComplex operation +// 3 real-valued mma operations +// A = (ar + j ai), B = (br +j bi), D = AB +// P1 = (ar + ai) * br, P2 = - ar * (br - bi), P3 = ai * (br + bi) +// D = dr + j di = (P1 - P3) + j (P1 + P2) +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Real-valued underlying type of complex-valued A operand + typename RealElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Real-valued underlying type of complex-valued B operand + typename RealElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Real-valued underlying type of complex-valued C operand + typename RealElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB> +struct DefaultMmaComplexTensorOp< + WarpShape_, + InstructionShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + TransformA, + TransformB, + arch::OpMultiplyAddGaussianComplex> { + + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + InstructionShape_, + 32, + RealElementA, + mctlass::layout::RowMajor, + RealElementB, + mctlass::layout::ColumnMajor, + RealElementC, + mctlass::layout::RowMajor, + arch::OpMultiplyAdd>, + mctlass::MatrixShape<1, 1> + >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaGaussianComplexTensorOp< + WarpShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + Policy, + TransformA, + TransformB>; +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Partial specialization - input and output types are complex*complex +// Use TF32 tensor operation internally +// 4 real-valued mma.sync.aligned.m16n8k8.f32.tf32.tf32.f32 operations on TF32 +// A = (ar + j ai), B (br +j bi), D = AB +// D = dr + j di = (ar*br - ai*bi) + j (ar*bi + ai*br) +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB> +struct DefaultMmaComplexTensorOp< + WarpShape_, + InstructionShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + TransformA, + TransformB, + arch::OpMultiplyAddComplex> { + + // Complex floating point tensor operation use mma.sync.aligned.m16n8k8.f32.tf32.tf32.f32 mma instruction + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + InstructionShape_, + 32, + tfloat32_t, + mctlass::layout::RowMajor, + tfloat32_t, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::RowMajor, + arch::OpMultiplyAdd>, + mctlass::MatrixShape<1, 1> + >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaComplexTensorOp< + WarpShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + Policy, + TransformA, + TransformB>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Partial specialization - input and output types are complex*complex +// Use BF16 tensor operation internally +// 4 real-valued mma.sync.aligned.m16n8k8.f32.bf16.bf16.f32 operations on BF16 +// A = (ar + j ai), B (br +j bi), D = AB +// D = dr + j di = (ar*br - ai*bi) + j (ar*bi + ai*br) +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB> +struct DefaultMmaComplexTensorOp< + WarpShape_, + InstructionShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + TransformA, + TransformB, + arch::OpMultiplyAddFastBF16> { + + // Complex floating point tensor operation use mma.sync.aligned.m16n8k8.f32.bf16.bf16.f32 mma instruction + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + InstructionShape_, + 32, + bfloat16_t, + mctlass::layout::RowMajor, + bfloat16_t, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::RowMajor, + arch::OpMultiplyAdd>, + mctlass::MatrixShape<1, 1> + >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaComplexTensorOp< + WarpShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + Policy, + TransformA, + TransformB>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Partial specialization - input and output types are complex*complex +// Use F16 tensor operation internally +// 4 real-valued mma.sync.aligned.m16n8k8.f32.f16.f16.f32 operations on F16 +// A = (ar + j ai), B (br +j bi), D = AB +// D = dr + j di = (ar*br - ai*bi) + j (ar*bi + ai*br) +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB> +struct DefaultMmaComplexTensorOp< + WarpShape_, + InstructionShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + TransformA, + TransformB, + arch::OpMultiplyAddFastF16> { + + // Complex floating point tensor operation use mma.sync.aligned.m16n8k8.f32.f16.f16.f32 mma instruction + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + InstructionShape_, + 32, + half_t, + mctlass::layout::RowMajor, + half_t, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::RowMajor, + arch::OpMultiplyAdd>, + mctlass::MatrixShape<1, 1> + >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaComplexTensorOp< + WarpShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + Policy, + TransformA, + TransformB>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// 3xTF32 or 4xTF32 (fast and accurate complex operation) +/// Partial specialization - input and output types are complex * complex +// Use 3xTF32 or 4xTF32 tensor operation internally +// 4 real-valued mma.sync.aligned.m16n8k8.f32.tf32.tf32.f32 operations on TF32 +// A = (ar + j ai), B (br +j bi), D = AB +// D = dr + j di = 3x[(ar*br - ai*bi) + j (ar*bi + ai*br)] +///////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB> +struct DefaultMmaComplexTensorOp< + WarpShape_, + InstructionShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + TransformA, + TransformB, + arch::OpMultiplyAddComplexFastF32> { + + // Complex floating point tensor operation use mma.sync.aligned.m16n8k8.f32.tf32.tf32.f32 mma instruction + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + InstructionShape_, + 32, + tfloat32_t, + mctlass::layout::RowMajor, + tfloat32_t, + mctlass::layout::ColumnMajor, + float, + mctlass::layout::RowMajor, + arch::OpMultiplyAdd>, + mctlass::MatrixShape<1, 1> + >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaComplexTensorOpFastF32< + WarpShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + Policy, + TransformA, + TransformB>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex*complex case +// 4 real-valued mma.sync.aligned.m16n8k4.f64.f64.f64.f64 operations +// A = (ar + j ai), B (br +j bi), D = AB +// D = dr + j di = (ar*br - ai*bi) + j (ar*bi + ai*br) +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename WarpShape_, + /// Real-valued underlying type of complex-valued A operand + typename RealElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Real-valued underlying type of complex-valued B operand + typename RealElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Real-valued underlying type of complex-valued C operand + typename RealElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB> +struct DefaultMmaComplexTensorOp< + WarpShape_, + GemmShape<16, 8, 4>, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + TransformA, + TransformB, + arch::OpMultiplyAddComplex> { + + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + GemmShape<16, 8, 4>, + 32, + RealElementA, + mctlass::layout::RowMajor, + RealElementB, + mctlass::layout::ColumnMajor, + RealElementC, + mctlass::layout::RowMajor, + arch::OpMultiplyAdd>, + mctlass::MatrixShape<1, 1> + >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaComplexTensorOp< + WarpShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + Policy, + TransformA, + TransformB, + true>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Partial specialization for complex*complex case using GaussianComplex operation +// 3 real-valued mma.sync.aligned.m16n8k4.f64.f64.f64.f64 operations +// A = (ar + j ai), B = (br +j bi), D = AB +// P1 = (ar + ai) * br, P2 = - ar * (br - bi), P3 = ai * (br + bi) +// D = dr + j di = (P1 - P3) + j (P1 + P2) +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename WarpShape_, + /// Real-valued underlying type of complex-valued A operand + typename RealElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Real-valued underlying type of complex-valued B operand + typename RealElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Real-valued underlying type of complex-valued C operand + typename RealElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB> +struct DefaultMmaComplexTensorOp< + WarpShape_, + GemmShape<16, 8, 4>, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + TransformA, + TransformB, + arch::OpMultiplyAddGaussianComplex> { + + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + GemmShape<16, 8, 4>, + 32, + RealElementA, + mctlass::layout::RowMajor, + RealElementB, + mctlass::layout::ColumnMajor, + RealElementC, + mctlass::layout::RowMajor, + arch::OpMultiplyAdd>, + mctlass::MatrixShape<1, 1> + >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaGaussianComplexTensorOp< + WarpShape_, + complex, + LayoutA, + complex, + LayoutB, + complex, + LayoutC, + Policy, + TransformA, + TransformB, + true>; +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/warp/default_mma_sparse_tensor_op.h b/csrc/mctlass/include/mctlass/gemm/warp/default_mma_sparse_tensor_op.h new file mode 100644 index 0000000..8434a85 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/default_mma_sparse_tensor_op.h @@ -0,0 +1,165 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Default warp-level GEMM operators selected by data type, size, and layouts of operands. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/warp/mma_sparse_tensor_op.h" + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Operator describing the tensor operation + typename Operator_ = arch::OpMultiplyAdd, + /// Number of partitions along K dimension + int PartitionsK = 1, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false +> +struct DefaultSparseMmaTensorOp; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial Specialization - inputs and output types are float - uses TF32 internally +template < + /// Shape of one matrix production operation (concept: GemmShape) + typename WarpShape_, + /// Shape of target matrix multiply instruction (concept: GemmShape) + typename InstructionShape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Number of partitions along K dimension + int PartitionsK, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor> +struct DefaultSparseMmaTensorOp< + WarpShape_, + InstructionShape_, + float, LayoutA, + float, LayoutB, + float, LayoutC, + arch::OpMultiplyAdd, PartitionsK, AccumulatorsInRowMajor> { + + // Uses TF32 internally + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::SparseMma< + InstructionShape_, + 32, + tfloat32_t, mctlass::layout::RowMajor, + tfloat32_t, mctlass::layout::ColumnMajor, + float, mctlass::layout::RowMajor, + arch::OpMultiplyAdd + >, + mctlass::MatrixShape<1, 1> >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::SparseMmaTensorOp< + WarpShape_, float, LayoutA, float, LayoutB, float, LayoutC, + Policy, PartitionsK, AccumulatorsInRowMajor>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for m-by-n-by-kgroup +template < + /// Shape of one matrix production operation (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A elements + typename ElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Data type of B elements + typename ElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Element type of C matrix + typename ElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Operator describing the tensor operation + typename Operator_, + /// Number of partitions along K dimension + int PartitionsK, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor> +struct DefaultSparseMmaTensorOp { + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::SparseMma, + mctlass::MatrixShape<1, 1> >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::SparseMmaTensorOp< + WarpShape_, ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, + Policy, PartitionsK, AccumulatorsInRowMajor>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/default_mma_tensor_op.h b/csrc/mctlass/include/mctlass/gemm/warp/default_mma_tensor_op.h new file mode 100644 index 0000000..d9f9859 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/default_mma_tensor_op.h @@ -0,0 +1,123 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Default warp-level GEMM operators selected by data type, size, and layouts of operands. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/warp/mma_tensor_op.h" + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Operator describing the tensor operation + typename Operator_ = arch::OpMultiplyAdd, + /// Number of partitions along K dimension + int PartitionsK = 1, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false> +struct DefaultMmaTensorOp; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for m-by-n-by-kgroup +template < + /// Shape of one matrix production operation (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A elements + typename ElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Data type of B elements + typename ElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Element type of C matrix + typename ElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Operator describing the tensor operation + typename Operator_, + /// Number of partitions along K dimension + int PartitionsK, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor> +struct DefaultMmaTensorOp { + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma, + mctlass::MatrixShape<1, 1> >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaTensorOp< + WarpShape_, ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, + Policy, PartitionsK, AccumulatorsInRowMajor>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "mctlass/gemm/warp/default_mma_tensor_op_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/default_mma_tensor_op_sm80.h b/csrc/mctlass/include/mctlass/gemm/warp/default_mma_tensor_op_sm80.h new file mode 100644 index 0000000..b8161b0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/default_mma_tensor_op_sm80.h @@ -0,0 +1,238 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Default warp-level GEMM operators selected by data type, size, and layouts of operands. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/arch/mma.h" +#include "mctlass/gemm/warp/mma_tensor_op.h" +#include "mctlass/gemm/warp/mma_tensor_op_fast_f32.h" +#include "mctlass/gemm/warp/default_mma_tensor_op.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial Specialization - inputs and output types are float - uses BF16 internally +template < + /// Shape of one matrix production operation (concept: GemmShape) + typename WarpShape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Number of partitions along K dimension + int PartitionsK, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor> +struct DefaultMmaTensorOp< + WarpShape_, + GemmShape<16, 8, 8>, + float, LayoutA, + float, LayoutB, + float, LayoutC, + arch::OpMultiplyAddFastBF16, + PartitionsK, AccumulatorsInRowMajor> { + + // Uses BF16 internally + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + GemmShape<16, 8, 8>, + 32, + bfloat16_t, mctlass::layout::RowMajor, + bfloat16_t, mctlass::layout::ColumnMajor, + float, mctlass::layout::RowMajor, + arch::OpMultiplyAdd + >, + mctlass::MatrixShape<1, 1> >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaTensorOp< + WarpShape_, float, LayoutA, float, LayoutB, float, LayoutC, + Policy, PartitionsK, AccumulatorsInRowMajor>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial Specialization - inputs and output types are float - uses F16 internally +template < + /// Shape of one matrix production operation (concept: GemmShape) + typename WarpShape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Number of partitions along K dimension + int PartitionsK, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor> +struct DefaultMmaTensorOp< + WarpShape_, + GemmShape<16, 8, 8>, + float, LayoutA, + float, LayoutB, + float, LayoutC, + arch::OpMultiplyAddFastF16, + PartitionsK, AccumulatorsInRowMajor> { + + // Uses F16 internally + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + GemmShape<16, 8, 8>, + 32, + half_t, mctlass::layout::RowMajor, + half_t, mctlass::layout::ColumnMajor, + float, mctlass::layout::RowMajor, + arch::OpMultiplyAdd + >, + mctlass::MatrixShape<1, 1> >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaTensorOp< + WarpShape_, float, LayoutA, float, LayoutB, float, LayoutC, + Policy, PartitionsK, AccumulatorsInRowMajor>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial Specialization - inputs and output types are float - uses TF32 internally +template < + /// Shape of one matrix production operation (concept: GemmShape) + typename WarpShape_, + /// Shape of target matrix multiply instruction (concept: GemmShape) + typename InstructionShape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Number of partitions along K dimension + int PartitionsK, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor> +struct DefaultMmaTensorOp< + WarpShape_, + InstructionShape_, + float, LayoutA, + float, LayoutB, + float, LayoutC, + arch::OpMultiplyAdd, PartitionsK, AccumulatorsInRowMajor> { + + // Uses TF32 internally + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + InstructionShape_, + 32, + tfloat32_t, mctlass::layout::RowMajor, + tfloat32_t, mctlass::layout::ColumnMajor, + float, mctlass::layout::RowMajor, + arch::OpMultiplyAdd + >, + mctlass::MatrixShape<1, 1> >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaTensorOp< + WarpShape_, float, LayoutA, float, LayoutB, float, LayoutC, + Policy, PartitionsK, AccumulatorsInRowMajor>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial Specialization - inputs and output types are float - uses TF32 for Fast Accurate FP32 +template < + /// Shape of one matrix production operation (concept: GemmShape) + typename WarpShape_, + /// Shape of target matrix multiply instruction (concept: GemmShape) + typename InstructionShape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Number of partitions along K dimension + int PartitionsK, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor> +struct DefaultMmaTensorOp< + WarpShape_, + InstructionShape_, + float, LayoutA, + float, LayoutB, + float, LayoutC, + arch::OpMultiplyAddFastF32, PartitionsK, AccumulatorsInRowMajor> { + + // Uses TF32 internally + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma< + InstructionShape_, + 32, + mctlass::tfloat32_t, mctlass::layout::RowMajor, + mctlass::tfloat32_t, mctlass::layout::ColumnMajor, + float, mctlass::layout::RowMajor, + arch::OpMultiplyAdd + >, + mctlass::MatrixShape<1, 1> >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaTensorOpFastF32< + WarpShape_, float, LayoutA, float, LayoutB, float, LayoutC, + Policy, PartitionsK, AccumulatorsInRowMajor>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "mctlass/gemm/warp/mma_complex_tensor_op_tile_iterator_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/default_mma_with_reduction_tensor_op.h b/csrc/mctlass/include/mctlass/gemm/warp/default_mma_with_reduction_tensor_op.h new file mode 100644 index 0000000..a188af3 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/default_mma_with_reduction_tensor_op.h @@ -0,0 +1,92 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Default warp-level GEMM operators selected by data type, size, and layouts of operands. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/warp/mma_with_reduction_tensor_op.h" + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A elements + typename ElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Data type of B elements + typename ElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Element type of C matrix + typename ElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Operator describing the tensor operation + typename Operator_, + /// Reduce operand A or B along K dimension + bool ReduceKForA_, + /// Number of partitions along K dimension + int PartitionsK = 1, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false> +struct DefaultMmaWithReductionTensorOp { + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Mma, + mctlass::MatrixShape<1, 1> >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaWithReductionTensorOp< + WarpShape_, ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, + Policy, ReduceKForA_, PartitionsK, AccumulatorsInRowMajor>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/default_mma_wmma_tensor_op.h b/csrc/mctlass/include/mctlass/gemm/warp/default_mma_wmma_tensor_op.h new file mode 100644 index 0000000..df6de2c --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/default_mma_wmma_tensor_op.h @@ -0,0 +1,130 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Default warp-level GEMM operators selected by data type, size, and layouts of operands. +*/ + +#pragma once + +#include "mctlass/arch/wmma.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) + +#include "mctlass/mctlass.h" +#include "mctlass/gemm/warp/mma_tensor_op_wmma.h" + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + ///< Size of the Gemm problem (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Operator describing the tensor operation + typename Operator_ = arch::OpMultiplyAdd, + /// Number of partitions along K dimension + int PartitionsK = 1 +> +struct DefaultMmaTensorOpWmma; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for m-by-n-by-kgroup +template < + ///< Shape of one matrix production operation (concept: GemmShape) + typename WarpShape_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Data type of A elements + typename ElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA, + /// Data type of B elements + typename ElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB, + /// Element type of C matrix + typename ElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC, + /// Operator describing the tensor operation + typename Operator_, + /// Number of partitions along K dimension + int PartitionsK> +struct DefaultMmaTensorOpWmma { + using Policy = mctlass::gemm::warp::MmaTensorOpPolicy< + mctlass::arch::Wmma< + InstructionShape_, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + Operator_>, + mctlass::MatrixShape<1, 1> >; + + // Define the warp-level tensor op + using Type = mctlass::gemm::warp::MmaTensorOpWmma< + WarpShape_, + ElementA, + LayoutA, + ElementB, + LayoutB, + ElementC, + LayoutC, + Policy, + PartitionsK>; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +#endif diff --git a/csrc/mctlass/include/mctlass/gemm/warp/layernorm_scale_bias_transform.h b/csrc/mctlass/include/mctlass/gemm/warp/layernorm_scale_bias_transform.h new file mode 100644 index 0000000..9271f68 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/layernorm_scale_bias_transform.h @@ -0,0 +1,141 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level per channel scale+bias+relu before + matrix multiply-accumulate operations targeting Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/platform/platform.h" + +#include "mctlass/numeric_conversion.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/arch/mma_sm75.h" +#include "mctlass/arch/mma_sm80.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/warp/mma_tensor_op_policy.h" + +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct LayernormScaleBiasTransform { + + using T = typename FragmentActivations::Element; + + static int const NumActivations = FragmentActivations::kElements; + static int const NumVarMean = FragmentVarMean::kElements; + static int const NumGammaBeta = FragmentGammaBeta::kElements; + static int const MmaElements = 2; + // One element has one scale and one bias + static int const MmaScaleBiasPair = 2; + // 16816 has 2 columns and 2 rows + static int const MmaCols = 2; + static int const MmaRows = 2; + + using MmaOperand = Array; + using VarMeanOperand = Array<__half2, MmaScaleBiasPair>; + using GammaBetaOperand = Array; + + MCTLASS_DEVICE + void transform(MmaOperand &activations, + VarMeanOperand const &var_mean, + GammaBetaOperand const &gamma_beta) { + +//#if defined(__MACA_ARCH__) +#if 0 + uint32_t *ptr_activations = reinterpret_cast(&activations); + uint32_t const *ptr_var_mean = reinterpret_cast(&var_mean); + uint32_t const *ptr_gamma_beta = reinterpret_cast(&gamma_beta); + + // Apply per channel scale+bias+relu if the data is not a special NaN + // (0x7eff). If it is a special NaN (0x7eff), hard code the output to 0. + + // We assumes the pair of FP16 are either both inbound or both out-of-bound. + // It requires C to be an even number. + asm volatile( + "{\n\t" + " fma.rn.f16x2 %0, %1, %2, %3;\n" + " fma.rn.f16x2 %0, %4, %0, %5;\n" + "}\n" + : "=r"(ptr_activations[0]) + : "r"(ptr_var_mean[0]), "r"(ptr_activations[0]), + "r"(ptr_var_mean[1]), + "r"(ptr_gamma_beta[0]), "r"(ptr_gamma_beta[1])); +#else + printf("layernorm_scale_bias_transform.h L112 this function cannot working correctly now.\n"); + assert(0); +#endif + } + + MCTLASS_DEVICE + void operator()(FragmentActivations &activations, + FragmentVarMean const &var_mean, + FragmentGammaBeta const &gamma_beta) { + MmaOperand *ptr_activations = reinterpret_cast(&activations); + VarMeanOperand const *ptr_var_mean = + reinterpret_cast(&var_mean); + GammaBetaOperand const *ptr_gamma_beta = + reinterpret_cast(&gamma_beta); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < (NumActivations / MmaElements); ++i) { + transform(ptr_activations[i], + ptr_var_mean[i / (MmaCols * MmaRows) * MmaRows + i % MmaRows], + ptr_gamma_beta[(i / MmaScaleBiasPair) % MmaCols]); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma.h b/csrc/mctlass/include/mctlass/gemm/warp/mma.h new file mode 100644 index 0000000..330336e --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma.h @@ -0,0 +1,61 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates exposing architecture support for warp-level multiply-add operations +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Query the number of threads per warp +template +struct WarpSize { + //static int const value = 32; + static int const value = 64; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_complex_tensor_op.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_complex_tensor_op.h new file mode 100644 index 0000000..69be516 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_complex_tensor_op.h @@ -0,0 +1,1171 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations targeting + Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" +#include "mctlass/complex.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/functional.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/arch/mma_sm75.h" +#include "mctlass/arch/mma_sm80.h" +#include "mctlass/arch/mma_sm90.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/warp/mma_tensor_op_policy.h" +#include "mctlass/gemm/warp/mma_tensor_op.h" + +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" +#include "mctlass/gemm/warp/mma_complex_tensor_op_tile_iterator_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +template < + /// Data type of real & imag members of complex numbers in the SourceFragment + typename RealElement, + /// Destination fragment required by the mma operation + typename DestinationFragment, + /// Source fragment holding complex elements + typename SourceFragment, + /// Number of mma operations performed + typename MmaIterations, + /// Shape of operand elements + typename MmaOperandShape, + /// Complex transform on A operand + ComplexTransform Transform_, + /// Operand A or Operand B + Operand Operand_, + /// Floating-point rounding style + FloatRoundStyle Round_> +struct UnpackComplexConvertAndPackForMma; + +// Partial specialization for OperandA and Congruous smem layout +template < + typename RealElement, + typename DestinationFragment, + typename SourceFragment, + typename MmaIterations, + typename MmaOperandShape, + ComplexTransform Transform_, + FloatRoundStyle Round_> +struct UnpackComplexConvertAndPackForMma < + RealElement, + DestinationFragment, + SourceFragment, + MmaIterations, + MmaOperandShape, + Transform_, + Operand::kA, + Round_> { + + // + // Type definitions + // + static Operand const kOperand = Operand::kA; + static ComplexTransform const kTransform = Transform_; + static FloatRoundStyle const kRound = Round_; + + // Data type of elements in the destination fragment + using MmaElement = typename DestinationFragment::Element; + + // Numeric convertor MmaElement <= RealElement + using Converter = NumericConverter; + + // Operand layout parameters + using SourceFragmentLayout = layout::ColumnMajor; + static int const kLdm = MmaIterations::kRow * MmaOperandShape::kRow; + + /// Ctor + MCTLASS_DEVICE + UnpackComplexConvertAndPackForMma() {} + + MCTLASS_DEVICE + void operator()(DestinationFragment *dest, SourceFragment const &source) { + + Converter convert_op; + SourceFragmentLayout layout(kLdm); + + MCTLASS_PRAGMA_UNROLL + for(int i=0; i and apply rounding on real and imag parts + MmaElement a = convert_op(source[layout(MatrixCoord{row,col})].real()); + MmaElement b = convert_op(source[layout(MatrixCoord{row,col})].imag()); + + // Unpack rounded complex and pack into DestinationFragment for mma operation + dest[i][pos] = a; + dest[i+MmaIterations::kRow][pos++] = (kTransform == ComplexTransform::kConjugate ? -b : b); + + } + } + } + } +}; + +// Partial specialization for OperandB and Congruous smem layout +template < + typename RealElement, + typename DestinationFragment, + typename SourceFragment, + typename MmaIterations, + typename MmaOperandShape, + ComplexTransform Transform_, + FloatRoundStyle Round_> +struct UnpackComplexConvertAndPackForMma < + RealElement, + DestinationFragment, + SourceFragment, + MmaIterations, + MmaOperandShape, + Transform_, + Operand::kB, + Round_> { + + // + // Type definitions + // + static Operand const kOperand = Operand::kB; + static ComplexTransform const kTransform = Transform_; + static FloatRoundStyle const kRound = Round_; + + // Data type of elements in the destination fragment + using MmaElement = typename DestinationFragment::Element; + + // Numeric convertor MmaElement <= RealElement + using Converter = NumericConverter; + + // Operand layout parameters + using SourceFragmentLayout = layout::RowMajor; + static int const kLdm = MmaIterations::kColumn * MmaOperandShape::kColumn; + + /// Ctor + MCTLASS_DEVICE + UnpackComplexConvertAndPackForMma() {} + + MCTLASS_HOST_DEVICE + void operator()(DestinationFragment *dest, SourceFragment const &source) { + + Converter convert_op; + SourceFragmentLayout layout(kLdm); + + MCTLASS_PRAGMA_UNROLL + for(int i=0; i apply rounding on real and imag parts + MmaElement a = convert_op(source[layout(MatrixCoord{row,col})].real()); + MmaElement b = convert_op(source[layout(MatrixCoord{row,col})].imag()); + + // Unpack rounded complex and pack into DestinationFragment for mma operation + dest[i][pos] = a; + dest[i+MmaIterations::kColumn][pos++] = (kTransform == ComplexTransform::kConjugate ? -b : b); + } + } + } + } +}; +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename RealElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename RealElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename RealElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Complex transform on A operand + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex transform on B operand + ComplexTransform TransformB = ComplexTransform::kNone, + /// Do source operands need more than one elements + bool GeneralizedOperatorElements = false, + /// Used for partial specialization + typename Enable = bool +> +class MmaComplexTensorOp; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex*complex+complex => complex using real-valued TensorOps +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename RealElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename RealElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename RealElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB +> +class MmaComplexTensorOp< + Shape_, + complex, + LayoutA_, + complex, + LayoutB_, + complex, + LayoutC_, + Policy_, + TransformA, + TransformB> { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of multiplicand A + using ElementA = complex; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of multiplicand B + using ElementB = complex; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of accumulator matrix C + using ElementC = complex; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicyTensorOp) + using Policy = Policy_; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Policy::Operator; + + /// Architecture tag from underlying instruction + using ArchTag = typename ArchMmaOperator::ArchTag; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassTensorOp; + + /// Shape of underlying instruction + using InstructionShape = typename ArchMmaOperator::Shape; + + /// Indicates math operator + using MathOperator = arch::OpMultiplyAddComplex; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = TransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = TransformB; + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 32; + +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kA, + ElementA, + LayoutA, + MatrixShape, + Policy::OpDelta::kRow, + 32, + 1 + >; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentA = FragmentA; + + /// Iterates over the B operand in memory + using IteratorB = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kB, + ElementB, + LayoutB, + MatrixShape, + Policy::OpDelta::kColumn, + 32, + 1 + >; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed B tile + using TransformedFragmentB = FragmentB; + + static_assert( + !(Shape::kM % ArchMmaOperator::Shape::kM) && + !(Shape::kN % ArchMmaOperator::Shape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + Shape::kM / ArchMmaOperator::Shape::kM, + Shape::kN / ArchMmaOperator::Shape::kN + >; + + /// Iterates over the C operand in memory + using IteratorC = MmaTensorOpAccumulatorTileIterator< + MatrixShape, + ElementC, + LayoutC, + typename ArchMmaOperator::Shape, + typename Policy::OpDelta>; + + /// Storage for C tile, the accumulator. Note, regardless of multiplicand type, this + /// storage arrangement is to be considered 'planar complex' in the sense that all real-valued + /// parts are stored consecutively followed by all imaginary parts. This matches the structure + /// of Tensor Cores which are always real-valued matrix multiplies. + using FragmentC = typename IteratorC::Fragment; + + static_assert( + FragmentC::kElements == 2 * MmaIterations::kCount * ArchMmaOperator::FragmentC::kElements, + "Unexpected planar complex fragment length."); + +private: + + // + // Data members + // + + /// Underlying real-valued matrix multiply operator (concept: arch::Mma) + ArchMmaOperator mma; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaComplexTensorOp() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + FragmentA const &A, + FragmentB const &B, + FragmentC const &C + ) const { + + // Alias types for underlying real-valued matrix multiply operator + using MmaOperandA = typename ArchMmaOperator::FragmentA; + using MmaOperandB = typename ArchMmaOperator::FragmentB; + using MmaOperandC = typename ArchMmaOperator::FragmentC; + + static_assert(MmaOperandA::kElements == 1, + "This implementation only supports math instructions in which exactly one element is needed for the A operand." + "We can geneneralize later."); + + static_assert(MmaOperandB::kElements == 1, + "This implementation only supports math instructions in which exactly one element is needed for the B operand." + "We can geneneralize later."); + + D = C; + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < MmaIterations::kRow; ++m) { + + // mma(accum.real(), a.real(), b.real(), accum.real()); + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_A; + MmaOperandB operand_B; + + operand_A[0] = A[m].real(); + operand_B[0] = B[n].real(); + + // Real-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow); + + mma(*accum, operand_A, operand_B, *accum); + } + + // mma(accum.imag(), a.real(), b.imag(), accum.imag()); + MCTLASS_PRAGMA_UNROLL + for (int n = MmaIterations::kColumn - 1; n >= 0; --n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_A; + MmaOperandB operand_B; + + operand_A[0] = A[m].real(); + operand_B[0] = (kTransformB == ComplexTransform::kConjugate ? -B[n].imag() : B[n].imag()); + + // Complex-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + MmaIterations::kCount; + + mma(*accum, operand_A, operand_B, *accum); + } + + // mma(accum.real(), -a.imag(), b.imag(), accum.real()) + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_A; + MmaOperandB operand_B; + + // A imaginary part is intentionally negated + operand_A[0] = (kTransformA == ComplexTransform::kConjugate ? A[m].imag() : -A[m].imag()); + operand_B[0] = (kTransformB == ComplexTransform::kConjugate ? -B[n].imag() : B[n].imag()); + + // Real-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow); + + mma(*accum, operand_A, operand_B, *accum); + } + + // mma(accum.imag(), a.imag(), b.real(), accum.imag()) + MCTLASS_PRAGMA_UNROLL + for (int n = MmaIterations::kColumn - 1; n >= 0; --n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_A; + MmaOperandB operand_B; + + operand_A[0] = (kTransformA == ComplexTransform::kConjugate ? -A[m].imag() : A[m].imag()); + operand_B[0] = B[n].real(); + + // Complex-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + MmaIterations::kCount; + + mma(*accum, operand_A, operand_B, *accum); + } + } + } + + /// Transform the mma operands to the required types + MCTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + dst_A = A; + dst_B = B; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex*complex+complex => complex: +// Operands data type: complex +// Rounding: float -> tfloat32_t (round half_ulp_truncate nearest) +// Math instruction: mma.sync.aligned.m16n8k8.f32.tf32.tf32.f32 +// Output data type: complex +// +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB +> +class MmaComplexTensorOp< + Shape_, + complex, + LayoutA_, + complex, + LayoutB_, + complex, + LayoutC_, + Policy_, + TransformA, + TransformB> { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of members of complex multiplicand A + using RealElementA = float; + + /// Data type of multiplicand A + using ElementA = complex; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of members of complex multiplicand B + using RealElementB = float; + + /// Data type of multiplicand B + using ElementB = complex; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of members of complex accumulator matrix C + using RealElementC = float; + + /// Data type of accumulator matrix C + using ElementC = complex; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Policy::Operator; + + /// Shape of underlying instruction + using InstructionShape = typename ArchMmaOperator::Shape; + + /// Underlying arch tag + using ArchTag = typename ArchMmaOperator::ArchTag; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassTensorOp; + + /// Indicates math operator + using MathOperator = typename arch::OpMultiplyAddComplex; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = TransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = TransformB; + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 32; + +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kA, + ElementA, + LayoutA, + MatrixShape, + Policy::OpDelta::kRow, + 32, + 1 + >; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentA = + Array; + + /// Iterates over the B operand in memory + using IteratorB = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kB, + ElementB, + LayoutB, + MatrixShape, + Policy::OpDelta::kColumn, + 32, + 1 + >; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed B tile + using TransformedFragmentB = + Array; + + static_assert( + !(Shape::kM % ArchMmaOperator::Shape::kM) && + !(Shape::kN % ArchMmaOperator::Shape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + + /// Number of complex products operations performed (one complex product needs four mma instructions) + using MmaIterations = MatrixShape< + Shape::kM / ArchMmaOperator::Shape::kM, + Shape::kN / ArchMmaOperator::Shape::kN + >; + + /// Iterates over the C operand in memory + using IteratorC = MmaTensorOpAccumulatorTileIterator< + MatrixShape, + ElementC, + LayoutC, + typename ArchMmaOperator::Shape, + typename Policy::OpDelta>; + + /// Storage for C tile, the accumulator. Note, regardless of multiplicand type, this + /// storage arrangement is to be considered 'planar complex' in the sense that all real-valued + /// parts are stored consecutively followed by all imaginary parts. This matches the structure + /// of Tensor Cores which are always real-valued matrix multiplies. + using FragmentC = typename IteratorC::Fragment; + +private: + + // + // Data members + // + + /// Underlying real-valued matrix multiply operator (concept: arch::Mma) + ArchMmaOperator mma; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaComplexTensorOp() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + TransformedFragmentA const &A, + TransformedFragmentB const &B, + FragmentC const &C + ) const { + + // Alias types for underlying real-valued matrix multiply operator + using InstMmaOperandA = typename ArchMmaOperator::FragmentA; + using InstMmaOperandB = typename ArchMmaOperator::FragmentB; + using MmaOperandC = typename ArchMmaOperator::FragmentC; + + static_assert(platform::is_same, typename ArchMmaOperator::Shape>::value, + "This implementation only supports mma.m16n8k8 math instructions."); + + static_assert(InstMmaOperandA::kElements == 4, + "This implementation only supports math instructions in which exactly four element is needed for the A operand." + "We can geneneralize later."); + + static_assert(InstMmaOperandB::kElements == 2, + "This implementation only supports math instructions in which exactly two element is needed for the B operand." + "We can geneneralize later."); + + // Instruction Operands A & B holding real part followed by imaginary part for mma operations + InstMmaOperandA const *operand_A = reinterpret_cast(&A); + InstMmaOperandB const *operand_B = reinterpret_cast(&B); + + // + // Accumulate in place + // + D = C; + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < MmaIterations::kRow; ++m) { + + // mma(accum.real(), a.real(), b.real(), accum.real()); + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // Real-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow); + + mma(*accum, operand_A[m], operand_B[n], *accum); + } + + // mma(accum.imag(), a.real(), b.imag(), accum.imag()); + MCTLASS_PRAGMA_UNROLL + for (int n = MmaIterations::kColumn - 1; n >= 0; --n) { + + // Complex-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + MmaIterations::kCount; + + mma(*accum, operand_A[m], operand_B[n+MmaIterations::kColumn], *accum); + } + + // mma(accum.real(), a.imag(), -b.imag(), accum.real()) + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // negate OperandB to accumulate -(a.imag()*b.imag()) + // negating OperandB emits less instrucitons than negating OperandA as OperandB has less elements + negate negate_op; + + // Real-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow); + + mma(*accum, operand_A[m+MmaIterations::kRow], negate_op(operand_B[n+MmaIterations::kColumn]), *accum); + } + + // mma(accum.imag(), a.imag(), b.real(), accum.imag()) + MCTLASS_PRAGMA_UNROLL + for (int n = MmaIterations::kColumn - 1; n >= 0; --n) { + + // Complex-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + MmaIterations::kCount; + + mma(*accum, operand_A[m+MmaIterations::kRow], operand_B[n], *accum); + } + } + } + + /// Transform the mma operands to the required types + MCTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + // Alias types for underlying real-valued matrix multiply operator + using InstMmaOperandA = typename ArchMmaOperator::FragmentA; + using InstMmaOperandB = typename ArchMmaOperator::FragmentB; + + // + // Define conversions from source type to instruction operands' type + // + + // #if defined(__MACA_ARCH__) && __MACA_ARCH__ >= 900 + #if 0 + FloatRoundStyle const kRoundA = FloatRoundStyle::round_to_nearest; + FloatRoundStyle const kRoundB = FloatRoundStyle::round_to_nearest; + #else + FloatRoundStyle const kRoundA = FloatRoundStyle::round_half_ulp_trunc_dntz; + FloatRoundStyle const kRoundB = FloatRoundStyle::round_half_ulp_trunc_dntz; + #endif + + detail::UnpackComplexConvertAndPackForMma < + RealElementA, + InstMmaOperandA, + FragmentA, + MmaIterations, + MatrixShape<2, 2>, + kTransformA, + Operand::kA, + kRoundA> convert_A; + + detail::UnpackComplexConvertAndPackForMma < + RealElementB, + InstMmaOperandB, + FragmentB, + MmaIterations, + MatrixShape<2, 1>, + kTransformB, + Operand::kB, + kRoundB> convert_B; + + // Convert Fragment[A|B] holding complex to InstMmaOperand[A|B] holding InstMmaOperand[A|B]::Element + convert_A(reinterpret_cast(&dst_A), A); + convert_B(reinterpret_cast(&dst_B), B); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Partial specialization for complex*complex+complex => complex: +// Operands data type: complex +// Math instruction: mma.sync.aligned.m16n8k4.f64.f64.f64.f64 +// Output data type: complex +// +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB +> +class MmaComplexTensorOp< + Shape_, + complex, + LayoutA_, + complex, + LayoutB_, + complex, + LayoutC_, + Policy_, + TransformA, + TransformB, + true> { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of members of complex multiplicand A + using RealElementA = double; + + /// Data type of multiplicand A + using ElementA = complex; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of members of complex multiplicand B + using RealElementB = double; + + /// Data type of multiplicand B + using ElementB = complex; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of members of complex accumulator matrix C + using RealElementC = double; + + /// Data type of accumulator matrix C + using ElementC = complex; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicyTensorOp) + using Policy = Policy_; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Policy::Operator; + + /// Shape of underlying instruction + using InstructionShape = typename ArchMmaOperator::Shape; + + /// Underlying arch tag + using ArchTag = typename ArchMmaOperator::ArchTag; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassTensorOp; + + /// Indicates math operator + using MathOperator = typename arch::OpMultiplyAddComplex; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = TransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = TransformB; + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 32; + +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kA, + ElementA, + LayoutA, + MatrixShape, + Policy::OpDelta::kRow, + 32, + 1 + >; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentA = FragmentA; + + /// Iterates over the B operand in memory + using IteratorB = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kB, + ElementB, + LayoutB, + MatrixShape, + Policy::OpDelta::kColumn, + 32, + 1 + >; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed B tile + using TransformedFragmentB = FragmentB; + + static_assert( + !(Shape::kM % ArchMmaOperator::Shape::kM) && + !(Shape::kN % ArchMmaOperator::Shape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + Shape::kM / ArchMmaOperator::Shape::kM, + Shape::kN / ArchMmaOperator::Shape::kN + >; + + /// Iterates over the C operand in memory + using IteratorC = MmaTensorOpAccumulatorTileIterator< + MatrixShape, + ElementC, + LayoutC, + typename ArchMmaOperator::Shape, + typename Policy::OpDelta>; + + /// Storage for C tile, the accumulator. Note, regardless of multiplicand type, this + /// storage arrangement is to be considered 'planar complex' in the sense that all real-valued + /// parts are stored consecutively followed by all imaginary parts. This matches the structure + /// of Tensor Cores which are always real-valued matrix multiplies. + using FragmentC = typename IteratorC::Fragment; + + static_assert( + FragmentC::kElements == 2 * MmaIterations::kCount * ArchMmaOperator::FragmentC::kElements, + "Unexpected planar complex fragment length."); + +private: + + // + // Data members + // + + /// Underlying real-valued matrix multiply operator (concept: arch::Mma) + ArchMmaOperator mma; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaComplexTensorOp() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + FragmentA const &A, + FragmentB const &B, + FragmentC const &C + ) const { + + // Alias types for underlying real-valued matrix multiply operator + using MmaOperandA = typename ArchMmaOperator::FragmentA; + using MmaOperandB = typename ArchMmaOperator::FragmentB; + using MmaOperandC = typename ArchMmaOperator::FragmentC; + + D = C; + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < MmaIterations::kRow; ++m) { + + // mma(accum.real(), a.real(), b.real(), accum.real()); + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_A; + MmaOperandB operand_B; + + MCTLASS_PRAGMA_UNROLL + for (int mk = 0; mk < MmaOperandA::kElements; ++mk) + operand_A[mk] = A[m*MmaOperandA::kElements + mk].real(); + + MCTLASS_PRAGMA_UNROLL + for (int nk = 0; nk < MmaOperandB::kElements; ++nk) + operand_B[nk] = B[n*MmaOperandB::kElements + nk].real(); + + // Real-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow); + + mma(*accum, operand_A, operand_B, *accum); + } + + // mma(accum.imag(), a.real(), b.imag(), accum.imag()); + MCTLASS_PRAGMA_UNROLL + for (int n = MmaIterations::kColumn - 1; n >= 0; --n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_A; + MmaOperandB operand_B; + + MCTLASS_PRAGMA_UNROLL + for (int mk = 0; mk < MmaOperandA::kElements; ++mk) + operand_A[mk] = A[m*MmaOperandA::kElements + mk].real(); + + MCTLASS_PRAGMA_UNROLL + for (int nk = 0; nk < MmaOperandB::kElements; ++nk) + operand_B[nk] = (kTransformB == ComplexTransform::kConjugate ? + -B[n*MmaOperandB::kElements + nk].imag() : B[n*MmaOperandB::kElements + nk].imag()); + + // Complex-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + MmaIterations::kCount; + + mma(*accum, operand_A, operand_B, *accum); + } + + // mma(accum.real(), -a.imag(), b.imag(), accum.real()) + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_A; + MmaOperandB operand_B; + + // A imaginary part is intentionally negated + MCTLASS_PRAGMA_UNROLL + for (int mk = 0; mk < MmaOperandA::kElements; ++mk) + operand_A[mk] = (kTransformA == ComplexTransform::kConjugate ? + A[m*MmaOperandA::kElements + mk].imag() : -A[m*MmaOperandA::kElements + mk].imag()); + + MCTLASS_PRAGMA_UNROLL + for (int nk = 0; nk < MmaOperandB::kElements; ++nk) + operand_B[nk] = (kTransformB == ComplexTransform::kConjugate ? + -B[n*MmaOperandB::kElements + nk].imag() : B[n*MmaOperandB::kElements + nk].imag()); + + // Real-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow); + + mma(*accum, operand_A, operand_B, *accum); + } + + // mma(accum.imag(), a.imag(), b.real(), accum.imag()) + MCTLASS_PRAGMA_UNROLL + for (int n = MmaIterations::kColumn - 1; n >= 0; --n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_A; + MmaOperandB operand_B; + + MCTLASS_PRAGMA_UNROLL + for (int mk = 0; mk < MmaOperandA::kElements; ++mk) + operand_A[mk] = (kTransformA == ComplexTransform::kConjugate ? + -A[m*MmaOperandA::kElements + mk].imag() : A[m*MmaOperandA::kElements + mk].imag()); + + MCTLASS_PRAGMA_UNROLL + for (int nk = 0; nk < MmaOperandB::kElements; ++nk) + operand_B[nk] = B[n*MmaOperandB::kElements + nk].real(); + + // Complex-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + MmaIterations::kCount; + + mma(*accum, operand_A, operand_B, *accum); + } + } + } + + /// Transform the mma operands to the required types + MCTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + dst_A = A; + dst_B = B; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_complex_tensor_op_fast_f32.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_complex_tensor_op_fast_f32.h new file mode 100644 index 0000000..a326922 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_complex_tensor_op_fast_f32.h @@ -0,0 +1,663 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations targeting + Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" +#include "mctlass/complex.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/functional.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/arch/mma_sm75.h" +#include "mctlass/arch/mma_sm80.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/warp/mma_tensor_op_policy.h" +#include "mctlass/gemm/warp/mma_tensor_op.h" + +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" +#include "mctlass/gemm/warp/mma_complex_tensor_op_tile_iterator_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +namespace detail { + +template < + /// Data type of real & imag members of complex numbers in the SourceFragment + typename RealElement, + /// Destination fragment required by the mma operation + typename DestinationFragment, + /// Source fragment holding complex elements + typename SourceFragment, + /// Number of mma operations performed + typename MmaIterations, + /// Shape of operand elements + typename MmaOperandShape, + /// Complex transform on A operand + ComplexTransform Transform_, + /// Operand A or Operand B + Operand Operand_, + /// Floating-point rounding style for big part + FloatRoundStyle RoundBig_, + /// Floating-point rounding style for small part + FloatRoundStyle RoundSmall_> +struct UnpackComplexConvertAndPackForMmaFastF32; + +// Partial specialization for OperandA and Congruous smem layout +template < + typename RealElement, + typename DestinationFragment, + typename SourceFragment, + typename MmaIterations, + typename MmaOperandShape, + ComplexTransform Transform_, + FloatRoundStyle RoundBig_, + FloatRoundStyle RoundSmall_> +struct UnpackComplexConvertAndPackForMmaFastF32 < + RealElement, + DestinationFragment, + SourceFragment, + MmaIterations, + MmaOperandShape, + Transform_, + Operand::kA, + RoundBig_, + RoundSmall_> { + + // + // Type definitions + // + static Operand const kOperand = Operand::kA; + static ComplexTransform const kTransform = Transform_; + static FloatRoundStyle const kRoundBig = RoundBig_; + static FloatRoundStyle const kRoundSmall = RoundSmall_; + + // Data type of elements in the destination fragment + using MmaElement = typename DestinationFragment::Element; + + // Numeric convertor MmaElementBig, MmaElementSmall <= RealElement + using Converter = NumericConverterFastF32; + + // Operand layout parameters + using SourceFragmentLayout = layout::ColumnMajor; + static int const kLdm = MmaIterations::kRow * MmaOperandShape::kRow; + + // BigSmall Fragment holding two TF32 elements (big, small) for every float + using BigSmallFragment = Array; + + /// Index in fargments for the big and small part + static int const kBigIndex = 0; + static int const kSmallIndex = 1; + + /// Ctor + MCTLASS_DEVICE + UnpackComplexConvertAndPackForMmaFastF32() {} + + MCTLASS_DEVICE + void operator()(DestinationFragment *dest, SourceFragment const &source) { + + Converter convert_op; + SourceFragmentLayout layout(kLdm); + + DestinationFragment *dest_big_ = reinterpret_cast(dest); + DestinationFragment *dest_small_ = reinterpret_cast(&dest[MmaIterations::kRow * 2]); + + MCTLASS_PRAGMA_UNROLL + for(int i=0; i and apply rounding on real and imag parts + BigSmallFragment a = convert_op(source[layout(MatrixCoord{row,col})].real()); + BigSmallFragment b = convert_op(source[layout(MatrixCoord{row,col})].imag()); + + // Unpack rounded complex and pack into DestinationFragment for mma operation + dest_big_[i][pos] = a[kBigIndex]; + dest_big_[i+MmaIterations::kRow][pos] = (kTransform == ComplexTransform::kConjugate ? -b[kBigIndex] : b[kBigIndex]); + + // Unpack rounded complex and pack into DestinationFragment for mma operation + dest_small_[i][pos] = a[kSmallIndex]; + dest_small_[i+MmaIterations::kRow][pos] = (kTransform == ComplexTransform::kConjugate ? -b[kSmallIndex] : b[kSmallIndex]); + + // Next position + pos++; + } + } + } + } +}; + +// Partial specialization for OperandB and Congruous smem layout +template < + typename RealElement, + typename DestinationFragment, + typename SourceFragment, + typename MmaIterations, + typename MmaOperandShape, + ComplexTransform Transform_, + FloatRoundStyle RoundBig_, + FloatRoundStyle RoundSmall_> +struct UnpackComplexConvertAndPackForMmaFastF32 < + RealElement, + DestinationFragment, + SourceFragment, + MmaIterations, + MmaOperandShape, + Transform_, + Operand::kB, + RoundBig_, + RoundSmall_> { + + // + // Type definitions + // + static Operand const kOperand = Operand::kB; + static ComplexTransform const kTransform = Transform_; + static FloatRoundStyle const kRoundBig = RoundBig_; + static FloatRoundStyle const kRoundSmall = RoundSmall_; + + // Data type of elements in the destination fragment + using MmaElement = typename DestinationFragment::Element; + + // Numeric convertor MmaElementBig, MmaElementSmall <= RealElement + using Converter = NumericConverterFastF32; + + // Operand layout parameters + using SourceFragmentLayout = layout::RowMajor; + static int const kLdm = MmaIterations::kColumn * MmaOperandShape::kColumn; + + // BigSmall Fragment holding two TF32 elements (big, small) for every float + using BigSmallFragment = Array; + + /// Index in fargments for the big and small part + static int const kBigIndex = 0; + static int const kSmallIndex = 1; + + /// Ctor + MCTLASS_DEVICE + UnpackComplexConvertAndPackForMmaFastF32() {} + + MCTLASS_HOST_DEVICE + void operator()(DestinationFragment *dest, SourceFragment const &source) { + + Converter convert_op; + SourceFragmentLayout layout(kLdm); + + DestinationFragment *dest_big_ = reinterpret_cast(dest); + DestinationFragment *dest_small_ = reinterpret_cast(&dest[MmaIterations::kColumn * 2]); + + MCTLASS_PRAGMA_UNROLL + for(int i=0; i apply rounding on real and imag parts + BigSmallFragment a = convert_op(source[layout(MatrixCoord{row,col})].real()); + BigSmallFragment b = convert_op(source[layout(MatrixCoord{row,col})].imag()); + + // Unpack rounded complex and pack into DestinationFragment for mma operation + dest_big_[i][pos] = a[kBigIndex]; + dest_big_[i+MmaIterations::kColumn][pos] = (kTransform == ComplexTransform::kConjugate ? -b[kBigIndex] : b[kBigIndex]); + + // Unpack rounded complex and pack into DestinationFragment for mma operation + dest_small_[i][pos] = a[kSmallIndex]; + dest_small_[i+MmaIterations::kColumn][pos] = (kTransform == ComplexTransform::kConjugate ? -b[kSmallIndex] : b[kSmallIndex]); + + // next position + pos++; + } + } + } + } +}; +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename RealElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename RealElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename RealElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Complex transform on A operand + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex transform on B operand + ComplexTransform TransformB = ComplexTransform::kNone, + /// Used for partial specialization + typename Enable = bool +> +class MmaComplexTensorOpFastF32; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex*complex+complex => complex: +// Operands data type: complex +// Rounding: float -> tfloat32_t (round half_ulp_truncate nearest) +// Math instruction: mma.sync.aligned.m16n8k8.f32.tf32.tf32.f32 +// Output data type: complex +// +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB, + /// Used for partial specialization + typename Enable +> +class MmaComplexTensorOpFastF32< + Shape_, + complex, + LayoutA_, + complex, + LayoutB_, + complex, + LayoutC_, + Policy_, + TransformA, + TransformB, + Enable> { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of members of complex multiplicand A + using RealElementA = float; + + /// Data type of multiplicand A + using ElementA = complex; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of members of complex multiplicand B + using RealElementB = float; + + /// Data type of multiplicand B + using ElementB = complex; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of members of complex accumulator matrix C + using RealElementC = float; + + /// Data type of accumulator matrix C + using ElementC = complex; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Policy::Operator; + + /// Shape of underlying instruction + using InstructionShape = typename ArchMmaOperator::Shape; + + /// Underlying arch tag + using ArchTag = typename ArchMmaOperator::ArchTag; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassTensorOp; + + /// Indicates math operator + using MathOperator = arch::OpMultiplyAddComplexFastF32; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = TransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = TransformB; + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 32; + + + /// Tune F32 to TF32 big small conversion for complex operation + /// Different combination of big small conversin can cause different tradeoff + /// between speed and accuracy. Generally, use round_half_ulp_truncate can + /// improve the performance but hur the accuracy. + using ComplexFastF32 = FastF32 < + FloatRoundStyle::round_toward_zero, // kRoundBigA + FloatRoundStyle::round_half_ulp_truncate, // kRoundSmallA + FloatRoundStyle::round_toward_zero, // kRoundBigB + FloatRoundStyle::round_half_ulp_truncate, // kRoundSmallB + TensorFloat32Op::k3xTF32 // Number of TF32 operations + >; + + /// Index in fargments for the big and small part + static int const kBigIndex = 0; + static int const kSmallIndex = 1; + +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kA, + ElementA, + LayoutA, + MatrixShape, + Policy::OpDelta::kRow, + 32, + 1 + >; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + // (4 times the original FragmentA::kElements) + // (real_big), (imag_big), (real_small), (imag_small) + using TransformedFragmentA = Array; + + // Fragment bisecting big and small sections + // (real_big, imag_big), (real_small, imag_small) + using AccessTypeFragmentA = Array; + + /// Iterates over the B operand in memory + using IteratorB = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kB, + ElementB, + LayoutB, + MatrixShape, + Policy::OpDelta::kColumn, + 32, + 1 + >; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed B tile + // (4 times the original FragmentB::kElements) + // (real_big), (imag_big), (real_small), (imag_small) + using TransformedFragmentB = Array; + + // Fragment bisecting big and small sections + // (real_big, imag_big), (real_small, imag_small) + using AccessTypeFragmentB = Array; + + static_assert( + !(Shape::kM % ArchMmaOperator::Shape::kM) && + !(Shape::kN % ArchMmaOperator::Shape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + + /// Number of complex products operations performed (one complex product needs four mma instructions) + using MmaIterations = MatrixShape< + Shape::kM / ArchMmaOperator::Shape::kM, + Shape::kN / ArchMmaOperator::Shape::kN + >; + + /// Iterates over the C operand in memory + using IteratorC = MmaTensorOpAccumulatorTileIterator< + MatrixShape, + ElementC, + LayoutC, + typename ArchMmaOperator::Shape, + typename Policy::OpDelta>; + + /// Storage for C tile, the accumulator. Note, regardless of multiplicand type, this + /// storage arrangement is to be considered 'planar complex' in the sense that all real-valued + /// parts are stored consecutively followed by all imaginary parts. This matches the structure + /// of Tensor Cores which are always real-valued matrix multiplies. + using FragmentC = typename IteratorC::Fragment; + + // + // Alias types for underlying real-valued matrix multiply operator + // + using InstMmaOperandA = typename ArchMmaOperator::FragmentA; + using InstMmaOperandB = typename ArchMmaOperator::FragmentB; + using MmaOperandC = typename ArchMmaOperator::FragmentC; + + static_assert(platform::is_same, typename ArchMmaOperator::Shape>::value, + "This implementation only supports mma.m16n8k8 math instructions."); + + static_assert(InstMmaOperandA::kElements == 4, + "This implementation only supports math instructions in which exactly four element is needed for the A operand." + "We can geneneralize later."); + + static_assert(InstMmaOperandB::kElements == 2, + "This implementation only supports math instructions in which exactly two element is needed for the B operand." + "We can geneneralize later."); + +private: + + // + // Data members + // + + /// Underlying real-valued matrix multiply operator (concept: arch::Mma) + ArchMmaOperator mma; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaComplexTensorOpFastF32() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + TransformedFragmentA const &A, + TransformedFragmentB const &B, + FragmentC const &C + ) const { + + AccessTypeFragmentA const *complex_A = reinterpret_cast(&A); + AccessTypeFragmentB const *complex_B = reinterpret_cast(&B); + + // + // Accumulate in place + // + D = C; + + + complex_mma_operator(D, complex_A[kSmallIndex], complex_B[kBigIndex], D); + + complex_mma_operator(D, complex_A[kBigIndex], complex_B[kSmallIndex], D); + + complex_mma_operator(D, complex_A[kBigIndex], complex_B[kBigIndex], D); + + if (ComplexFastF32::kPrecision == TensorFloat32Op::k4xTF32) + complex_mma_operator(D, complex_A[kSmallIndex], complex_B[kSmallIndex], D); + } + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void complex_mma_operator( + FragmentC &D, + AccessTypeFragmentA const &complex_A, + AccessTypeFragmentB const &complex_B, + FragmentC const &C + ) const { + + // Instruction Operands A & B holding real part followed by imaginary part for mma operations + InstMmaOperandA const *operand_A = reinterpret_cast(&complex_A); + InstMmaOperandB const *operand_B = reinterpret_cast(&complex_B); + + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < MmaIterations::kRow; ++m) { + + // mma(accum.real(), a.real(), b.real(), accum.real()); + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // Real-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow); + + mma(*accum, operand_A[m], operand_B[n], *accum); + } + + // mma(accum.imag(), a.real(), b.imag(), accum.imag()); + MCTLASS_PRAGMA_UNROLL + for (int n = MmaIterations::kColumn - 1; n >= 0; --n) { + + // Complex-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + MmaIterations::kCount; + + mma(*accum, operand_A[m], operand_B[n+MmaIterations::kColumn], *accum); + } + + // mma(accum.real(), a.imag(), -b.imag(), accum.real()) + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // negate OperandB to accumulate -(a.imag()*b.imag()) + // negating OperandB emits less instrucitons than negating OperandA as OperandB has less elements + negate negate_op; + + // Real-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow); + + mma(*accum, operand_A[m+MmaIterations::kRow], negate_op(operand_B[n+MmaIterations::kColumn]), *accum); + } + + // mma(accum.imag(), a.imag(), b.real(), accum.imag()) + MCTLASS_PRAGMA_UNROLL + for (int n = MmaIterations::kColumn - 1; n >= 0; --n) { + + // Complex-valued accumulator part + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + MmaIterations::kCount; + + mma(*accum, operand_A[m+MmaIterations::kRow], operand_B[n], *accum); + } + } + } + + /// Transform the mma operands to the required types + MCTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + + detail::UnpackComplexConvertAndPackForMmaFastF32 < + RealElementA, + InstMmaOperandA, + FragmentA, + MmaIterations, + MatrixShape<2, 2>, + kTransformA, + Operand::kA, + ComplexFastF32::kRoundBigA, + ComplexFastF32::kRoundSmallA> convert_A; + + detail::UnpackComplexConvertAndPackForMmaFastF32 < + RealElementB, + InstMmaOperandB, + FragmentB, + MmaIterations, + MatrixShape<2, 1>, + kTransformB, + Operand::kB, + ComplexFastF32::kRoundBigB, + ComplexFastF32::kRoundSmallB> convert_B; + + // Convert Fragment[A|B] holding complex to InstMmaOperand[A|B] holding InstMmaOperand[A|B]::Element + convert_A(reinterpret_cast(&dst_A), A); + convert_B(reinterpret_cast(&dst_B), B); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_complex_tensor_op_tile_iterator_sm80.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_complex_tensor_op_tile_iterator_sm80.h new file mode 100644 index 0000000..e524714 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_complex_tensor_op_tile_iterator_sm80.h @@ -0,0 +1,2661 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines iterators used by warp-level matrix multiply operations targeting Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/gemm/gemm.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor_op_multiplicand_sm80.h" + +#include "mctlass/platform/platform.h" +#include "mctlass/fast_math.h" + +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for loading 128b vectors of 128b elements. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::TensorOpMultiplicandCongruous128b, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + static_assert(!(Shape::kContiguous % 8) && !(Shape::kStrided % 4), "Divisibility."); + + static_assert(sizeof_bits::value == 128, "This is specialized for 128b accesses."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCongruous128b; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Load two elements per access + static int const kElementsPerAccess = 1; + + /// Policy defining internal details of tile iterator + struct Policy { + + /// Shape of one access + using Delta = layout::PitchLinearShape<8, 4>; + + /// Number of iterations to load + using Iterations = layout::PitchLinearShape< + Shape::kContiguous / Delta::kContiguous, + InstructionShape::kStrided / Delta::kStrided + >; + }; + +private: + + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = AlignedArray; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + +private: + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(): stride_(0), byte_offset_(0) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): + stride_(ref.stride(0) / kElementsPerAccess), byte_offset_(0) { + + int quad_pair = lane_id / 8; + int quad = lane_id / 4; + int lane = lane_id % 4; + + int row = (quad & 1) * 4 + (lane ^ quad_pair); + + byte_offset_ = (row + quad_pair * stride_) * sizeof(AccessType); + + pointer_= reinterpret_cast(ref.data()); + + //Special compute for mma shape with m8n8k4f64 + if (platform::is_same>::value == true) { + + const int quad_pair = lane_id & 0x7; + const int quad = (lane_id >> 4); + const int lane = quad_pair / 2; + const int row = ((quad_pair & 1) * 4) | (quad ^ lane); + byte_offset_ = (row + lane * stride_) * sizeof(AccessType); + + } + + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + pointer_ += offset; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + int offset = + (tile_offset.contiguous() * Shape::kContiguous) + + (tile_offset.strided() * InstructionShape::kStrided * stride_); + + add_pointer_offset(offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + pointer_ += stride_ * InstructionShape::kStrided; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_byte_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + AccessType *fetch_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::Iterations::kStrided; ++s) { + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::Iterations::kContiguous; ++c) { + + int access_idx = c + s * Policy::Iterations::kContiguous; + + AccessType const *source_ptr = pointer_ + + Policy::Delta::kContiguous * c + + Policy::Delta::kStrided * s * stride_; + + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_; + + AccessType const *source = reinterpret_cast(source_byte_ptr); + + fetch_ptr[access_idx] = *source; + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = + tile_offset.contiguous() * Shape::kContiguous + + tile_offset.strided() * InstructionShape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + + } +}; + +//////////////////////////////////////////////////////////////////////////////// +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::RowMajorTensorOpMultiplicandCongruous128b, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajorTensorOpMultiplicandCongruous128b; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCongruous128b, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): iterator_({ref.data(), ref.stride()}, lane_id) { + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(layout::PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(layout::PitchLinearCoord(-tile_offset.column(), -tile_offset.row())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.strided(), tile_offset.contiguous()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::ColumnMajorTensorOpMultiplicandCongruous128b, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::ColumnMajorTensorOpMultiplicandCongruous128b; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCongruous128b, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): iterator_({ref.data(), ref.stride()}, lane_id) { + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(layout::PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(layout::PitchLinearCoord(-tile_offset.row(), -tile_offset.column())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.contiguous(), tile_offset.strided()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// +/// Partial specialization for complex +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of underlying field of reals. + typename RealElement, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_> +class MmaTensorOpAccumulatorTileIterator< + Shape_, complex, mctlass::layout::RowMajor, InstructionShape_, OpDelta_> { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = complex; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + using OpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kRow % InstructionShape::kM) && + !(Shape::kColumn % InstructionShape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + + static_assert(platform::is_same::value, + "Layouts must be defined for logical MatrixCoord coordinate space."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape; + }; + +private: + + // Assume accumulator tile is an arrangement of 8-by-8 tiles replicated over the entire + // shape, with each quad mapped to one row and each thread mapped to 1/4 of the elements + // of that row. The accumulators within one row are assumed to be consecutive. + static int const kElementsPerAccess = InstructionShape::kN / 4; + static int const kRowsPerTile = 8; + static int const kAccumulatorRows = InstructionShape::kM / kRowsPerTile; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile. It is assumed that the accumulators + /// are stored in a planar complex arrangement with the real parts as entirely contiguous + /// followed by the imaginary parts. + using Fragment = Array; + + static int const kRealIndex = 0; + static int const kImaginaryIndex = Shape::kCount / kThreads; + +private: + + /// Reference to output tensor + TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator( + TensorRef const &ref, + int lane_id + ): + ref_(ref) { + + if (platform::is_same>::value == true) { + //Special process for mma shape: m8n8k4f64 + if (InstructionShape::kM == 8 && InstructionShape::kN == 8 + && InstructionShape::kK == 4) { + + const int lane_in_quad = lane_id & 0x7; + const int quad = (lane_id >> 4) & 0x3; + MatrixCoord lane_offset(quad, lane_in_quad); + ref_.add_coord_offset(lane_offset); + + } + else { + //To Do + printf("Not Impl for other InstructionShape...\n"); + } + } + else if (platform::is_same>::value == true) { + const int quad = (lane_id >> 4) << 2; + const int lane_in_quad = lane_id & 0x7; + MatrixCoord lane_offset(quad, lane_in_quad); + ref_.add_coord_offset(lane_offset); + } + else { + + const int quad = (lane_id >> 2); + const int lane_in_quad = (lane_id & 3); + MatrixCoord lane_offset(quad, lane_in_quad * kElementsPerAccess); + ref_.add_coord_offset(lane_offset); + + } + + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + ref_.add_coord_offset(tile_offset * make_Coord(Shape::kRow, Shape::kColumn)); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator++() { + // deliberate no-op + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator--() { + // deliberate no-op + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to load from the tensor + Index pointer_offset) const { ///< loads a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + + Element z = offset_ref.at({accum_m, accum_n}); + + frag[mma_accum_start + row * kElementsPerAccess + col + kRealIndex] = z.real(); + frag[mma_accum_start + row * kElementsPerAccess + col + kImaginaryIndex] = z.imag(); + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + Fragment &frag, ///< fragment to load from the tensor + Index byte_offset) const { ///< loads a tile with a linear offset + + load_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset) const { ///< loads a tile with a logical offset in units of whole tiles + + load(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset, ///< loads a tile with a logical offset in units of whole tiles + Index pointer_offset) const { ///< loads a tile with a logical offset AND a pointer offset + + load_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } + + /// Stores a fragment to memory + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index pointer_offset) const { ///< store a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_m = 0; + int accum_n = 0; + int idx = 0; + if (platform::is_same>::value == true) { + //Special process for mma shape: m8n8k4f64 + if (InstructionShape::kM == 8 &&InstructionShape::kN == 8 + && InstructionShape::kK == 4) { + accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + row * kRowsPerTile + + col * (InstructionShape::kM / kElementsPerAccess); + accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn; + idx =mma_accum_start + row * kElementsPerAccess + col; + } + else { + //To Do + printf("Not Impl now.\n"); + } + } + else if (platform::is_same>::value == true) { + accum_m = mma_m * InstructionShape::kM; + accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn ; + int t_idx = row * kElementsPerAccess + col; + idx = mma_accum_start + t_idx; + accum_m += t_idx; + } + else { + accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile; + accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + idx =mma_accum_start + row * kElementsPerAccess + col; + } + + Element z(frag[kRealIndex + idx], frag[kImaginaryIndex + idx]); + + offset_ref.at({accum_m, accum_n}) = z; + } + } + } + } + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_byte_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index byte_offset) const { ///< store a tile with a linear offset + + store_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Stores a fragment to memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + Fragment &frag, ///< fragment to store to the tensor + TensorCoord const &tile_offset) const { ///< stores a tile with a logical offset in units of whole tiles + + store(frag, tile_offset, 0); + } + + /// Stores a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + /// fragment to store to the tensor + Fragment const &frag, + /// stores a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// stores a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + store_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////////////////////////////////////////////////// + + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for loading 128b vectors of 128b elements. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::TensorOpMultiplicandCrosswise128x4, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + static_assert(!(Shape::kContiguous % 4) && !(Shape::kStrided % 8), "Divisibility."); + + static_assert(sizeof_bits::value == 128, "This is specialized for 128b accesses."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCrosswise128x4; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Load two elements per access + static int const kElementsPerAccess = 1; + + /// Policy defining internal details of tile iterator + struct Policy { + + /// Shape of one access + using Delta = layout::PitchLinearShape<4, 8>; + + /// Number of iterations to load + using Iterations = layout::PitchLinearShape< + InstructionShape::kContiguous / Delta::kContiguous, + Shape::kStrided / Delta::kStrided + >; + }; + +private: + + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = AlignedArray; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + +private: + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(): stride_(0), byte_offset_(0) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): + stride_(ref.stride(0) / kElementsPerAccess), byte_offset_(0) { + + int quad = lane_id / 4; + int liq = lane_id % 4; + + int c = liq + (quad & 1) * 4; + int s = (quad / 2); + + byte_offset_ = (c + s * stride_) * sizeof(AccessType); + + pointer_= reinterpret_cast(ref.data()); + + //Special process for mma shape: m8n8k4f64 + if (platform::is_same>::value == true) { + + const int quad = lane_id & 0x7; + const int liq = (lane_id >> 4) & 0x3; + const int c = liq + (quad & 1) * 4; + const int s = (quad / 4) * 2 + (quad & 2) / 2; + byte_offset_ = (c + s * stride_) * sizeof(AccessType); + + } + + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + pointer_ += offset; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + // Compute the offset in units of elements. Note, the external coordinate system is + // approximately transposed with respect to the tiled internal structure + int offset = + (tile_offset.contiguous() * InstructionShape::kContiguous) * stride_ + + (tile_offset.strided() * Shape::kStrided); + + add_pointer_offset(offset); + + byte_offset_ ^= (tile_offset.contiguous() & 1) * 4 * sizeof(AccessType); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + pointer_ += stride_ * InstructionShape::kContiguous; + + byte_offset_ ^= 4 * sizeof(AccessType); + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_byte_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + AccessType *fetch_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::Iterations::kContiguous; ++c) { + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::Iterations::kStrided; ++s) { + + int access_idx = s + c * Policy::Iterations::kStrided; + + AccessType const *source_ptr = pointer_ + + Policy::Delta::kContiguous * c * stride_ + + Policy::Delta::kStrided * s; + + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_; + + AccessType const *source = reinterpret_cast(source_byte_ptr); + + fetch_ptr[access_idx] = *source; + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = + tile_offset.contiguous() * InstructionShape::kContiguous * stride_ + + tile_offset.strided() * Shape::kStrided; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + + } +}; + + +//////////////////////////////////////////////////////////////////////////////// +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::RowMajorTensorOpMultiplicandCrosswise128x4, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajorTensorOpMultiplicandCrosswise128x4; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCrosswise128x4, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): iterator_({ref.data(), ref.stride()}, lane_id) { + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(layout::PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(layout::PitchLinearCoord(-tile_offset.column(), -tile_offset.row())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.strided(), tile_offset.contiguous()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + + +//////////////////////////////////////////////////////////////////////////////// +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::ColumnMajorTensorOpMultiplicandCrosswise128x4, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::ColumnMajorTensorOpMultiplicandCrosswise128x4; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCrosswise128x4, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): iterator_({ref.data(), ref.stride()}, lane_id) { + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(layout::PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(layout::PitchLinearCoord(-tile_offset.row(), -tile_offset.column())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.contiguous(), tile_offset.strided()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Congruous shared memory layout +// Warp-level iterators for complex*complex + complex => complex +// The underlying iterators are similar to that for MMA f64*f64 + f64 = f64 +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for loading 128b vectors of 64b elements. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, mctlass::complex, + mctlass::layout::TensorOpMultiplicandCongruous64b, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + static_assert(!(Shape::kContiguous % 16) && !(Shape::kStrided % 8), "Divisibility."); + + /// Element type + using Element = mctlass::complex; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCongruous64b; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Load two elements per access + static int const kElementsPerAccess = 2; + + /// Policy defining internal details of tile iterator + struct Policy { + + /// Shape of one access + using Delta = layout::PitchLinearShape<8, 4>; + + /// Number of iterations to load + using Iterations = layout::PitchLinearShape< + Shape::kContiguous / kElementsPerAccess / Delta::kContiguous, + InstructionShape::kStrided / Delta::kStrided + >; + + }; + +private: + + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = AlignedArray; + + /// Internal counter used to jump to next K partition + int k_group_idx_; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + +private: + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_[kElementsPerAccess]; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(): stride_(0), byte_offset_(0) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): + stride_(ref.stride(0) / kElementsPerAccess), byte_offset_(0), + k_group_idx_(0) { + + /* int access_strided = lane_id / Policy::Delta::kContiguous; + int access_contiguous = (lane_id % Policy::Delta::kContiguous) ^ access_strided; + pointer_ = reinterpret_cast(ref.data()) + + access_contiguous + access_strided * stride_; + */ + int access_strided = lane_id / Policy::Delta::kContiguous; + for (int i = 0; i < kElementsPerAccess; ++i) { + int lane_id_offset = ((__lane_id() >> 4) & 0x1) - 2 * (__lane_id() >> 5) + i; + int access_contiguous = ((lane_id + lane_id_offset) % Policy::Delta::kContiguous) ^ access_strided; + pointer_[i] = reinterpret_cast(ref.data()) + + access_contiguous + access_strided * stride_; + } + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + byte_offset_ += offset * sizeof(Element); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + int offset = + (tile_offset.strided() * InstructionShape::kStrided) * stride_ * kElementsPerAccess + + tile_offset.contiguous() * Shape::kContiguous; + + add_pointer_offset(offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + add_tile_offset({0, 1}); + + return *this; + } + + /// Advances the iterator along the opposite of the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + add_tile_offset({0, -1}); + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_byte_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + AccessType *fetch_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::Iterations::kStrided; ++s) { + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::Iterations::kContiguous; ++c) { + + int access_idx = c + s * Policy::Iterations::kContiguous; + + /* AccessType const *source_ptr = pointer_ + + Policy::Delta::kContiguous * c + + Policy::Delta::kStrided * s * stride_; + + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_; + + AccessType const *source = reinterpret_cast(source_byte_ptr); + + fetch_ptr[access_idx] = *source; + */ + if (kOperand == Operand::kA) { + AccessType const *source_ptr[kElementsPerAccess]; + char const *source_byte_ptr[kElementsPerAccess]; + Element const *source_element[kElementsPerAccess]; + for (int i = 0; i < kElementsPerAccess; ++i) { + source_ptr[i] = pointer_[i] + Policy::Delta::kStrided * (__lane_id() >> 5) * stride_; + source_byte_ptr[i] = reinterpret_cast(source_ptr[i]) + byte_offset + byte_offset_ + 16 * access_idx * Policy::Delta::kContiguous; + source_element[i] = reinterpret_cast(source_byte_ptr[i] + ((__lane_id() >> 3) & 0x1) * Policy::Delta::kContiguous); + } + Element source_element_ptr[kElementsPerAccess] ={*source_element[0], *source_element[1]}; + char const *source_byte_temp_ptr = reinterpret_cast(source_element_ptr); + AccessType const *source = reinterpret_cast(source_byte_temp_ptr); + fetch_ptr[access_idx] = *source; + } + else { + int index = access_idx & 0x1; + int offset = access_idx / kElementsPerAccess; + AccessType const *source_ptr = pointer_[index] + Policy::Delta::kStrided * (__lane_id() >> 5) * stride_; + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_ + + 16 * offset * Policy::Delta::kContiguous; + AccessType const *source = reinterpret_cast(source_byte_ptr); + if(Policy::Iterations::kContiguous > 1) { + fetch_ptr[index * 2 + offset] = *source; + } + else { + fetch_ptr[index] = *source; + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + + Index pointer_offset = + tile_offset.contiguous() * Shape::kContiguous / Layout::kElementsPerAccess + + tile_offset.strided() * InstructionShape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Crosswise shared memory layout +// Warp-level iterators for complex*complex + complex => complex +// The underlying iterators are similar to that for f64*f64 + f64 = f64 +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for loading 128b vectors of 64b elements. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, complex, + mctlass::layout::TensorOpMultiplicand64bCrosswise, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + static_assert(!(Shape::kContiguous % 4) && !(Shape::kStrided % 16), "Divisibility."); + + static_assert(sizeof_bits>::value == 64, "This is specialized for 64b accesses."); + + /// Element type + using Element = complex; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicand64bCrosswise; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Load two elements per access + static int const kElementsPerAccess = 2; + + /// Policy defining internal details of tile iterator + struct Policy { + + /// Shape of one access + using Delta = layout::PitchLinearShape<4, 16>; + + /// Number of iterations to load + using Iterations = layout::PitchLinearShape< + InstructionShape::kContiguous / Delta::kContiguous, + Shape::kStrided / Delta::kStrided + >; + + }; + +private: + + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = AlignedArray; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + +private: + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + /// Internal counter for tracking K-group + Index k_group_idx_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(): stride_(0), byte_offset_(0) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): + stride_(ref.stride(0) / kElementsPerAccess), byte_offset_(0), + k_group_idx_(0) { + /* + int access_strided = lane_id / 8; + int access_contiguous = (lane_id % 8); + + byte_offset_ = (access_contiguous + access_strided * stride_) * sizeof(AccessType); + + pointer_= reinterpret_cast(ref.data()); + */ + lane_id = __lane_id(); + if (kOperand == Operand::kA) { + int row = lane_id & 0xf; + int col = (lane_id >> 4) << 1; + TensorCoord coord(col, row); + Layout lo(ref.stride(0)); + byte_offset_ = lo(coord) * (sizeof_bits::value / 8); + + } + else { + int row = (lane_id >> 4) << 1; + int col = lane_id & 0x7; + TensorCoord coord(row, col); + Layout lo(ref.stride(0)); + byte_offset_ = lo(coord) * (sizeof_bits::value / 8); + } + pointer_= reinterpret_cast(ref.data()); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + pointer_ += offset / kElementsPerAccess; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + int offset = (tile_offset.contiguous() * InstructionShape::kContiguous) * + stride_ * kElementsPerAccess + + tile_offset.strided() * Shape::kStrided; + + add_pointer_offset(offset); + + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative(TensorCoord const &tile_offset) { + + add_tile_offset(tile_offset); + + if (k_group_idx_ & 1) + byte_offset_ ^= 0x40; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + pointer_ += stride_ * InstructionShape::kContiguous; + + // xor ptr + byte_offset_ ^= 0x40; + + ++k_group_idx_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_byte_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + AccessType *fetch_ptr = reinterpret_cast(&frag); + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::Iterations::kContiguous; ++c) { + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::Iterations::kStrided; ++s) { + + int access_idx = c * Policy::Iterations::kStrided + s; + + /* AccessType const *source_ptr = pointer_ + + Policy::Delta::kContiguous * c * stride_ + + Policy::Delta::kStrided * s / kElementsPerAccess; + + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_; + + AccessType const *source = reinterpret_cast(source_byte_ptr); + + fetch_ptr[access_idx] = *source; + */ + /// For a 4x4 A-matrix: 1 1 2 3 + // 4 5 6 7 + // 8 9 10 11 + // 12 13 14 15 + // the reading order is (__lane_id==0)1 1 ; (__lane_id==1)4 5 ; (__lane_id==2)8 9 ; (__lane_id==3)12 13 + // (__lane_id==16)2 3 ; (__lane_id==17)6 7 ; (__lane_id==18)10 11 ; (__lane_id==19)14 15 + + /// For a 4x4 B-matrix: 1 1 2 3 + // 4 5 6 7 + // 8 9 10 11 + // 12 13 14 15 + // the reading order is (__lane_id==0)1 4 ; (__lane_id==1)1 5 ; (__lane_id==2)2 6 ; (__lane_id==3)3 7 + // (__lane_id==16)8 12 ; (__lane_id==17)9 13 ; (__lane_id==18)10 14 ; (__lane_id==19)11 15 + + if (kOperand == Operand::kA) { + AccessType const *source_ptr = pointer_ + + Policy::Delta::kContiguous * c * stride_ + + Policy::Delta::kStrided * s / kElementsPerAccess; + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_ ; + AccessType const *source = reinterpret_cast(source_byte_ptr); + char const *source_byte_ptr1 = reinterpret_cast(source_ptr) + byte_offset + byte_offset_ + sizeof(AccessType); + AccessType const *source1 = reinterpret_cast(source_byte_ptr1); + Element source_element[kElementsPerAccess]; + + source_element[0] = *reinterpret_cast(source); + source_element[1] = *reinterpret_cast(source1); + + fetch_ptr[access_idx] = *reinterpret_cast(source_element); + const mctlass::complex* ptr= fetch_ptr[access_idx].data(); + } + else { + AccessType const *source_ptr = pointer_ + + Policy::Delta::kStrided * s / kElementsPerAccess; + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_ + c * sizeof(AccessType); + AccessType const *source = reinterpret_cast(source_byte_ptr); + char const *source_byte_ptr1 = source_byte_ptr + sizeof(AccessType); + AccessType const *source1 = reinterpret_cast(source_byte_ptr1); + Element source_element[kElementsPerAccess]; + source_element[0] = *reinterpret_cast(source); + if (__lane_id() >= 32) { + source_element[1] = *(reinterpret_cast(source) - 1); + } + else { + source_element[1] = *(reinterpret_cast(source1) - 1); + } + fetch_ptr[access_idx] = *reinterpret_cast(source_element); + } + } + } + + /* Element *exchange_ptr = reinterpret_cast(&frag); + + // exchange on 64b granularity only for fragments held in k=8/2 to k=8 + MCTLASS_PRAGMA_UNROLL + for (int i = Fragment::kElements/2; i < Fragment::kElements; i += 2) { + Element tmp = exchange_ptr[i]; + exchange_ptr[i] = exchange_ptr[i + 1]; + exchange_ptr[i + 1] = tmp; + } + */ + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = tile_offset.contiguous() * + InstructionShape::kContiguous / + Layout::kElementsPerAccess + + tile_offset.strided() * Shape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + k_group_idx_ = k_group; + } +}; + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_gaussian_complex_tensor_op.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_gaussian_complex_tensor_op.h new file mode 100644 index 0000000..ba26bef --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_gaussian_complex_tensor_op.h @@ -0,0 +1,642 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations targeting + Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" +#include "mctlass/complex.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/arch/mma_sm75.h" +#include "mctlass/arch/mma_sm80.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/warp/mma_tensor_op_policy.h" +#include "mctlass/gemm/warp/mma_tensor_op.h" + +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" +#include "mctlass/gemm/warp/mma_gaussian_complex_tensor_op_tile_iterator_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename RealElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename RealElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename RealElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Complex transform on A operand + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex transform on B operand + ComplexTransform TransformB = ComplexTransform::kNone, + /// Do source operands need more than one elements + bool GeneralizedOperatorElements = false, + /// Used for partial specialization + typename Enable = bool +> +class MmaGaussianComplexTensorOp; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex*complex+complex => complex using real-valued TensorOps +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename RealElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename RealElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename RealElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB +> +class MmaGaussianComplexTensorOp< + Shape_, + complex, + LayoutA_, + complex, + LayoutB_, + complex, + LayoutC_, + Policy_, + TransformA, + TransformB> { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of multiplicand A + using ElementA = complex; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of multiplicand B + using ElementB = complex; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of accumulator matrix C + using ElementC = complex; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Policy::Operator; + + /// Shape of underlying instruction + using InstructionShape = typename ArchMmaOperator::Shape; + + /// Underlying arch tag + using ArchTag = typename ArchMmaOperator::ArchTag; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassTensorOp; + + /// Indicates math operator + using MathOperator = arch::OpMultiplyAddGaussianComplex; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = TransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = TransformB; + + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 32; + +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kA, + ElementA, + LayoutA, + MatrixShape, + Policy::OpDelta::kRow, + 32, + 1 + >; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentA = FragmentA; + + /// Iterates over the B operand in memory + using IteratorB = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kB, + ElementB, + LayoutB, + MatrixShape, + Policy::OpDelta::kColumn, + 32, + 1 + >; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed B tile + using TransformedFragmentB = FragmentB; + + static_assert( + !(Shape::kM % ArchMmaOperator::Shape::kM) && + !(Shape::kN % ArchMmaOperator::Shape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + Shape::kM / ArchMmaOperator::Shape::kM, + Shape::kN / ArchMmaOperator::Shape::kN + >; + + /// Iterates over the C operand in memory + using IteratorC = MmaTensorOpGaussianComplexAccumulatorTileIterator< + MatrixShape, + ElementC, + LayoutC, + typename ArchMmaOperator::Shape, + typename Policy::OpDelta>; + + /// Storage for C tile, the accumulator. Note, regardless of multiplicand type, this + /// storage arrangement is to be considered 'gaussian complex' in the sense that the accumulation is + /// done in three parts namely part1, part2, and part3. The parts 1, 2, and 3 are stored consecutively + /// in InteratorC::Frament. This matches the structure of Tensor Cores which are always real-valued matrix multiplies. + using FragmentC = typename IteratorC::Fragment; + + static_assert( + FragmentC::kElements == 3 * MmaIterations::kCount * ArchMmaOperator::FragmentC::kElements, + "Unexpected gaussian complex fragment length."); + +private: + + // + // Data members + // + + /// Underlying real-valued matrix multiply operator (concept: arch::Mma) + ArchMmaOperator mma; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaGaussianComplexTensorOp() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + FragmentA const &A, + FragmentB const &B, + FragmentC const &C + ) const { + + // Alias types for underlying real-valued matrix multiply operator + using MmaOperandA = typename ArchMmaOperator::FragmentA; + using MmaOperandB = typename ArchMmaOperator::FragmentB; + using MmaOperandC = typename ArchMmaOperator::FragmentC; + + static_assert(MmaOperandA::kElements == 1, + "This implementation only supports math instructions in which exactly one element is needed for the A operand." + "We can geneneralize later."); + + static_assert(MmaOperandB::kElements == 1, + "This implementation only supports math instructions in which exactly one element is needed for the B operand." + "We can geneneralize later."); + + D = C; + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < MmaIterations::kRow; ++m) { + + // mma(accum.part1(), (a.real() + a.imag()), b.real(), accum.part1()); + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_Asum; + MmaOperandB operand_Br; + + operand_Asum[0] = A[m].real() + ((kTransformA == ComplexTransform::kConjugate) ? -A[m].imag() : +A[m].imag()); + operand_Br[0] = B[n].real(); + + // accumulator part1 + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow); + + mma(*accum, operand_Asum, operand_Br, *accum); + } + + // mma(accum.part2(), -a.real(), (b.real() - b.imag()), accum.part2()); + MCTLASS_PRAGMA_UNROLL + for (int n = MmaIterations::kColumn - 1; n >= 0; --n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_Ar; + MmaOperandB operand_Bdiff; + + operand_Ar[0] = -A[m].real(); + operand_Bdiff[0] = B[n].real() - ((kTransformB == ComplexTransform::kConjugate) ? -B[n].imag() : +B[n].imag()); + + // accumulator part2 + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + MmaIterations::kCount; + + mma(*accum, operand_Ar, operand_Bdiff, *accum); + } + + // mma(accum.part3(), a.imag(), (b.real() + b.imag()), accum.part3()) + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_Ai; + MmaOperandB operand_Bsum; + + operand_Ai[0] = (kTransformA == ComplexTransform::kConjugate) ? -A[m].imag() : +A[m].imag(); + operand_Bsum[0] = B[n].real() + ((kTransformB == ComplexTransform::kConjugate) ? -B[n].imag() : +B[n].imag()); + + // accumulator part3 + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + 2 * MmaIterations::kCount; + + mma(*accum, operand_Ai, operand_Bsum, *accum); + } + } + } + + /// Transform the mma operands to the required types + MCTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + dst_A = A; + dst_B = B; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for complex*complex+complex => complex using real-valued TensorOps +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename RealElementA, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename RealElementB, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename RealElementC, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Complex transform on A operand + ComplexTransform TransformA, + /// Complex transform on B operand + ComplexTransform TransformB +> +class MmaGaussianComplexTensorOp< + Shape_, + complex, + LayoutA_, + complex, + LayoutB_, + complex, + LayoutC_, + Policy_, + TransformA, + TransformB, + true> { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of multiplicand A + using ElementA = complex; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of multiplicand B + using ElementB = complex; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of accumulator matrix C + using ElementC = complex; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Policy::Operator; + + /// Shape of underlying instruction + using InstructionShape = typename ArchMmaOperator::Shape; + + /// Underlying arch tag + using ArchTag = typename ArchMmaOperator::ArchTag; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassTensorOp; + + /// Indicates math operator + using MathOperator = arch::OpMultiplyAddGaussianComplex; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = TransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = TransformB; + + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 32; + +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kA, + ElementA, + LayoutA, + MatrixShape, + Policy::OpDelta::kRow, + 32, + 1 + >; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentA = FragmentA; + + /// Iterates over the B operand in memory + using IteratorB = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kB, + ElementB, + LayoutB, + MatrixShape, + Policy::OpDelta::kColumn, + 32, + 1 + >; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed B tile + using TransformedFragmentB = FragmentB; + + static_assert( + !(Shape::kM % ArchMmaOperator::Shape::kM) && + !(Shape::kN % ArchMmaOperator::Shape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + Shape::kM / ArchMmaOperator::Shape::kM, + Shape::kN / ArchMmaOperator::Shape::kN + >; + + /// Iterates over the C operand in memory + using IteratorC = MmaTensorOpGaussianComplexAccumulatorTileIterator< + MatrixShape, + ElementC, + LayoutC, + typename ArchMmaOperator::Shape, + typename Policy::OpDelta>; + + /// Storage for C tile, the accumulator. Note, regardless of multiplicand type, this + /// storage arrangement is to be considered 'gaussian complex' in the sense that the accumulation is + /// done in three parts namely part1, part2, and part3. The parts 1, 2, and 3 are stored consecutively + /// in InteratorC::Frament. This matches the structure of Tensor Cores which are always real-valued matrix multiplies. + using FragmentC = typename IteratorC::Fragment; + + static_assert( + FragmentC::kElements == 3 * MmaIterations::kCount * ArchMmaOperator::FragmentC::kElements, + "Unexpected gaussian complex fragment length."); + +private: + + // + // Data members + // + + /// Underlying real-valued matrix multiply operator (concept: arch::Mma) + ArchMmaOperator mma; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaGaussianComplexTensorOp() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + FragmentA const &A, + FragmentB const &B, + FragmentC const &C + ) const { + + // Alias types for underlying real-valued matrix multiply operator + using MmaOperandA = typename ArchMmaOperator::FragmentA; + using MmaOperandB = typename ArchMmaOperator::FragmentB; + using MmaOperandC = typename ArchMmaOperator::FragmentC; + + D = C; + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < MmaIterations::kRow; ++m) { + + // mma(accum.part1(), (a.real() + a.imag()), b.real(), accum.part1()); + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_Asum; + MmaOperandB operand_Br; + + MCTLASS_PRAGMA_UNROLL + for (int mk = 0; mk < MmaOperandA::kElements; ++mk) + operand_Asum[mk] = A[m*MmaOperandA::kElements + mk].real() + ((kTransformA == ComplexTransform::kConjugate) ? + -A[m*MmaOperandA::kElements + mk].imag() : +A[m*MmaOperandA::kElements + mk].imag()); + + MCTLASS_PRAGMA_UNROLL + for (int nk = 0; nk < MmaOperandB::kElements; ++nk) + operand_Br[nk] = B[n*MmaOperandB::kElements + nk].real(); + + // accumulator part1 + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow); + + mma(*accum, operand_Asum, operand_Br, *accum); + } + + // mma(accum.part2(), -a.real(), (b.real() - b.imag()), accum.part2()); + MCTLASS_PRAGMA_UNROLL + for (int n = MmaIterations::kColumn - 1; n >= 0; --n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_Ar; + MmaOperandB operand_Bdiff; + + MCTLASS_PRAGMA_UNROLL + for (int mk = 0; mk < MmaOperandA::kElements; ++mk) + operand_Ar[mk] = -A[m*MmaOperandA::kElements + mk].real(); + + MCTLASS_PRAGMA_UNROLL + for (int nk = 0; nk < MmaOperandB::kElements; ++nk) + operand_Bdiff[nk] = B[n*MmaOperandB::kElements + nk].real() - ((kTransformB == ComplexTransform::kConjugate) ? + -B[n*MmaOperandB::kElements + nk].imag() : +B[n*MmaOperandB::kElements + nk].imag()); + + // accumulator part2 + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + MmaIterations::kCount; + + mma(*accum, operand_Ar, operand_Bdiff, *accum); + } + + // mma(accum.part3(), a.imag(), (b.real() + b.imag()), accum.part3()) + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // Pack operands together. This may result in actual MOVs + MmaOperandA operand_Ai; + MmaOperandB operand_Bsum; + + MCTLASS_PRAGMA_UNROLL + for (int mk = 0; mk < MmaOperandA::kElements; ++mk) + operand_Ai[mk] = (kTransformA == ComplexTransform::kConjugate) ? + -A[m*MmaOperandA::kElements + mk].imag() : +A[m*MmaOperandA::kElements + mk].imag(); + + MCTLASS_PRAGMA_UNROLL + for (int nk = 0; nk < MmaOperandB::kElements; ++nk) + operand_Bsum[nk] = B[n*MmaOperandB::kElements + nk].real() + ((kTransformB == ComplexTransform::kConjugate) ? + -B[n*MmaOperandB::kElements + nk].imag() : +B[n*MmaOperandB::kElements + nk].imag()); + + // accumulator part3 + MmaOperandC *accum = reinterpret_cast(&D) + + (m + n * MmaIterations::kRow) + 2 * MmaIterations::kCount; + + mma(*accum, operand_Ai, operand_Bsum, *accum); + } + } + } + + /// Transform the mma operands to the required types + MCTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + dst_A = A; + dst_B = B; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_gaussian_complex_tensor_op_tile_iterator_sm80.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_gaussian_complex_tensor_op_tile_iterator_sm80.h new file mode 100644 index 0000000..18ee839 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_gaussian_complex_tensor_op_tile_iterator_sm80.h @@ -0,0 +1,395 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines iterators used by warp-level matrix multiply operations targeting Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/gemm/gemm.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor_op_multiplicand_sm80.h" +#include "mctlass/gemm/warp/mma_complex_tensor_op_tile_iterator_sm80.h" + +#include "mctlass/platform/platform.h" +#include "mctlass/fast_math.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Layout of operand in memory + typename Layout_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_> +class MmaTensorOpGaussianComplexAccumulatorTileIterator; + +//////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// +/// Partial specialization for complex +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of underlying field of reals. + typename RealElement, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_> +class MmaTensorOpGaussianComplexAccumulatorTileIterator< + Shape_, complex, mctlass::layout::RowMajor, InstructionShape_, OpDelta_> { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = complex; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + using OpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kRow % InstructionShape::kM) && + !(Shape::kColumn % InstructionShape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + + static_assert(platform::is_same::value, + "Layouts must be defined for logical MatrixCoord coordinate space."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape; + }; + +private: + + // Assume accumulator tile is an arrangement of 8-by-8 tiles replicated over the entire + // shape, with each quad mapped to one row and each thread mapped to 1/4 of the elements + // of that row. The accumulators within one row are assumed to be consecutive. + static int const kElementsPerAccess = InstructionShape::kN / 4; + static int const kRowsPerTile = 8; + static int const kAccumulatorRows = InstructionShape::kM / kRowsPerTile; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile. It is assumed that the accumulators + /// are stored in a gaussian complex arrangement with parts 1, 2, and 3 as entirely contiguous + /// arranged as [part1, part2, part3] + using Fragment = Array; + + static int const kPart1Index = (Shape::kCount / kThreads) * 0; + static int const kPart2Index = (Shape::kCount / kThreads) * 1; + static int const kPart3Index = (Shape::kCount / kThreads) * 2; + +private: + + /// Reference to output tensor + TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpGaussianComplexAccumulatorTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpGaussianComplexAccumulatorTileIterator( + TensorRef const &ref, + int lane_id + ): + ref_(ref) { + + // int quad = (lane_id >> 2); + // int lane_in_quad = (lane_id & 3); + // MatrixCoord lane_offset(quad, lane_in_quad * kElementsPerAccess); + int lane_in_quad = lane_id & 0x7; + int quad = (lane_id >> 4) & 0x3; + MatrixCoord lane_offset(quad, lane_in_quad); + + ref_.add_coord_offset(lane_offset); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpGaussianComplexAccumulatorTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpGaussianComplexAccumulatorTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + ref_.add_coord_offset(tile_offset * make_Coord(Shape::kRow, Shape::kColumn)); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpGaussianComplexAccumulatorTileIterator & operator++() { + // deliberate no-op + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpGaussianComplexAccumulatorTileIterator & operator--() { + // deliberate no-op + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpGaussianComplexAccumulatorTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpGaussianComplexAccumulatorTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to load from the tensor + Index pointer_offset) const { ///< loads a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + + Element z = offset_ref.at({accum_m, accum_n}); + + frag[mma_accum_start + row * kElementsPerAccess + col + kPart1Index] = z.real() + z.imag(); + frag[mma_accum_start + row * kElementsPerAccess + col + kPart2Index] = -z.real(); + frag[mma_accum_start + row * kElementsPerAccess + col + kPart3Index] = z.imag(); + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + Fragment &frag, ///< fragment to load from the tensor + Index byte_offset) const { ///< loads a tile with a linear offset + + load_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset) const { ///< loads a tile with a logical offset in units of whole tiles + + load(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset, ///< loads a tile with a logical offset in units of whole tiles + Index pointer_offset) const { ///< loads a tile with a logical offset AND a pointer offset + + load_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } + + /// Stores a fragment to memory + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index pointer_offset) const { ///< store a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + // int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + // row * kRowsPerTile; + // int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + row * kRowsPerTile + + col * (InstructionShape::kM / kElementsPerAccess); + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn; + int idx = mma_accum_start + row * kElementsPerAccess + col; + + Element z(frag[kPart1Index + idx] - frag[kPart3Index + idx], + frag[kPart1Index + idx] + frag[kPart2Index + idx]); + + offset_ref.at({accum_m, accum_n}) = z; + } + } + } + } + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_byte_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index byte_offset) const { ///< store a tile with a linear offset + + store_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Stores a fragment to memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + Fragment &frag, ///< fragment to store to the tensor + TensorCoord const &tile_offset) const { ///< stores a tile with a logical offset in units of whole tiles + + store(frag, tile_offset, 0); + } + + /// Stores a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + /// fragment to store to the tensor + Fragment const &frag, + /// stores a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// stores a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + store_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_planar_complex.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_planar_complex.h new file mode 100644 index 0000000..5c2937f --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_planar_complex.h @@ -0,0 +1,182 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/complex.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/gemm/gemm.h" + +#include "mctlass/array_planar_complex.h" +#include "mctlass/gemm/warp/tile_iterator_planar_complex.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Underlying real-valued warp-level matrix multiply + typename Operator_, + /// Transformation applied to A operand (typically folded into math instruction) + ComplexTransform TransformA = ComplexTransform::kNone, + /// Transformation applied to B operand (typically folded into math instruction) + ComplexTransform TransformB = ComplexTransform::kNone +> +class MmaPlanarComplex { +public: + + /// Underlying real-valued warp-level matrix multiply + using Operator = Operator_; + + /// Shape of warp-level matrix multipy + using Shape = typename Operator::Shape; + + /// Transformation applied to A operand (typically folded into math instruction) + static ComplexTransform const kTransformA = TransformA; + + /// Transformation applied to B operand (typically folded into math instruction) + static ComplexTransform const kTransformB = TransformB; + + /// Fragment of elements + using FragmentA = ArrayPlanarComplex; + + /// Iterator into planar complex + using IteratorA = TileIteratorPlanarComplex; + + /// Layout in memory of the A operand + using LayoutA = typename Operator::LayoutA; + + using FragmentB = ArrayPlanarComplex; + + /// Iterator into planar complex + using IteratorB = TileIteratorPlanarComplex; + + /// Layout in memory of the B operand + using LayoutB = typename Operator::LayoutB; + + /// Tile iterator for accumulator + using IteratorC = TileIteratorPlanarComplex; + + /// Accumulator fragment + using FragmentC = ArrayPlanarComplex; + + /// Layout of accumulator fragment in memory + using LayoutC = typename Operator::LayoutC; + +private: + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + Operator::Shape::kM / Operator::Policy::Operator::Shape::kM, + Operator::Shape::kN / Operator::Policy::Operator::Shape::kN + >; + +public: + /// Ctor + MCTLASS_DEVICE + MmaPlanarComplex() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + FragmentA const &A_in, + FragmentB const &B_in, + FragmentC const &C) const { + + D.real = C.real; + D.imag = C.imag; + + // + // Transform fragments based on conjugate operations. + // + + negate neg_A; + + FragmentA frag_A; + frag_A.real = A_in.real; + + if (kTransformA == ComplexTransform::kConjugate) { + frag_A.imag = neg_A(frag_A.imag); + } + else { + frag_A.imag = frag_A.imag; + } + + FragmentB frag_B; + frag_B.real = B_in.real; + + if (kTransformB == ComplexTransform::kConjugate) { + negate neg; + frag_B.imag = neg(frag_B.imag); + } + else { + frag_B.imag = frag_B.imag; + } + + // + // Accumulated real-valued matrix multiplies + // + + Operator real_mma; + + // D.i += A.i * B.r + real_mma(D.imag, frag_A.imag, frag_B.real, D.imag); + + // D.r += A.r * B.r + real_mma(D.real, frag_A.real, frag_B.real, D.real); + + // D.i += A.r * B.i + real_mma(D.imag, frag_A.real, frag_B.imag, D.imag); + + // D.r += -A.i * B.i + frag_A.imag = neg_A(frag_A.imag); + real_mma(D.real, frag_A.imag, frag_B.imag, D.real); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_simt.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_simt.h new file mode 100644 index 0000000..31c16f7 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_simt.h @@ -0,0 +1,263 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/thread/mma.h" + +#include "mctlass/gemm/warp/mma_simt_tile_iterator.h" +#include "mctlass/gemm/warp/mma_simt_policy.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK = 1, + /// Complex transformation on operand A + ComplexTransform TransformA = ComplexTransform::kNone, + /// Complex transformation on operand B + ComplexTransform TransformB = ComplexTransform::kNone, + /// Used for partial specialization + typename Enable = bool +> +class MmaSimt { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of multiplicand A + using ElementA = ElementA_; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of multiplicand B + using ElementB = ElementB_; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of accumulator matrix C + using ElementC = ElementC_; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassSimt; + + /// Hard-coded for now + using ArchTag = arch::Sm50; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = TransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = TransformB; + + /// Layout of threads + using ThreadLayoutA = typename platform::conditional< platform::is_same< layout::ColumnMajorInterleaved<4>, LayoutA >::value, + layout::ColumnMajor, + typename platform::conditional < platform::is_same< layout::RowMajorInterleaved<4>, LayoutA >::value, + layout::RowMajor, + LayoutA>::type + >::type; + + using ThreadLayoutB = typename platform::conditional< platform::is_same< layout::ColumnMajorInterleaved<4>, LayoutB >::value, + layout::ColumnMajor, + typename platform::conditional < platform::is_same< layout::RowMajorInterleaved<4>, LayoutB >::value, + layout::RowMajor, + LayoutB>::type + >::type; + + static constexpr bool use_dp4a = (platform::is_same< layout::ColumnMajorInterleaved<4>, LayoutA>::value || + platform::is_same< layout::RowMajorInterleaved<4>, LayoutA >::value) && + platform::is_same< ElementA, int8_t >::value && + platform::is_same< ElementB, int8_t >::value; + + using dp4a_type = typename platform::conditional< use_dp4a , int8_t, bool >::type; + + /// Thread-level matrix multiply accumulate operator + using ThreadMma = thread::Mma< + GemmShape< + Shape::kM / Policy::WarpShape::kRow, + Shape::kN / Policy::WarpShape::kColumn, + Policy::LaneMmaShape::kK>, + ElementA, + ThreadLayoutA, + ElementB, + ThreadLayoutB, + ElementC, + LayoutC, + arch::OpMultiplyAdd, + dp4a_type + >; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename ThreadMma::ArchMmaOperator; + + /// Indicates math operator + using MathOperator = typename ArchMmaOperator::Operator; + + /// Shape of the underlying instruction + using InstructionShape = GemmShape<1,1,use_dp4a ? 4 : 1>; + +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaSimtTileIterator< + MatrixShape, + Operand::kA, + ElementA, + LayoutA, + Policy, + PartitionsK, + Shape::kK + >; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentA = FragmentA; + + /// Iterates over the B operand in memory + using IteratorB = MmaSimtTileIterator< + MatrixShape, + Operand::kB, + ElementB, + LayoutB, + Policy, + PartitionsK, + Shape::kK + >; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentB = FragmentB; + + /// Iterates over the C operand in memory + using IteratorC = MmaSimtTileIterator< + MatrixShape, + Operand::kC, + ElementC, + LayoutC, + Policy + >; + + /// Storage for C tile + using FragmentC = typename ThreadMma::FragmentC; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaSimt() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &d, + FragmentA a, + FragmentB b, + FragmentC const &c, int group_idx = 0) const { + + ThreadMma mma; + + if (kTransformA == ComplexTransform::kConjugate) { + conjugate conj_a; + a = conj_a(a); + } + + if (kTransformB == ComplexTransform::kConjugate) { + conjugate conj_b; + b = conj_b(b); + } + + mma(d, a, b, c); + } + + /// Transform the mma operands to the required types + MCTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + dst_A = A; + dst_B = B; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_simt_policy.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_simt_policy.h new file mode 100644 index 0000000..822e5cf --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_simt_policy.h @@ -0,0 +1,69 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Describes the lane policy used by warp-level matrix multiply operators targeting SIMT + instructions +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Describes the arrangement and configuration of per-lane operations in warp-level matrix multiply +template < + typename WarpShape_, ///< shape of the warp in lanes (concept: MatrixShape) + typename LaneLayout_, ///< layout function of lanes + typename LaneMmaShape_ ///< size of each lane's thread-level matrix product (concept: GemmShape) +> +struct MmaSimtPolicy { + using WarpShape = WarpShape_; + using LaneLayout = LaneLayout_; + using LaneMmaShape = LaneMmaShape_; + using MmaShape = LaneMmaShape; + + /// Returns a layout functor mapping lane position in the warp to thread ID + MCTLASS_HOST_DEVICE + static LaneLayout get_lane_layout() { + return LaneLayout::packed({WarpShape::kRow, WarpShape::kColumn}); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_simt_tile_iterator.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_simt_tile_iterator.h new file mode 100644 index 0000000..8ca77ea --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_simt_tile_iterator.h @@ -0,0 +1,1892 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Describes the lane policy used by warp-level matrix multiply operators targeting SIMT + instructions +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" + +#include "mctlass/layout/matrix.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma_simt_policy.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Iterates over operands to warp-level matrix multiply operations targeting SIMT instructions +/// +/// concept: MutableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Operand identity + Operand Operand, + /// Data type of A elements + typename Element_, + /// Layout of operand + typename Layout_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Number of partitions along K dimension - used in sliced-K + int PartitionsK = 1, + /// Group Size along kPartition - used in sliced-K + int PartitionGroupSize = 1 +> +class MmaSimtTileIterator; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization for A operands of column-major layouts +/// +/// Concept: MutableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of A elements + typename Element_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Number of partitions along K dimension - used in sliced-K + int PartitionsK, + /// Group Size along kPartition - used in sliced-K + int PartitionGroupSize +> +class MmaSimtTileIterator { +public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kA; + + /// Element type + using Element = Element_; + + /// Layout of policy + using Layout = layout::ColumnMajor; + + /// Decomposition of elements among threads + using Policy = Policy_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + // + // Derived quantities + // + + static_assert(!(Shape::kRow % Policy::WarpShape::kRow), + "The warp-level GEMM M size must be divisible by the number of threads arranged along the M dimension."); + + static_assert(Shape::kRow > 0, "Shape::kRow must be greater than zero."); + static_assert(Shape::kColumn > 0, "Shape::kColumn must be greater than zero."); + static_assert(Policy::WarpShape::kRow > 0, "Policy::WarpShape::kRow must be greater than zero."); + static_assert(Shape::kRow / Policy::WarpShape::kRow > 0, "Shape::kRow / Policy::WarpShape::kRow must be greater than zero."); + + /// Thread-level shape of a fragment + using ThreadShape = MatrixShape< + Shape::kRow / Policy::WarpShape::kRow, + Shape::kColumn + >; + + static_assert(!(ThreadShape::kRow % Policy::LaneMmaShape::kM), + "Thread-level GEMM must be divisible by Policy::LaneMmaShape."); + + /// Number of individual loads + using Iterations = MatrixShape< + ThreadShape::kRow / Policy::LaneMmaShape::kM, + ThreadShape::kColumn + >; + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + +private: + + /// Internal reference + mctlass::TensorRef, layout::ColumnMajor> ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaSimtTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaSimtTileIterator( + TensorRef ref, + int lane_id + ) { + + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + MatrixCoord(Policy::LaneMmaShape::kM, 0); + + ref.add_coord_offset(lane_offset); + + ref_.reset( + reinterpret_cast *>(ref.data()), + ref.stride(0) / Policy::LaneMmaShape::kM); + } + + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_tile_offset(TensorCoord const &coord) { + + ref_.add_coord_offset({ + coord.row() * Shape::kRow / Policy::LaneMmaShape::kM, + coord.column() * Shape::kColumn}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator++() { + + ref_.add_coord_offset({0, Shape::kColumn}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator--() { + + ref_.add_coord_offset({0, -Shape::kColumn}); + + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. (vector loads) + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + Array *dst_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kColumn; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Iterations::kRow; ++m) { + + // This logic has been replaced with calls to inline PTX to guarantee vectorization. + #if 1 + dst_ptr[m + k * Iterations::kRow] = + *(ref_.data() + ref_.offset({m * Policy::WarpShape::kRow, k}) + pointer_offset / Policy::LaneMmaShape::kM); + #else + + auto ptr = ref_.data() + ref_.offset({m * Policy::WarpShape::kRow, k}) + pointer_offset / Policy::LaneMmaShape::kM; + arch::shared_load(dst_ptr[m + k * Iterations::kRow], ptr); + #endif + } + } + } + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) const { + + Array const *src_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kN; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Iterations::kM; ++m) { + *(ref_.data() + ref_.offset(m * Policy::WarpShape::kM, k) + pointer_offset / Policy::LaneMmaShape::kM) = + src_ptr[m + k * Iterations::kM]; + } + } + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation here + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization for A operands of row-major layouts +/// +/// Concept: MutableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of A elements + typename Element_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Number of partitions along K dimension - used in sliced-K + int PartitionsK, + /// Group Size along kPartition - used in sliced-K + int PartitionGroupSize +> +class MmaSimtTileIterator { +public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kA; + + /// Element type + using Element = Element_; + + /// Layout of policy + using Layout = layout::RowMajor; + + /// Decomposition of elements among threads + using Policy = Policy_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + // + // Derived quantities + // + + static_assert(!(Shape::kRow % Policy::WarpShape::kRow), + "The warp-level GEMM M size must be divisible by the number of threads arranged along the M dimension."); + + static_assert(Shape::kRow > 0, "Shape::kRow must be greater than zero."); + static_assert(Shape::kColumn > 0, "Shape::kColumn must be greater than zero."); + static_assert(Policy::WarpShape::kRow > 0, "Policy::WarpShape::kRow must be greater than zero."); + static_assert(Shape::kRow / Policy::WarpShape::kRow > 0, "Shape::kRow / Policy::WarpShape::kRow must be greater than zero."); + + /// Thread-level shape of a fragment + using ThreadShape = MatrixShape< + Shape::kRow / Policy::WarpShape::kRow, + Shape::kColumn + >; + + static_assert(!(ThreadShape::kRow % Policy::LaneMmaShape::kM), + "Thread-level GEMM must be divisible by Policy::LaneMmaShape."); + + /// Number of individual loads (scalar loads) + using Iterations = MatrixShape< + ThreadShape::kRow / Policy::LaneMmaShape::kM, + ThreadShape::kColumn + >; + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + +private: + + /// Internal reference + mctlass::TensorRef ref_; + + /// Extent of tensor + MatrixCoord extent_; + + /// Origin + MatrixCoord origin_; + + /// Used to conditionally enable extents checking + bool divisible_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaSimtTileIterator() : divisible_(true) { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaSimtTileIterator( + TensorRef ref, + int lane_id + ) : extent_(Shape::kRow, Shape::kColumn), divisible_ (true) { + + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + MatrixCoord(Policy::LaneMmaShape::kM, 0); + + origin_ = lane_offset; + + ref.add_coord_offset(lane_offset); + + ref_.reset(ref.data(), ref.stride(0)); + + } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaSimtTileIterator( + TensorRef ref, + TensorCoord extent, + int lane_id + ) : extent_(extent), divisible_ (false) { + + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + MatrixCoord(Policy::LaneMmaShape::kM, 0); + + origin_ = lane_offset; + + ref.add_coord_offset(lane_offset); + + ref_.reset(ref.data(), ref.stride(0)); + + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_tile_offset(TensorCoord const &coord) { + + TensorCoord coord_offset( + coord.row() * Shape::kRow, + coord.column() * Shape::kColumn); + + origin_ += coord_offset; + + ref_.add_coord_offset(coord_offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator++() { + + ref_.add_coord_offset({0, Shape::kColumn}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator--() { + + ref_.add_coord_offset({0, -Shape::kColumn}); + + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. (scalar loads) + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kColumn; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Iterations::kRow; ++m) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Policy::LaneMmaShape::kM; i++) { + + MatrixCoord offset(m * Policy::WarpShape::kRow * Policy::LaneMmaShape::kM + i, k); + + MatrixCoord access_coord = origin_ + offset; + + int frag_idx = m * Policy::LaneMmaShape::kM + i + k * Iterations::kRow; + + if (divisible_ || + (access_coord.row() < extent_.row() && access_coord.column() < extent_.column())) { + + frag[frag_idx] = *(ref_.data() + ref_.offset(offset) + pointer_offset); + } + else { + frag[frag_idx] = Element(); + } + } + } + } + } + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) const { + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kColumn; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Iterations::kRow; ++m) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Policy::LaneMmaShape::kM; i++) { + + *(ref_.data() + ref_.offset(m * Policy::WarpShape::kM * Policy::LaneMmaShape::kM + i, k) + pointer_offset) = + frag[m * Policy::LaneMmaShape::kM + i + k * Iterations::kM]; + } + } + } + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation here + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization for B operands of row-major layouts +/// +/// Concept: MutableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of A elements + typename Element_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK, + /// Group Size along kPartition - used in sliced-K + int PartitionGroupSize +> +class MmaSimtTileIterator { +public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kB; + + /// Element type + using Element = Element_; + + /// Layout of policy + using Layout = layout::RowMajor; + + /// Decomposition of elements among threads + using Policy = Policy_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + // + // Derived quantities + // + + static_assert(!(Shape::kColumn % Policy::WarpShape::kColumn), + "The warp-level GEMM N size must be divisible by the number of threads arranged along the N dimension."); + + static_assert(Shape::kRow > 0, "Shape::kRow must be greater than zero."); + static_assert(Shape::kColumn > 0, "Shape::kColumn must be greater than zero."); + static_assert(Policy::WarpShape::kColumn > 0, "Policy::WarpShape::kColumn must be greater than zero."); + static_assert(Shape::kColumn / Policy::WarpShape::kColumn > 0, "Shape::kColumn / Policy::WarpShape::kColumn must be greater than zero."); + + /// Thread-level shape of a fragment + using ThreadShape = MatrixShape< + Shape::kRow, + Shape::kColumn / Policy::WarpShape::kColumn + >; + + static_assert(!(ThreadShape::kColumn % Policy::LaneMmaShape::kN), + "Thread-level GEMM must be divisible by Policy::LaneMmaShape."); + + /// Number of individual loads + using Iterations = MatrixShape< + ThreadShape::kRow, + ThreadShape::kColumn / Policy::LaneMmaShape::kN + >; + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + +protected: + + /// Internal reference + mctlass::TensorRef, layout::RowMajor> ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaSimtTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaSimtTileIterator( + TensorRef ref, + int lane_id + ) { + + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + MatrixCoord(0, Policy::LaneMmaShape::kN); + + ref.add_coord_offset(lane_offset); + + ref_.reset( + reinterpret_cast *>(ref.data()), + ref.stride(0) / Policy::LaneMmaShape::kN); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_tile_offset(TensorCoord const &coord) { + + ref_.add_coord_offset({ + coord.row() * Shape::kRow, + coord.column() * Shape::kColumn / Policy::LaneMmaShape::kN}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator++() { + + ref_.add_coord_offset({Shape::kRow, 0}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator--() { + + ref_.add_coord_offset({-Shape::kRow, 0}); + + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. (vector loads) + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + Array *dst_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kRow; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Iterations::kColumn; ++n) { + + #if 1 + dst_ptr[n + k * Iterations::kColumn] = + *(ref_.data() + ref_.offset({k, n * Policy::WarpShape::kColumn}) + pointer_offset / Policy::LaneMmaShape::kN); + #else + + void const *ptr = ref_.data() + ref_.offset({k, n * Policy::WarpShape::kColumn}) + pointer_offset / Policy::LaneMmaShape::kN; + arch::shared_load(dst_ptr[n + k * Iterations::kColumn], ptr); + #endif + } + } + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) const { + + Array const *src_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kM; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Iterations::kN; ++n) { + *(ref_.data() + ref_.offset({k, n * Policy::WarpShape::kN}) + pointer_offset / Policy::LaneMmaShape::kN) = + src_ptr[n + k * Iterations::kN]; + } + } + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store(Fragment const &frag, Index pointer_offset) const { + store_with_pointer_offset(frag, 0); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation here + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization for B operands of column-major layouts +/// +/// Concept: MutableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of A elements + typename Element_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK, + /// Group Size along kPartition - used in sliced-K + int PartitionGroupSize +> +class MmaSimtTileIterator { +public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kB; + + /// Element type + using Element = Element_; + + /// Layout of policy + using Layout = layout::ColumnMajor; + + /// Decomposition of elements among threads + using Policy = Policy_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + // + // Derived quantities + // + + static_assert(!(Shape::kColumn % Policy::WarpShape::kColumn), + "The warp-level GEMM N size must be divisible by the number of threads arranged along the N dimension."); + + static_assert(Shape::kRow > 0, "Shape::kRow must be greater than zero."); + static_assert(Shape::kColumn > 0, "Shape::kColumn must be greater than zero."); + static_assert(Policy::WarpShape::kColumn > 0, "Policy::WarpShape::kColumn must be greater than zero."); + static_assert(Shape::kColumn / Policy::WarpShape::kColumn > 0, "Shape::kColumn / Policy::WarpShape::kColumn must be greater than zero."); + + /// Thread-level shape of a fragment + using ThreadShape = MatrixShape< + Shape::kRow, + Shape::kColumn / Policy::WarpShape::kColumn + >; + + static_assert(!(ThreadShape::kColumn % Policy::LaneMmaShape::kN), + "Thread-level GEMM must be divisible by Policy::LaneMmaShape."); + + /// Number of individual loads + using Iterations = MatrixShape< + ThreadShape::kRow, + ThreadShape::kColumn / Policy::LaneMmaShape::kN + >; + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + +private: + + /// Internal reference + mctlass::TensorRef ref_; + + /// Extent of tensor + MatrixCoord extent_; + + /// Origin + MatrixCoord origin_; + + /// Used to conditionally enable extents checking + bool divisible_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaSimtTileIterator(): divisible_(true) { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaSimtTileIterator( + TensorRef ref, + int lane_id + ): extent_(Shape::kRow, Shape::kColumn), divisible_(true) { + + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + MatrixCoord(0, Policy::LaneMmaShape::kN); + + origin_ = lane_offset; + + ref.add_coord_offset(lane_offset); + + ref_.reset(ref.data(), ref.stride(0)); + } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaSimtTileIterator( + TensorRef ref, + TensorCoord extent, + int lane_id + ): extent_(extent), divisible_(false) { + + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + MatrixCoord(0, Policy::LaneMmaShape::kN); + + origin_ = lane_offset; + + ref.add_coord_offset(lane_offset); + + ref_.reset(ref.data(), ref.stride(0)); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_tile_offset(TensorCoord const &coord) { + + TensorCoord coord_offset( + coord.row() * Shape::kRow, + coord.column() * Shape::kColumn); + + origin_ += coord_offset; + + ref_.add_coord_offset(coord_offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator++() { + + ref_.add_coord_offset({Shape::kRow, 0}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator--() { + + ref_.add_coord_offset({-Shape::kRow, 0}); + + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. (scalar loads) + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kRow; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Iterations::kColumn; ++n) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Policy::LaneMmaShape::kN; ++i) { + + MatrixCoord offset(k, n * Policy::WarpShape::kColumn * Policy::LaneMmaShape::kN + i); + + MatrixCoord access_coord = origin_ + offset; + + int frag_idx = n * Policy::LaneMmaShape::kN + i + k * Iterations::kColumn; + + if (divisible_ || + (access_coord.row() < extent_.row() && access_coord.column() < extent_.column())) { + + frag[frag_idx] = *(ref_.data() + ref_.offset(offset) + pointer_offset); + } + else { + frag[frag_idx] = Element(); + } + } + } + } + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) const { + + Array const *src_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kM; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Iterations::kN; ++n) { + *(ref_.data() + ref_.offset({k, n * Policy::WarpShape::kN}) + pointer_offset / Policy::LaneMmaShape::kN) = + src_ptr[n + k * Iterations::kN]; + } + } + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store(Fragment const &frag, Index pointer_offset) const { + store_with_pointer_offset(frag, 0); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation here + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization for C operands of column-major layouts +/// +/// Concept: MutableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of A elements + typename Element_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_ +> +class MmaSimtTileIterator { +public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = Element_; + + /// Layout of accumulators in memory + using Layout = layout::ColumnMajor; + + /// Decomposition of elements among threads + using Policy = Policy_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + // + // Derived quantities + // + + static_assert( + (!(Shape::kRow % Policy::WarpShape::kRow)) && (!(Shape::kColumn % Policy::WarpShape::kColumn)), + "Warp-level GEMM shape must be divisible by the arrangement of threads in the warp."); + + static_assert(Shape::kRow > 0, "Shape::kRow must be greater than zero."); + static_assert(Shape::kColumn > 0, "Shape::kColumn must be greater than zero."); + static_assert(Policy::WarpShape::kRow > 0, "Policy::WarpShape::kRow must be greater than zero."); + static_assert(Policy::WarpShape::kColumn > 0, "Policy::WarpShape::kColumn must be greater than zero."); + static_assert(Shape::kRow / Policy::WarpShape::kRow > 0, "Shape::kRow / Policy::WarpShape::kRow must be greater than zero."); + static_assert(Shape::kColumn / Policy::WarpShape::kColumn > 0, "Shape::kColumn / Policy::WarpShape::kColumn must be greater than zero."); + + /// Thraed-level shape of a fragment + using ThreadShape = MatrixShape< + Shape::kRow / Policy::WarpShape::kRow, + Shape::kColumn / Policy::WarpShape::kColumn + >; + + static_assert( + (!(ThreadShape::kRow % Policy::LaneMmaShape::kM)) && (!(ThreadShape::kColumn % Policy::LaneMmaShape::kN)), + "Warp-level GEMM shape must be divisible by the arrangement of threads in the warp."); + + /// Number of individual loads + using Iterations = MatrixShape< + ThreadShape::kRow / Policy::LaneMmaShape::kM, + ThreadShape::kColumn / Policy::LaneMmaShape::kN + >; + + using Delta = MatrixShape< + Policy::WarpShape::kRow * Policy::LaneMmaShape::kM, + Policy::WarpShape::kColumn * Policy::LaneMmaShape::kN + >; + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + +private: + + TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaSimtTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaSimtTileIterator( + TensorRef const &ref, + int lane_id + ): + ref_(ref) { + + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + MatrixCoord(Policy::LaneMmaShape::kM, Policy::LaneMmaShape::kN); + + ref_.add_coord_offset(lane_offset); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_tile_offset(TensorCoord const &coord) { + + ref_.add_coord_offset({ + coord.row() * Shape::kRow, + coord.column() * Shape::kColumn}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator++() { + + ref_.add_coord_offset({Shape::kRow, 0}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator--() { + + ref_.add_coord_offset({-Shape::kRow, 0}); + + return *this; + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_HOST_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to be loaded from memory + Index pointer_offset) const { ///< linear offset (in units of Element) when loading + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Iterations::kN; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::LaneMmaShape::kN; ++n) { + + Array const *src_ptr = + reinterpret_cast const *>( + ref_.data() + pointer_offset + ref_.offset({0, mma_n * Delta::kN + n})); + + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Iterations::kM; ++mma_m) { + + Array *dst_ptr = + reinterpret_cast *>(&frag) + + mma_m + Iterations::kM * (n + mma_n * Policy::LaneMmaShape::kN); + + *dst_ptr = src_ptr[mma_m * Policy::WarpShape::kM]; + } + } + } + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) const { + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Iterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Policy::LaneMmaShape::kN; ++n) { + + Array *dst_ptr= + reinterpret_cast *>( + ref_.data() + pointer_offset + ref_.offset({0, mma_n * Delta::kColumn + n})); + + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Iterations::kRow; ++mma_m) { + + Array const *src_ptr = + reinterpret_cast const *>(&frag) + + mma_m + Iterations::kRow * (n + mma_n * Policy::LaneMmaShape::kN); + + dst_ptr[mma_m * Policy::WarpShape::kRow] = *src_ptr; + } + } + } + } + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization for C operands of row-major layouts +/// +/// Concept: MutableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of A elements + typename Element_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_ +> +class MmaSimtTileIterator { +public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = Element_; + + /// Layout of accumulators in memory + using Layout = layout::RowMajor; + + /// Decomposition of elements among threads + using Policy = Policy_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + // + // Derived quantities + // + + static_assert( + (!(Shape::kRow % Policy::WarpShape::kRow)) && (!(Shape::kColumn % Policy::WarpShape::kColumn)), + "Warp-level GEMM shape must be divisible by the arrangement of threads in the warp."); + + static_assert(Shape::kRow > 0, "Shape::kRow must be greater than zero."); + static_assert(Shape::kColumn > 0, "Shape::kColumn must be greater than zero."); + static_assert(Policy::WarpShape::kRow > 0, "Policy::WarpShape::kRow must be greater than zero."); + static_assert(Policy::WarpShape::kColumn > 0, "Policy::WarpShape::kColumn must be greater than zero."); + static_assert(Shape::kRow / Policy::WarpShape::kRow > 0, "Shape::kRow / Policy::WarpShape::kRow must be greater than zero."); + static_assert(Shape::kColumn / Policy::WarpShape::kColumn > 0, "Shape::kColumn / Policy::WarpShape::kColumn must be greater than zero."); + + /// Thraed-level shape of a fragment + using ThreadShape = MatrixShape< + Shape::kRow / Policy::WarpShape::kRow, + Shape::kColumn / Policy::WarpShape::kColumn + >; + + static_assert( + (!(ThreadShape::kRow % Policy::LaneMmaShape::kM)) && (!(ThreadShape::kColumn % Policy::LaneMmaShape::kN)), + "Warp-level GEMM shape must be divisible by the arrangement of threads in the warp."); + + /// Number of individual loads + using Iterations = MatrixShape< + ThreadShape::kRow / Policy::LaneMmaShape::kM, + ThreadShape::kColumn / Policy::LaneMmaShape::kN + >; + + using Delta = MatrixShape< + Policy::WarpShape::kRow * Policy::LaneMmaShape::kM, + Policy::WarpShape::kColumn * Policy::LaneMmaShape::kN + >; + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + +private: + + TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaSimtTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaSimtTileIterator( + TensorRef const &ref, + int lane_id + ): + ref_(ref) { + + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + MatrixCoord(Policy::LaneMmaShape::kM, Policy::LaneMmaShape::kN); + + ref_.add_coord_offset(lane_offset); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_tile_offset(TensorCoord const &coord) { + + ref_.add_coord_offset({ + coord.row() * Shape::kRow, + coord.column() * Shape::kColumn}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator++() { + + ref_.add_coord_offset({Shape::kRow, 0}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator--() { + + ref_.add_coord_offset({-Shape::kRow, 0}); + + return *this; + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_HOST_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to be loaded from memory + Index pointer_offset) const { ///< linear offset (in units of Element) when loading + + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Iterations::kRow; ++mma_m) { + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Policy::LaneMmaShape::kM; ++m) { + + Array const *src_ptr = + reinterpret_cast const *>( + ref_.data() + pointer_offset + ref_.offset({mma_m * Delta::kRow + m, 0})); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Iterations::kColumn; ++mma_n) { + + Array *dst_ptr = + reinterpret_cast *>(&frag) + + mma_n + Iterations::kColumn * (m + mma_m * Policy::LaneMmaShape::kM); + + *dst_ptr = src_ptr[mma_n * Policy::WarpShape::kColumn]; + } + } + } + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) const { + + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Iterations::kRow; ++mma_m) { + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Policy::LaneMmaShape::kM; ++m) { + + Array *dst_ptr = + reinterpret_cast *>( + ref_.data() + pointer_offset + ref_.offset({mma_m * Delta::kRow + m, 0})); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Iterations::kColumn; ++mma_n) { + + Array const *src_ptr = + reinterpret_cast const *>(&frag) + + mma_n + Iterations::kColumn * (m + mma_m * Policy::LaneMmaShape::kM); + + dst_ptr[mma_n * Policy::WarpShape::kColumn] = *src_ptr; + } + } + } + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization for A operands of column-major-K interleaved layouts +/// +/// Concept: MutableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of A elements + typename Element_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK, + /// Number of KGroups per kPartition + int PartitionGroupSize +> +class MmaSimtTileIterator, Policy_, PartitionsK, PartitionGroupSize> { +public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kA; + + /// Element type + using Element = Element_; + + /// Layout of policy + using Layout = layout::ColumnMajorInterleaved<4> ; + + /// Decomposition of elements among threads + using Policy = Policy_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Iterleave factor + static const int kInterleave = 4; + + /// Number of partitions along K dimension + static const int kPartitionsK = PartitionsK; + + /// Number of KGroups per kPartition + static const int kGroupPerTile = PartitionGroupSize / Shape::kColumn; + + // + // Derived quantities + // + + static_assert(!(Shape::kRow % Policy::WarpShape::kRow), + "The warp-level GEMM M size must be divisible by the number of threads arranged along the M dimension."); + + static_assert(Shape::kRow > 0, "Shape::kRow must be greater than zero."); + static_assert(Shape::kColumn > 0, "Shape::kColumn must be greater than zero."); + static_assert(Policy::WarpShape::kRow > 0, "Policy::WarpShape::kRow must be greater than zero."); + static_assert(Shape::kRow / Policy::WarpShape::kRow > 0, "Shape::kRow / Policy::WarpShape::kRow must be greater than zero."); + + /// Thread-level shape of a fragment + using ThreadShape = MatrixShape< + Shape::kRow / Policy::WarpShape::kRow, + Shape::kColumn + >; + + static_assert(!(ThreadShape::kRow % Policy::LaneMmaShape::kM) && !(ThreadShape::kColumn % Policy::LaneMmaShape::kK), + "Thread-level GEMM must be divisible by Policy::LaneMmaShape."); + + /// Number of individual loads + using Iterations = MatrixShape< + ThreadShape::kRow / Policy::LaneMmaShape::kM, + ThreadShape::kColumn / Policy::LaneMmaShape::kK + >; + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + +private: + + /// Internal reference + mctlass::TensorRef, layout::ColumnMajorInterleaved<4>> ref_; + + /// group index within tile + int k_group_idx_; + +public: + MCTLASS_HOST_DEVICE + MmaSimtTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaSimtTileIterator( + TensorRef ref, + int lane_id + ) { + + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + MatrixCoord(Policy::LaneMmaShape::kM, 0); + + ref.add_coord_offset(lane_offset); + + k_group_idx_ = 0; + ref_.reset(reinterpret_cast *>(ref.data()), ref.stride(0)/Policy::LaneMmaShape::kMK); + } + + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_tile_offset(TensorCoord const &coord) { + + ref_.add_coord_offset({ + coord.row() * Shape::kRow / Policy::LaneMmaShape::kMK, + coord.column() * Shape::kColumn}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator++() { + + add_tile_offset({0, 1}); + + if (kPartitionsK > 1) { + ++k_group_idx_; + // Jump to next stage + if (k_group_idx_ == kGroupPerTile) { + k_group_idx_ = 0; + add_tile_offset({0, kGroupPerTile * (kPartitionsK-1)}); + } + } + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator--() { + + ref_.add_coord_offset({0, -Shape::kColumn}); + + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + Array *dst_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kColumn; ++k) { + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Iterations::kRow; ++m) { + + dst_ptr[m + k * Iterations::kRow] = + *((ref_.data() + ref_.offset({m * Policy::WarpShape::kRow / kInterleave, + k*Policy::LaneMmaShape::kK}) + pointer_offset / Policy::LaneMmaShape::kM)); + } + } + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) const { + + Array const *src_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kN; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Iterations::kM; ++m) { + *(ref_.data() + ref_.offset(m * Policy::WarpShape::kM, k) + pointer_offset / Policy::LaneMmaShape::kM) = + src_ptr[m + k * Iterations::kM]; + } + } + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation here + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization for B operands of row-major k-interleaved layouts +/// +/// Concept: MutableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of A elements + typename Element_, + /// Shape of the warp in units of thread (concept: MmaSimtPolicy) + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK, + /// Number of KGroups per kPartition + int PartitionGroupSize +> +class MmaSimtTileIterator, Policy_, PartitionsK, PartitionGroupSize> { +public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kB; + + /// Element type + using Element = Element_; + + /// Layout of policy + using Layout = layout::RowMajorInterleaved<4>; + + /// Decomposition of elements among threads + using Policy = Policy_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Interleave factor + static const int kInterleave = 4; + + /// Number of partitions along K dimension + static const int kPartitionsK = PartitionsK; + + /// Number of KGroups per kPartition + static const int kGroupPerTile = PartitionGroupSize / Shape::kRow; + + // + // Derived quantities + // + + static_assert(!(Shape::kColumn % Policy::WarpShape::kColumn), + "The warp-level GEMM N size must be divisible by the number of threads arranged along the N dimension."); + + static_assert(Shape::kRow > 0, "Shape::kRow must be greater than zero."); + static_assert(Shape::kColumn > 0, "Shape::kColumn must be greater than zero."); + static_assert(Policy::WarpShape::kColumn > 0, "Policy::WarpShape::kColumn must be greater than zero."); + static_assert(Shape::kColumn / Policy::WarpShape::kColumn > 0, "Shape::kColumn / Policy::WarpShape::kColumn must be greater than zero."); + + /// Thread-level shape of a fragment + using ThreadShape = MatrixShape< + Shape::kRow, + Shape::kColumn / Policy::WarpShape::kColumn + >; + + static_assert(!(ThreadShape::kColumn % Policy::LaneMmaShape::kN) && !(ThreadShape::kRow % Policy::LaneMmaShape::kK), + "Thread-level GEMM must be divisible by Policy::LaneMmaShape."); + + /// Number of individual loads + using Iterations = MatrixShape< + ThreadShape::kRow / Policy::LaneMmaShape::kK, + ThreadShape::kColumn / Policy::LaneMmaShape::kN + >; + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + + +private: + + /// Internal reference + mctlass::TensorRef, layout::RowMajorInterleaved<4>> ref_; + + /// group index within tile + int k_group_idx_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaSimtTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaSimtTileIterator( + TensorRef ref, + int lane_id + ) { + + // compute offset based on thread ID and lane layout + typename Policy::LaneLayout lane_layout = Policy::get_lane_layout(); + + MatrixCoord lane_offset = lane_layout.inverse(lane_id) * + MatrixCoord(0, Policy::LaneMmaShape::kN); + + ref.add_coord_offset(lane_offset); + + k_group_idx_ = 0; + + ref_.reset( + reinterpret_cast *>(ref.data()), + ref.stride(0) / Policy::LaneMmaShape::kKN); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaSimtTileIterator &add_tile_offset(TensorCoord const &coord) { + + ref_.add_coord_offset({ + coord.row() * Shape::kRow, + coord.column() * Shape::kColumn / Policy::LaneMmaShape::kKN}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator++() { + + add_tile_offset({1, 0}); + + if (kPartitionsK > 1) { + ++k_group_idx_; + // Jump to next stage + if (k_group_idx_ == kGroupPerTile) { + k_group_idx_ = 0; + add_tile_offset({kGroupPerTile * (kPartitionsK-1), 0}); + } + } + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaSimtTileIterator & operator--() { + + ref_.add_coord_offset({-Shape::kRow, 0}); + + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + Array *dst_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kRow; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Iterations::kColumn; ++n) { + dst_ptr[n + k * Iterations::kColumn] = + *(ref_.data() + ref_.offset({k * Policy::LaneMmaShape::kK, + n * Policy::WarpShape::kColumn / kInterleave}) + pointer_offset / Policy::LaneMmaShape::kN); + } + } + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) const { + + Array const *src_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kM; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Iterations::kN; ++n) { + *(ref_.data() + ref_.offset({k, n * Policy::WarpShape::kN}) + pointer_offset / Policy::LaneMmaShape::kN) = + src_ptr[n + k * Iterations::kN]; + } + } + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store(Fragment const &frag, Index pointer_offset) const { + store_with_pointer_offset(frag, 0); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation here + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_sparse_tensor_op.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_sparse_tensor_op.h new file mode 100644 index 0000000..998eb62 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_sparse_tensor_op.h @@ -0,0 +1,341 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate + operations targeting sparse Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/platform/platform.h" + +#include "mctlass/numeric_conversion.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/arch/mma_sm75.h" +#include "mctlass/arch/mma_sm80.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/warp/mma_tensor_op_policy.h" +#include "mctlass/gemm/warp/mma_tensor_op.h" + +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sparse.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK_ = 1, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Used for partial specialization + typename Enable = bool +> +class SparseMmaTensorOp { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of multiplicand A + using ElementA = ElementA_; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of multiplicand B + using ElementB = ElementB_; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of accumulator matrix C + using ElementC = ElementC_; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Equivalant base dense mma + using Base = MmaTensorOp; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Base::ArchMmaOperator; + + /// Indicates math operator + using MathOperator = typename ArchMmaOperator::Operator; + + /// Architecture tag from underlying instruction + using ArchTag = typename Base::ArchTag; + + /// Indicates class of matrix operator + using OperatorClass = typename Base::OperatorClass; + + /// Shape of underlying instruction + using InstructionShape = typename Base::InstructionShape; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = Base::kTransformA; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = Base::kTransformB; + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// Sparsity in Operand A + static int const kSparse = Policy::Operator::kSparse; + + /// Meta data size in bits + static int const kMetaSizeInBits = Policy::Operator::kMetaSizeInBits; + + /// Max ID2 + static int const kMaxID2 = Policy::Operator::kMaxID2; + + /// Data type of meta E that is moved at the same time + using ElementE = + typename mctlass::platform::conditional::type; + + /// Number of ElementA that is associated with one ElementE + static int const kElementsPerElementE = + 128 / mctlass::sizeof_bits::value; + + /// Meta data is essentially interleaved but mapped to ColumnMajor internally + static int const kInterleaved = 2; + + /// Layout of meta E + using LayoutE = mctlass::layout::ColumnMajor; + + public: + + /// Iterates over the A operand in memory + using IteratorA = MmaTensorOpMultiplicandTileIterator< + MatrixShape, Operand::kA, ElementA, + LayoutA, + MatrixShape, + Policy::OpDelta::kRow, kThreadCount, kPartitionsK>; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentA = + Array; + + /// Iterates over the B operand in memory + using IteratorB = typename Base::IteratorB; + + /// Storage for B tile + using FragmentB = typename Base::FragmentB; + + /// Storage for transformed B tile + using TransformedFragmentB = typename Base::TransformedFragmentB; + + /// Iterates over the C operand in memory + using IteratorC = typename Base::IteratorC; + + /// Storage for C tile + using FragmentC = typename Base::FragmentC; + + /// Iterates over the E operand in memory + using IteratorE = SparseMmaTensorOpMetaTileIterator< + MatrixShape, + ElementE, LayoutE, + MatrixShape, + Policy::OpDelta::kRow, kThreadCount, kPartitionsK>; + + /// Storage for E tile + using FragmentE = typename IteratorE::Fragment; + + /// Number of mma operations performed + using MmaIterations = typename Base::MmaIterations; + +public: + + /// Underlying matrix multiply operator (concept: arch::Mma) + ArchMmaOperator mma; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + SparseMmaTensorOp() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + TransformedFragmentA const &A, + TransformedFragmentB const &B, + FragmentC const &C, + FragmentE const &E + ) const { + + using MmaOperandA = typename Policy::Operator::FragmentA; + using MmaOperandB = typename Policy::Operator::FragmentB; + using MmaOperandC = typename Policy::Operator::FragmentC; + using MmaOperandE = typename Policy::Operator::FragmentE; + + //#if defined(__MACA_ARCH__) + #if 0 //mcTlass not support sparse + + D = C; + + MmaOperandA const *ptr_A = reinterpret_cast(&A); + MmaOperandB const *ptr_B = reinterpret_cast(&B); + MmaOperandC *ptr_D = reinterpret_cast(&D); + MmaOperandE const *ptr_E = reinterpret_cast(&E); + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < MmaIterations::kRow; ++m) { + + int id2 = m % kMaxID2; + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n); + + if (AccumulatorsInRowMajor) { // matrix B is reordered + mma( + ptr_D[n_serpentine + m * MmaIterations::kColumn], + ptr_A[m], + ptr_B[n_serpentine], + ptr_D[n_serpentine + m * MmaIterations::kColumn], + ptr_E[(m / kMaxID2)], + id2); + } else { + mma(ptr_D[m + n_serpentine * MmaIterations::kRow], + ptr_A[m], + ptr_B[n_serpentine], + ptr_D[m + n_serpentine * MmaIterations::kRow], + ptr_E[(m / kMaxID2)], + id2); + } + } + } + #else + assert(0); + #endif + } + + /// Transform the mma operands to the required types + MCTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + + //#if defined(__MACA_ARCH__) + #if 0 //mcTlass not support sparse + // + // Define conversions from source type to instruction type + // + FloatRoundStyle const kRoundA = + PreferredRoundingMode::kRound; + FloatRoundStyle const kRoundB = + PreferredRoundingMode::kRound; + detail::ConvertAndPack + convert_A; + NumericArrayConverter + convert_B; + Array const *ptr_A = + reinterpret_cast const *>(&A); + Array * + ptr_dst_A = reinterpret_cast *>(&dst_A); + + dst_B = convert_B(B); + + ptr_dst_A[0] = convert_A(ptr_A[0]); + ptr_dst_A[1] = convert_A(ptr_A[1]); + #else + assert(0); + #endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op.h new file mode 100644 index 0000000..d8a9ef7 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op.h @@ -0,0 +1,433 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations targeting + Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/platform/platform.h" + +#include "mctlass/numeric_conversion.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/arch/mma_sm75.h" +#include "mctlass/arch/mma_sm80.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/warp/mma_tensor_op_policy.h" + +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +template +struct ConvertAndPack { + + using Converter = NumericArrayConverter; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &source) { + Converter converter; + + return converter(source); + } +}; + +template +struct ConvertAndPack { + + MCTLASS_HOST_DEVICE + Array operator()(Array const &source) { + return source; + } +}; + +template +struct ConvertAndPack { + + using Converter = NumericArrayConverter; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &source) { + Converter converter; + + Array tmp; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + int idx = (((i << 1) & 2) | ((i >> 1) & 1) | (i & 0xfffffffc)); + tmp[i] = source[idx]; + } + + return converter(tmp); + } +}; + +template +struct ConvertAndPack { + + using Converter = NumericArrayConverter; + + MCTLASS_HOST_DEVICE + Array operator()(Array const &source) { + Converter converter; + + Array tmp; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + int idx = (((i << 1) & 2) | ((i >> 1) & 1) | (i & 0xfffffffc)); + tmp[i] = source[idx]; + } + + return converter(tmp); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK_ = 1, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Used for partial specialization + typename Enable = bool +> +class MmaTensorOp { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of multiplicand A + using ElementA = ElementA_; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of multiplicand B + using ElementB = ElementB_; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of accumulator matrix C + using ElementC = ElementC_; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Policy::Operator; + + /// Indicates math operator + using MathOperator = typename ArchMmaOperator::Operator; + + /// Architecture tag from underlying instruction + using ArchTag = typename ArchMmaOperator::ArchTag; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassTensorOp; + + /// Shape of underlying instruction + using InstructionShape = typename ArchMmaOperator::Shape; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = ComplexTransform::kNone; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = ComplexTransform::kNone; + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaTensorOpMultiplicandTileIterator< + MatrixShape, Operand::kA, ElementA, LayoutA, + MatrixShape, + Policy::OpDelta::kRow, kThreadCount, kPartitionsK>; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentA = + Array; + + /// Iterates over the B operand in memory + using IteratorB = MmaTensorOpMultiplicandTileIterator< + MatrixShape, Operand::kB, ElementB, LayoutB, + MatrixShape, + Policy::OpDelta::kRow, kThreadCount, kPartitionsK>; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed B tile + using TransformedFragmentB = + Array; + + /// Iterates over the C operand in memory + using IteratorC = MmaTensorOpAccumulatorTileIterator< + MatrixShape, ElementC, LayoutC, + typename ArchMmaOperator::Shape, typename Policy::OpDelta>; + + /// Storage for C tile + using FragmentC = typename IteratorC::Fragment; + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + (Shape::kM + ArchMmaOperator::Shape::kM - 1) / ArchMmaOperator::Shape::kM, + (Shape::kN + ArchMmaOperator::Shape::kN - 1) / ArchMmaOperator::Shape::kN + >; + +public: + + /// Underlying matrix multiply operator (concept: arch::Mma) + ArchMmaOperator mma; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaTensorOp() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + TransformedFragmentA const &A, + TransformedFragmentB const &B, + FragmentC const &C + ) const { + + using MmaOperandA = typename ArchMmaOperator::FragmentA; + using MmaOperandB = typename ArchMmaOperator::FragmentB; + using MmaOperandC = typename ArchMmaOperator::FragmentC; + + D = C; + + MmaOperandA const *ptr_A = reinterpret_cast(&A); + MmaOperandB const *ptr_B = reinterpret_cast(&B); + MmaOperandC *ptr_D = reinterpret_cast(&D); + + //#if defined(__MACA_ARCH__) && (__TLASS_ARCH__ < 800) + #if 0 + // Serpentine visitation order maximizing reuse of Rb + // The visitation order is like + // _ + // | | | | + // | | | | + // |_| |_| + // + // Down Up Down Up + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < MmaIterations::kRow; ++m) { + + int m_serpentine = ((n % 2) ? (MmaIterations::kRow - 1 - m) : m); + + if (AccumulatorsInRowMajor) { // matrix B is reordered + mma( + ptr_D[n + m_serpentine * MmaIterations::kColumn], + ptr_A[m_serpentine], + ptr_B[n], + ptr_D[n + m_serpentine * MmaIterations::kColumn]); + } else { + mma( + ptr_D[m_serpentine + n * MmaIterations::kRow], + ptr_A[m_serpentine], + ptr_B[n], + ptr_D[m_serpentine + n * MmaIterations::kRow]); + } + } + } + #elif defined(__MACA_ARCH__) + // Serpentine visitation order maximizing reuse of Ra + // The visitation order is like + // _________ + // _________| + // |_________ + // __________| + // + // Right Left Right Left + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < MmaIterations::kRow; ++m) { + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n); + + if (AccumulatorsInRowMajor) { // matrix B is reordered + mma( + ptr_D[n_serpentine + m * MmaIterations::kColumn], + ptr_A[m], + ptr_B[n_serpentine], + ptr_D[n_serpentine + m * MmaIterations::kColumn]); + } else { + mma(ptr_D[m + n_serpentine * MmaIterations::kRow], + ptr_A[m], + ptr_B[n_serpentine], + ptr_D[m + n_serpentine * MmaIterations::kRow]); + } + } + } + #else + assert(0); + #endif + } + + /// Transform the mma operands to the required types + MCTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + + // + // Define conversions from source type to instruction type + // + FloatRoundStyle const kRoundA = + PreferredRoundingMode::kRound; + FloatRoundStyle const kRoundB = + PreferredRoundingMode::kRound; + //#if defined(__MACA_ARCH__) && (__TLASS_ARCH__ < 800) + #if 0 + detail::ConvertAndPack + convert_A; + NumericArrayConverter + convert_B; + Array const *ptr_B = + reinterpret_cast const *>(&B); + Array * + ptr_dst_B = reinterpret_cast *>(&dst_B); + + dst_A = convert_A(A); + + ptr_dst_B[0] = convert_B(ptr_B[0]); + ptr_dst_B[1] = convert_B(ptr_B[1]); + + #elif defined(__MACA_ARCH__) + detail::ConvertAndPack + convert_A; + NumericArrayConverter + convert_B; + Array const *ptr_A = + reinterpret_cast const *>(&A); + Array * + ptr_dst_A = reinterpret_cast *>(&dst_A); + + dst_B = convert_B(B); + + ptr_dst_A[0] = convert_A(ptr_A[0]); + ptr_dst_A[1] = convert_A(ptr_A[1]); + #else + assert(0); + #endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "mctlass/gemm/warp/mma_tensor_op_fast_f32.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_fast_f32.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_fast_f32.h new file mode 100644 index 0000000..0a7c39f --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_fast_f32.h @@ -0,0 +1,471 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations targeting + Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/platform/platform.h" + +#include "mctlass/numeric_conversion.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/mma_sm80.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/warp/mma_tensor_op_policy.h" +#include "mctlass/gemm/warp/mma_tensor_op.h" + +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +enum class TensorFloat32Op { + k3xTF32, + k4xTF32 +}; + +template < + /// Floating-point rounding style + FloatRoundStyle RoundBigA_, + /// Floating-point rounding style + FloatRoundStyle RoundSmallA_, + /// Floating-point rounding style + FloatRoundStyle RoundBigB_ = RoundBigA_, + /// Floating-point rounding style + FloatRoundStyle RoundSmallB_ = RoundSmallA_, + /// Precision for TensorFloat32Op + // (k3xTF32: BigxBig, BigxSmall, SmallxBig) + // (k4xTF32: BigxBig, BigxSmall, SmallxBig, SmallxSmall) + TensorFloat32Op Precision_ = TensorFloat32Op::k3xTF32 + > +struct FastF32 { + + static FloatRoundStyle const kRoundBigA = RoundBigA_; + static FloatRoundStyle const kRoundSmallA = RoundSmallA_; + static FloatRoundStyle const kRoundBigB = RoundBigB_; + static FloatRoundStyle const kRoundSmallB = RoundSmallB_; + static TensorFloat32Op const kPrecision = Precision_; +}; + + +namespace detail { + + template< + int N, + FloatRoundStyle RoundBig = FloatRoundStyle::round_toward_zero, + FloatRoundStyle RoundSmall = FloatRoundStyle::round_half_ulp_truncate + > + struct ConvertAndPackAccurateF32 { + + /// Rounding styles for big and small part + static FloatRoundStyle const kRoundBig = RoundBig; + static FloatRoundStyle const kRoundSmall = RoundSmall; + + /// Converter type + using Converter = NumericConverterFastF32; + + /// Source fragement + using SourceFragment = Array; + + /// Destination fragment + using DestinationFragment = Array; + + /// Converter Fragment holding two tfloat32_t elements for every float + using ConverterFragment = Array; + + /// Index in fargments for the big and small part + static int const kBigIndex = 0; + static int const kSmallIndex = 1; + + MCTLASS_HOST_DEVICE + void operator()(SourceFragment const &source, + DestinationFragment &dst_big, + DestinationFragment &dst_small) { + + Converter convert_; + ConverterFragment result_; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + // convert source to result fragment + result_ = convert_(source[i]); + + // store converted result fragments to destination fragment + dst_big[i] = result_[kBigIndex]; + dst_small[i] = result_[kSmallIndex]; + } + } + }; +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK_ = 1, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Used for partial specialization + typename Enable = bool +> +class MmaTensorOpFastF32; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for float*float+float => float using TF32 TensorOps +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK_, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor, + /// Used for partial specialization + typename Enable +> +class MmaTensorOpFastF32< + Shape_, + float, LayoutA_, + float, LayoutB_, + float, LayoutC_, + Policy_, PartitionsK_, + AccumulatorsInRowMajor, Enable> { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of multiplicand A + using ElementA = float; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of multiplicand B + using ElementB = float; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of accumulator matrix C + using ElementC = float; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Policy::Operator; + + /// Indicates math operator + using MathOperator = arch::OpMultiplyAddFastF32; + + /// Architecture tag from underlying instruction + using ArchTag = typename ArchMmaOperator::ArchTag; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassTensorOp; + + /// Shape of underlying instruction + using InstructionShape = typename ArchMmaOperator::Shape; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = ComplexTransform::kNone; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = ComplexTransform::kNone; + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// Tune F32 to TF32 big small conversion for float operation + /// Different combination of big small conversin can cause different tradeoff + /// between speed and accuracy. Generally, use round_half_ulp_truncate can + /// improve the performance but hur the accuracy. + using MmaFastF32 = FastF32 < + FloatRoundStyle::round_toward_zero, // kRoundBigA + FloatRoundStyle::round_half_ulp_truncate, // kRoundSmallA + FloatRoundStyle::round_toward_zero, // kRoundBigB + FloatRoundStyle::round_half_ulp_truncate, // kRoundSmallB + TensorFloat32Op::k3xTF32 // Number of TF32 operations + >; + +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kA, + ElementA, + LayoutA, + MatrixShape, + Policy::OpDelta::kRow, + kThreadCount, + kPartitionsK + >; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentA = + Array; + + /// Fragment bisecting big and small sections + using AccessTypeFragmentA = + Array; + + /// Iterates over the B operand in memory + using IteratorB = MmaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kB, + ElementB, + LayoutB, + MatrixShape, + Policy::OpDelta::kRow, + kThreadCount, + kPartitionsK + >; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed B tile + using TransformedFragmentB = + Array; + + /// Fragment bisecting big and small sections + using AccessTypeFragmentB = + Array; + + /// Index in fargments for the big and small part + static int const kBigIndex = 0; + static int const kSmallIndex = 1; + + /// Iterates over the C operand in memory + using IteratorC = MmaTensorOpAccumulatorTileIterator< + MatrixShape, ElementC, LayoutC, + typename ArchMmaOperator::Shape, typename Policy::OpDelta>; + + /// Storage for C tile + using FragmentC = typename IteratorC::Fragment; + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + (Shape::kM + ArchMmaOperator::Shape::kM - 1) / ArchMmaOperator::Shape::kM, + (Shape::kN + ArchMmaOperator::Shape::kN - 1) / ArchMmaOperator::Shape::kN + >; + +public: + + /// Underlying matrix multiply operator (concept: arch::Mma) + ArchMmaOperator mma; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaTensorOpFastF32() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + TransformedFragmentA const &A, + TransformedFragmentB const &B, + FragmentC const &C + ) const { + + AccessTypeFragmentA const *ptr_A = reinterpret_cast(&A); + AccessTypeFragmentB const *ptr_B = reinterpret_cast(&B); + + // + // Accumulate in place + // + D = C; + + mma_operator(D, ptr_A[kSmallIndex], ptr_B[kBigIndex], D); + + mma_operator(D, ptr_A[kBigIndex], ptr_B[kSmallIndex], D); + + mma_operator(D, ptr_A[kBigIndex], ptr_B[kBigIndex], D); + + if (MmaFastF32::kPrecision == TensorFloat32Op::k4xTF32) + mma_operator(D, ptr_A[kSmallIndex], ptr_B[kSmallIndex], D); + } + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void mma_operator( + FragmentC &D, + AccessTypeFragmentA const &A, + AccessTypeFragmentB const &B, + FragmentC const &C + ) const { + + #if defined(__MACA_ARCH__) + + using MmaOperandA = typename ArchMmaOperator::FragmentA; + using MmaOperandB = typename ArchMmaOperator::FragmentB; + using MmaOperandC = typename ArchMmaOperator::FragmentC; + + MmaOperandA const *ptr_A = reinterpret_cast(&A); + MmaOperandB const *ptr_B = reinterpret_cast(&B); + MmaOperandC *ptr_D = reinterpret_cast(&D); + + // Serpentine visitation order maximizing reuse of Ra + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < MmaIterations::kRow; ++m) { + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + // This allows to reuse of Rb when at serpentine turns + int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n); + + if (AccumulatorsInRowMajor) { // matrix B is reordered + mma( + ptr_D[n_serpentine + m * MmaIterations::kColumn], + ptr_A[m], + ptr_B[n_serpentine], + ptr_D[n_serpentine + m * MmaIterations::kColumn]); + } else { + mma( + ptr_D[m + n_serpentine * MmaIterations::kRow], + ptr_A[m], + ptr_B[n_serpentine], + ptr_D[m + n_serpentine * MmaIterations::kRow]); + } + } // end n loop + } // end m loop + #else + assert(0); + #endif + } + + /// Transform the mma operands to the required types + MCTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + + // + // Define conversions from source type to instruction type + // + #if defined(__MACA_ARCH__) + + detail::ConvertAndPackAccurateF32< + FragmentA::kElements / 2, + MmaFastF32::kRoundBigA, + MmaFastF32::kRoundSmallA> convert_A; + + detail::ConvertAndPackAccurateF32< + FragmentB::kElements, + MmaFastF32::kRoundBigB, + MmaFastF32::kRoundSmallB> convert_B; + + Array *ptr_dst_B = + reinterpret_cast *>(&dst_B); + + convert_B(B, ptr_dst_B[0], ptr_dst_B[1]); + + Array *ptr_dst_A = + reinterpret_cast *>(&dst_A); + + Array const *ptr_A = + reinterpret_cast const *>(&A); + + convert_A(ptr_A[0], ptr_dst_A[0], ptr_dst_A[2]); + + convert_A(ptr_A[1], ptr_dst_A[1], ptr_dst_A[3]); + #else + assert(0); + #endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_fragment_iterator.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_fragment_iterator.h new file mode 100644 index 0000000..2e356aa --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_fragment_iterator.h @@ -0,0 +1,528 @@ +/*! \file + \brief This defines a "fragment" iterator for visiting the fragments of a warp tile + that participate in one warp-level mma operation. + + Typically, this is used to access the accumulator tile/fragement of a warp-level mma operation. + The accumulator tile is then partitioned into smaller tiles/fragments that can be fed into + next warp-level mma operation. + + This iterator is necessary to accomplish warp-level mma fusion where the accumulator tile is + reused as multiplicand tile for the next mma. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/numeric_conversion.h" + +namespace mctlass { +namespace gemm { +namespace warp { + + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Size of the accumulation tile shape (concept: MatrixShape) + typename AccumulatorShape_, + /// KBlocks columns to compute residual + int KBlocksColumn_, + /// Accumulator Element type + typename ElementAccumulator_, + /// Element type + typename Element_, + /// Layout of operand in memory + typename Layout_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Output operation on the fragment + typename OutputOp_> +class MmaTensorOpFragmentIterator; + + +// Partial specialization for col-major accumulator tile + +template < + /// Shape of warp tile to load (concept: MatrixShape) + typename Shape_, + /// Shape of the warp accumulation tile (concept: MatrixShape) + typename AccumulatorShape_, + /// KBlocks columns to compute residual + int KBlocksColumn_, + /// Accumulator Element type + typename ElementAccumulator_, + /// Element type + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Output operation on fragment + typename OutputOp_> +class MmaTensorOpFragmentIterator { + public: + + /// Shape of warp tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Shape of the warp accumulation tile (concept: MatrixShape) + using AccumulatorShape = AccumulatorShape_; + + /// KBlocks columns to compute residual + static int const kKBlockColumn = KBlocksColumn_; + + /// Accumulator Element type + using ElementAccumulator = ElementAccumulator_; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::ColumnMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Output operation on fragment + using OutputOp = OutputOp_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kRow % InstructionShape::kM) && + !(Shape::kColumn % InstructionShape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + static_assert( + AccumulatorShape::kRow == Shape::kRow, + "Rows of Warp Accumulator must be the same as rows of warp"); + static_assert( + !(AccumulatorShape::kColumn % Shape::kColumn), + "Shape of Warp Accumulator must be divisible by warp shape."); + static_assert( + !(kKBlockColumn % Shape::kColumn), + "KBlock size must be divisible by warp shape."); + + /// Number of times this iterator can be incremented + static int const kIterations = AccumulatorShape::kCount / Shape::kCount; + }; + +private: + + static int const kElementsPerAccess = InstructionShape::kM * InstructionShape::kN / kThreads; + + /// Number of mma operations performed by a warp + using MmaIterations = MatrixShape; + /// Number of mma operations performed by the entire accumulator + using AccumulatorIterations = MatrixShape; + + /// Number of K iterations + static int const kKBlockIterations = (AccumulatorShape::kColumn + kKBlockColumn - 1) / kKBlockColumn; + static int const kResidualColumn = AccumulatorShape::kColumn - (kKBlockIterations - 1) * kKBlockColumn; + static int const kKBlockColumnIterations = kKBlockColumn / Shape::kColumn + * (AccumulatorShape::kRow / Shape::kRow); + static int const kResidualIndex = kResidualColumn / Shape::kColumn + * (AccumulatorShape::kRow / Shape::kRow); + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array; + + /// Accumulator Fragment object + using AccumulatorFragment = Array; + + /// Scale Bias Element Type + using ElementScaleBias = typename OutputOp::ElementCompute; + + /// Scale Bias Fragment object + using ScaleBiasFragment = Array; + + +private: + + /// Internal access type + using AccessType = Array; + using FragmentAccessType = Array; + + using ScaleBiasAccessType = Array; + +private: + // + // Data members + // + + /// Accumulator tile + AccessType const *accumulators_; + + /// Internal index + int index_; + + /// Used to access residual tile first + bool is_residual_tile_; + +public: + /// Constructs an iterator + MCTLASS_HOST_DEVICE + MmaTensorOpFragmentIterator(AccumulatorFragment const &accum) + : accumulators_(reinterpret_cast(&accum)), + index_(0), is_residual_tile_(true) {} + + /// Add offset + MCTLASS_HOST_DEVICE + void add_offset(int index_offset) { + index_ += index_offset; + if(is_residual_tile_ && index_ >= kKBlockColumnIterations) { + index_ = index_ - kKBlockColumnIterations + kResidualIndex; + is_residual_tile_ = false; + } + } + + /// Increments + MCTLASS_HOST_DEVICE + MmaTensorOpFragmentIterator &operator++() { + add_offset(1); + return *this; + } + + /// Decrements + MCTLASS_HOST_DEVICE + MmaTensorOpFragmentIterator &operator--() { + add_offset(-1); + return *this; + } + + /// Loads a fragment from the referenced part of the accumulator tile + MCTLASS_HOST_DEVICE + void load(Fragment &frag, OutputOp output_op) const { + + if (output_op.is_source_needed()) //beta must be zero + assert(0); + + FragmentAccessType *frag_ptr = reinterpret_cast(&frag); + + int index = index_ * MmaIterations::kCount; + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; n++) { + for (int m = 0; m < MmaIterations::kRow; m++) { + int accumulator_access_offset = + n * AccumulatorIterations::kRow + m + index; + + frag_ptr[m * MmaIterations::kColumn + n].clear(); + if(!(is_residual_tile_ && index_ >= kResidualIndex)) + frag_ptr[m * MmaIterations::kColumn + n] = output_op(accumulators_[accumulator_access_offset]); + } + } + } + + /// Loads a fragment from the referenced part of the accumulator tile + /// Then apply per-channel scale and bias + MCTLASS_HOST_DEVICE + void load(Fragment &frag, ScaleBiasFragment &scale, + ScaleBiasFragment &bias, OutputOp output_op) const { + + if (output_op.is_source_needed()) //beta must be zero + assert(0); + + FragmentAccessType *frag_ptr = reinterpret_cast(&frag); + ScaleBiasAccessType * scale_ptr = reinterpret_cast(&scale); + ScaleBiasAccessType * bias_ptr = reinterpret_cast(&bias); + + int index = index_ * MmaIterations::kCount; + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; n++) { + for (int m = 0; m < MmaIterations::kRow; m++) { + int accumulator_access_offset = + n * AccumulatorIterations::kRow + m + index; + + frag_ptr[m * MmaIterations::kColumn + n].clear(); + if(!(is_residual_tile_ && index_ >= kResidualIndex)) + frag_ptr[m * MmaIterations::kColumn + n] = + output_op(accumulators_[accumulator_access_offset], + scale_ptr[n] /*scale*/, bias_ptr[n] /*bias*/); + } + } + } + + + +}; + +// Partial specialization for row-major accumulator tile + +template < + /// Shape of warp tile to load (concept: MatrixShape) + typename Shape_, + /// Shape of the warp accumulation tile (concept: MatrixShape) + typename AccumulatorShape_, + /// KBlocks columns to compute residual + int KBlocksColumn_, + /// Accumulator Element type + typename ElementAccumulator_, + /// Element type + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Output operation on fragment + typename OutputOp_> +class MmaTensorOpFragmentIterator { + public: + + /// Shape of warp tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Shape of the warp accumulation tile (concept: MatrixShape) + using AccumulatorShape = AccumulatorShape_; + + /// KBlocks columns to compute residual + static int const kKBlockColumn = KBlocksColumn_; + + /// Accumulator Element type + using ElementAccumulator = ElementAccumulator_; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Output operation on fragment + using OutputOp = OutputOp_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kRow % InstructionShape::kM) && + !(Shape::kColumn % InstructionShape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + static_assert( + AccumulatorShape::kRow == Shape::kRow, + "Rows of Warp Accumulator must be the same as rows of warp"); + static_assert( + !(AccumulatorShape::kColumn % Shape::kColumn), + "Shape of Warp Accumulator must be divisible by warp shape."); + static_assert( + !(kKBlockColumn % Shape::kColumn), + "KBlock size must be divisible by warp shape."); + + /// Number of times this iterator can be incremented + static int const kIterations = AccumulatorShape::kCount / Shape::kCount; + }; + +private: + + static int const kRowsPerIteration = 8; + static int const kColumnsPerIteration = 16; + static int const kElementsPerIteration = kRowsPerIteration * InstructionShape::kN / kThreads; + static int const kElementsPerAccess = kRowsPerIteration * kColumnsPerIteration / kThreads; + static int const kIterationsPerAccess = kElementsPerAccess / kElementsPerIteration; + + // Number of iterations per actual instruction + static int const kIterationsPerInstruction = InstructionShape::kM / kRowsPerIteration; + + static int const kAccessStride = kIterationsPerInstruction; + + /// Number of mma operations performed by a warp + using MmaIterations = MatrixShape; + /// Number of mma operations performed by the entire accumulator + using AccumulatorIterations = MatrixShape; + + /// Number of Accesses in a warp + using AccessIterations = MatrixShape; + + /// Number of K iterations + static int const kKBlockIterations = (AccumulatorShape::kColumn + kKBlockColumn - 1) / kKBlockColumn; + static int const kResidualColumn = AccumulatorShape::kColumn - (kKBlockIterations - 1) * kKBlockColumn; + static int const kKBlockColumnIterations = kKBlockColumn / Shape::kColumn; + static int const kResidualIndex = kResidualColumn / Shape::kColumn; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + /// This is the fragment size produced by one access of the iterator. + using Fragment = Array; + + /// Accumulator Fragment object + using AccumulatorFragment = Array; + + /// Scale Bias Element Type + using ElementScaleBias = typename OutputOp::ElementCompute; + + /// Scale Bias Fragment object + using ScaleBiasFragment = Array; + + +private: + + /// Internal access type + using AccessType = Array; + using FragmentAccessType = Array; + using ScaleBiasAccessType = Array; + +private: + // + // Data members + // + + /// Accumulator tile + AccessType const *accumulators_; + + /// Internal index + int index_; + + /// Used to access residual tile first + bool is_residual_tile_; + +public: + /// Constructs an iterator + MCTLASS_HOST_DEVICE + MmaTensorOpFragmentIterator(AccumulatorFragment const &accum) + : accumulators_(reinterpret_cast(&accum)), + index_(0), is_residual_tile_(true) {} + + /// Add offset + MCTLASS_HOST_DEVICE + void add_offset(int index_offset) { + index_ += index_offset; + if(is_residual_tile_ && index_ >= kKBlockColumnIterations) { + index_ = index_ - kKBlockColumnIterations + kResidualIndex; + is_residual_tile_ = false; + } + } + + /// Increments + MCTLASS_HOST_DEVICE + MmaTensorOpFragmentIterator &operator++() { + add_offset(1); + return *this; + } + + /// Decrements + MCTLASS_HOST_DEVICE + MmaTensorOpFragmentIterator &operator--() { + add_offset(-1); + return *this; + } + + MCTLASS_HOST_DEVICE + void set_index(int idx) { + index_ = idx; + } + + /// Loads a fragment from the referenced part of the accumulator tile + MCTLASS_HOST_DEVICE + void load(Fragment &frag, OutputOp output_op) const { + + if (output_op.is_source_needed()) //beta must be zero + assert(0); + + FragmentAccessType *frag_ptr = reinterpret_cast(&frag); + + int index = index_ * AccessIterations::kCount; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < AccessIterations::kCount; i++) { + + int accumulator_access_offset = index / AccessIterations::kCount * (MmaIterations::kColumn * kIterationsPerInstruction) + + (index % AccessIterations::kCount) / (AccessIterations::kColumn * kIterationsPerInstruction) * + AccumulatorIterations::kColumn * kIterationsPerInstruction + + (index % (AccessIterations::kColumn * kIterationsPerInstruction)) / kIterationsPerInstruction * + (kIterationsPerInstruction * kIterationsPerAccess) + + (index % kIterationsPerInstruction); + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < kIterationsPerAccess; j++) { + + frag_ptr[i*kIterationsPerAccess + j].clear(); + if(!(is_residual_tile_ && index_ >= kResidualIndex)) + frag_ptr[i*kIterationsPerAccess + j] = output_op(accumulators_[accumulator_access_offset + j * kAccessStride]); + } + index++; + } + } + + /// Loads a fragment from the referenced part of the accumulator tile + /// Then apply per-channel scale and bias + MCTLASS_HOST_DEVICE + void load(Fragment &frag, ScaleBiasFragment &scale, + ScaleBiasFragment & bias, OutputOp output_op) const { + + if (output_op.is_source_needed()) //beta must be zero + assert(0); + + FragmentAccessType *frag_ptr = reinterpret_cast(&frag); + ScaleBiasAccessType * scale_ptr = reinterpret_cast(&scale); + ScaleBiasAccessType * bias_ptr = reinterpret_cast(&bias); + + int index = index_ * AccessIterations::kCount; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < AccessIterations::kCount; i++) { + + int accumulator_access_offset = index / AccessIterations::kCount * (MmaIterations::kColumn * kIterationsPerInstruction) + + (index % AccessIterations::kCount) / (AccessIterations::kColumn * kIterationsPerInstruction) * + AccumulatorIterations::kColumn * kIterationsPerInstruction + + (index % (AccessIterations::kColumn * kIterationsPerInstruction)) / kIterationsPerInstruction * + (kIterationsPerInstruction * kIterationsPerAccess) + + (index % kIterationsPerInstruction); + + int scale_bias_offset = (index + % (kIterationsPerInstruction * AccessIterations::kColumn)) + * kIterationsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int j = 0; j < kIterationsPerAccess; j++) { + + + frag_ptr[i*kIterationsPerAccess + j].clear(); + if(!(is_residual_tile_ && index_ >= kResidualIndex)) + frag_ptr[i*kIterationsPerAccess + j] = output_op( + accumulators_[accumulator_access_offset + j * kAccessStride], + scale_ptr[scale_bias_offset + j], bias_ptr[scale_bias_offset + j]); + } + index++; + } + } + +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_policy.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_policy.h new file mode 100644 index 0000000..e04a647 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_policy.h @@ -0,0 +1,65 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Policy describing implementation details of warp-level GEMM targeting Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/gemm/gemm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Policy +template < + typename Operator_, ///< hardware instruction(s) performing TensorOp (concept: arch::Mma) + typename OpDelta_ ///< distance between operations (concept: MatrixShape) +> +struct MmaTensorOpPolicy { + + using Operator = Operator_; ///< hardware instruction(s) performing TensorOp (concept: arch::Mma) + using OpDelta = OpDelta_; ///< distance between operations (concept: MatrixShape) + using MmaShape = typename Operator::Shape; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_sm70.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_sm70.h new file mode 100644 index 0000000..e8b474b --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_sm70.h @@ -0,0 +1,280 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations targeting + Tensor Cores. + + This is a work in progress. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/mma.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/warp/mma_tensor_op_policy.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm70.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Used for partial specialization + typename Enable = bool +> +class MmaVoltaTensorOp { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of multiplicand A + using ElementA = ElementA_; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of multiplicand B + using ElementB = ElementB_; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of accumulator matrix C + using ElementC = ElementC_; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassTensorOp; + + /// Architecture tag + using ArchTag = arch::Sm70; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Policy::Operator; + + /// Indicates math operator + using MathOperator = typename ArchMmaOperator::Operator; + + /// Underlying instruction shape + using InstructionShape = typename ArchMmaOperator::Shape; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = ComplexTransform::kNone; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = ComplexTransform::kNone; + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 32; + + /// interleaved 32x32 tiles + using InterleavedTileShape = GemmShape<32, 32, 4>; + + static_assert(!(Shape::kM % InterleavedTileShape::kM) && + !(Shape::kN % InterleavedTileShape::kN), + "Shape must be a multiple of InterleavedTileShape."); +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaVoltaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kA, + ElementA, + LayoutA, + MatrixShape< + ArchMmaOperator::Shape::kM, + ArchMmaOperator::Shape::kK + >, + Policy::OpDelta::kRow, + kThreadCount + >; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Iterates over the B operand in memory + using IteratorB = MmaVoltaTensorOpMultiplicandTileIterator< + MatrixShape, + Operand::kB, + ElementB, + LayoutB, + MatrixShape< + ArchMmaOperator::Shape::kK, + ArchMmaOperator::Shape::kN + >, + Policy::OpDelta::kRow, + kThreadCount + >; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Iterates over the C operand in memory + using IteratorC = MmaVoltaTensorOpAccumulatorTileIterator< + MatrixShape, + ElementC, + LayoutC, + typename ArchMmaOperator::Shape, + typename Policy::OpDelta + >; + + /// Storage for C tile + using FragmentC = typename IteratorC::Fragment; + +private: + + static_assert( + !(Shape::kM % ArchMmaOperator::Shape::kM) && + !(Shape::kN % ArchMmaOperator::Shape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + InterleavedTileShape::kM / ArchMmaOperator::Shape::kM, + InterleavedTileShape::kN / ArchMmaOperator::Shape::kN + >; + using TileIterations = MatrixShape< + Shape::kM / InterleavedTileShape::kM, + Shape::kN / InterleavedTileShape::kN + >; + + // Whether matrix B is reordered + bool reorder_B_; + +public: + + /// Underlying matrix multiply operator (concept: arch::Mma) + ArchMmaOperator mma; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaVoltaTensorOp() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + FragmentA const &A, + FragmentB const &B, + FragmentC const &C) { + + using MmaOperandA = typename ArchMmaOperator::FragmentA; + using MmaOperandB = typename ArchMmaOperator::FragmentB; + using MmaOperandC = typename ArchMmaOperator::FragmentC; + + D = C; + + MmaOperandA const *ptr_A = reinterpret_cast(&A); + MmaOperandB const *ptr_B = reinterpret_cast(&B); + MmaOperandC *ptr_D = reinterpret_cast(&D); + + MCTLASS_PRAGMA_UNROLL + for (int outer_col = 0; outer_col < TileIterations::kColumn; ++outer_col) { + MCTLASS_PRAGMA_UNROLL + for (int inner_col = 0; inner_col < MmaIterations::kColumn; ++inner_col) { + MCTLASS_PRAGMA_UNROLL + for (int outer_row = 0; outer_row < TileIterations::kRow; ++outer_row) { + MCTLASS_PRAGMA_UNROLL + + for (int inner_row = 0; inner_row < MmaIterations::kRow; ++inner_row) { + + int op_col = inner_col + MmaIterations::kColumn * outer_col; + + // Column-major serpentine sequence to maximize reuse of A operand. + int inner_row_serp = inner_row; + int outer_row_serp = outer_row; + if (op_col & 1) { + inner_row_serp = MmaIterations::kRow - inner_row - 1; + outer_row_serp = TileIterations::kRow - outer_row - 1; + } + int op_row = inner_row_serp + MmaIterations::kRow * outer_row_serp; + int op_idx = inner_row_serp + MmaIterations::kRow * + (inner_col + MmaIterations::kColumn * + (outer_row_serp + TileIterations::kRow * outer_col)); + mma( + ptr_D[op_idx], + ptr_A[op_row], + ptr_B[op_col], + ptr_D[op_idx]); + + } + } + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_access_iterator.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_access_iterator.h new file mode 100644 index 0000000..21f96e1 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_access_iterator.h @@ -0,0 +1,362 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines iterators used by warp-level matrix multiply operations targeting Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/gemm/gemm.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor_op_multiplicand_sm80.h" + +#include "mctlass/platform/platform.h" +#include "mctlass/fast_math.h" + +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + + +/// Tile access iterator +/// Each iteration acess in the tile is +/// used as multiplicand for one +/// warp-level matrix multiplication +template < + /// Size of the tile (concept: MatrixShape) + typename Shape_, + /// Operand identity + Operand Operand_, + /// Data type of A elements + typename Element_, + /// Layout of operand + typename Layout_, + /// Shape of one matrix production operation (concept: MatrixShape) + typename InstructionShape_, + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + int OpDelta_, + /// Number of threads participating in one matrix operation + int Threads = 32, + /// Enable Residual Support + bool EnableResidual = false, + /// Number of partitions along K dimension + int PartitionsK_ = 1 +> +class MmaTensorOpMultiplicandTileAccessIterator { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + /// Basic check + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = Layout_; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Number of elements accessed per Shared Memory load + static int const kElementsPerAccess = + (sizeof_bits::value >= 32 ? 1 : 32 / sizeof_bits::value); + + using InstructionCount = MatrixShape< + Shape::kRow / InstructionShape::kRow, + Shape::kColumn / InstructionShape::kColumn + >; + + static int const kIterations = (kOperand == Operand::kA) ? + InstructionCount::kColumn : InstructionCount::kRow; + + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array< + Element, + (kOperand == Operand::kA) ? + (Shape::kRow * InstructionShape::kColumn / kThreads) : + (Shape::kColumn * InstructionShape::kRow / kThreads) + >; + + /// Memory access type + using AccessType = AlignedArray; + +private: + + /// Underlying tensor reference + TensorRef ref_; + + /// Extent of tensor + MatrixCoord extent_; + + /// Origin + MatrixCoord origin_; + + /// Used to load residual tile + bool is_residual_; + + /// residual offset of each thread + TensorCoord residual_offset_; + + /// Iterations in a tile + int iterations_; + +public: + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileAccessIterator( + TensorRef const &ref, + TensorCoord extent, + int lane_id + ): ref_(ref), extent_(extent), is_residual_(false), iterations_(0) { + + if (kOperand == Operand::kA) { + origin_ = MatrixCoord(lane_id / 4, (lane_id % 4) * kElementsPerAccess); + } + else { + origin_ = MatrixCoord((lane_id % 4) * kElementsPerAccess, lane_id / 4); + } + + ref_.add_coord_offset(origin_); + + if(EnableResidual) { + // compute residual offset + if (kOperand == Operand::kA) { + typename TensorCoord::Index residual_size = + extent_.column() % Shape::kColumn; + if(residual_size) { + is_residual_ = true; + residual_offset_ = make_Coord(0, residual_size); + } + } + else { + typename TensorCoord::Index residual_size = + extent_.row() % Shape::kRow; + if(residual_size) { + is_residual_ = true; + residual_offset_ = make_Coord(residual_size, 0); + } + } + } + } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileAccessIterator( + TensorRef const &ref, + int lane_id + ): MmaTensorOpMultiplicandTileAccessIterator(ref, + {Shape::kRow, Shape::kColumn}, lane_id) { + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileAccessIterator &add_tile_offset(TensorCoord const &tile_offset) { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + origin_ += coord_offset; + + ref_.add_coord_offset(coord_offset); + + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + void advance() { + + if(EnableResidual && is_residual_) { + is_residual_ = false; + + origin_ += residual_offset_; + ref_.add_coord_offset(residual_offset_); + + } + + else { + if (kOperand == Operand::kA) { + add_tile_offset({0, 1}); + } + else { + add_tile_offset({1, 0}); + } + } + + iterations_ = 0; + } + + /// increase iterations in a tile + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileAccessIterator & operator++() { + + iterations_++; + + if(iterations_ >= kIterations) + advance(); + + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + int const kWarpShapeDivisibleInner = + (kOperand == Operand::kA ? InstructionShape::kColumn : InstructionShape::kRow); + + // Take advantage of Tensor Op's 8 x 4T access pattern + int const kAccessesInner = (kWarpShapeDivisibleInner / kElementsPerAccess) / 4; + + AccessType *access_ptr = reinterpret_cast(&frag); + + if (kOperand == Operand::kA) { + int const kTilesPerInstruction = InstructionShape::kRow / 8; + + MCTLASS_PRAGMA_UNROLL + for (int inst_m_idx = 0; inst_m_idx < InstructionCount::kRow; ++inst_m_idx) { + + MCTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { + + MCTLASS_PRAGMA_UNROLL + for (int access_m_idx = 0; access_m_idx < kTilesPerInstruction; ++access_m_idx) { + int access_idx = + access_m_idx + kTilesPerInstruction * (inner_idx + kAccessesInner * inst_m_idx); + + MatrixCoord offset( + access_m_idx * 8 + inst_m_idx * InstructionShape::kRow, + inner_idx * 4 * kElementsPerAccess + iterations_ * InstructionShape::kColumn); + + MatrixCoord access_coord = origin_ + offset; + +// if(access_coord.row() < extent_.row() && access_coord.column() < extent_.column()) { + + access_ptr[access_idx] = *reinterpret_cast( + ref_.data() + ref_.offset(offset)); +// } +// else { +// AccessType zero; +// zero.clear(); +// access_ptr[access_idx] = zero; +// } + } + } + } + } + else { + MCTLASS_PRAGMA_UNROLL + for (int inst_n_idx = 0; inst_n_idx < InstructionCount::kColumn; ++inst_n_idx) { + + MCTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { + int access_idx = inner_idx + kAccessesInner * inst_n_idx; + + MatrixCoord offset( + inner_idx * 4 * kElementsPerAccess + iterations_ * InstructionShape::kRow, + inst_n_idx * 8); + + MatrixCoord access_coord = origin_ + offset; + +// if(access_coord.row() < extent_.row() && access_coord.column() < extent_.column()) { + + access_ptr[access_idx] = *reinterpret_cast( + ref_.data() + ref_.offset(offset)); +// } +// else { +// AccessType zero; +// zero.clear(); +// access_ptr[access_idx] = zero; +// } + } + } + } + } + +}; + + + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator.h new file mode 100644 index 0000000..b7208d1 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator.h @@ -0,0 +1,10135 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines iterators used by warp-level matrix multiply operations targeting Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/gemm/gemm.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" + +#include "mctlass/platform/platform.h" +#include "mctlass/fast_math.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Operand identity + Operand Operand, + /// Data type of A elements + typename Element_, + /// Layout of operand + typename Layout_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + int OpDelta_, + /// Number of threads participating in one matrix operation + int Threads, + /// Number of partitions along K dimension + int PartitionsK_ = 1> +class MmaTensorOpMultiplicandTileIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to load from shared +/// memory and therefore must be initialized with a TensorRef to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::TensorOpMultiplicandCongruous::value, + 64>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCongruous< + sizeof_bits::value, 64>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + // Determine number of elements along outer dimension per individual LDSM op + static int const kLdsmOpOuter = Layout::kElementsPerAccess; + static int const kLdsmOpInner = 8; + + static_assert(!(Shape::kContiguous % kLdsmOpOuter), + "Shape of warp-level mma must be divisible by LDSM's fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsmOpInner), + "Shape of warp-level mma must be divisible by LDSM's fundamental tile size."); + + /// Shape of one individual LDSM instruction + static int const LdsmShapeStrided = + InstructionShape::kStrided / kLdsmOpInner; + static int const LdsmShapeContiguous = 4 / LdsmShapeStrided; + using LdsmShape = + layout::PitchLinearShape; + + /// Number and arrangement of LDSM instructions + using LdsmIterations = layout::PitchLinearShape< + Shape::kContiguous / Layout::kElementsPerAccess / LdsmShapeContiguous, + 1>; + + /// Number of groups for each tile + static int const kGroupsPerTile = + Shape::kStrided / InstructionShape::kStrided; + }; + +private: + + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Number of internal pointers needed to reference shared memory + static int const kPointerCount = + Layout::TileShape::kContiguous / Policy::LdsmShape::kContiguous; + + /// Pointer type used for accesses + using AccessType = Array; + + /// Internal counter used to jump to next K partition + int k_group_idx_; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + +private: + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_[kPointerCount]; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(): stride_(0), byte_offset_(0) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): + stride_(ref.stride(0) / Layout::kElementsPerAccess), byte_offset_(0), + k_group_idx_(0) { + + int quad_pair = (lane_id >> 3); + int quad_quad = (lane_id >> 4); + int lane_in_quad = (lane_id & 3); + int lane_in_quad_pair = (lane_id & 7); + int lane_in_quad_quad = (lane_id & 15); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPointerCount; ++i) { + int partition_contiguous_idx = -1; + int access_contiguous_idx = -1; + int access_strided_idx = -1; + + if (Policy::LdsmShape::kContiguous == 4) { + // Matrix multiply 1688 A/B + // Q0 Q1 Q2 Q3 (Q stands for 1 8x128bit block). + // Four blocks are next to each other in the contiguous dimension. + partition_contiguous_idx = ((lane_in_quad_pair >> 2) ^ i); + access_contiguous_idx = (quad_pair ^ lane_in_quad); + access_strided_idx = lane_in_quad_pair; + } + else if (Policy::LdsmShape::kContiguous == 2 && + kOperand == Operand::kA) { + // Matrix multiply 16816 A + // Q0 Q1 + // Q2 Q3 + partition_contiguous_idx = ((lane_in_quad_pair >> 2) ^ (i >> 1)); + access_contiguous_idx = + (((quad_pair & 1) + ((i & 1) << 1)) ^ lane_in_quad); + access_strided_idx = lane_in_quad_pair + (lane_id >> 4 << 3); + } else if (Policy::LdsmShape::kContiguous == 2 && + kOperand == Operand::kB) { + // Matrix multiply 16816 B + // Q0 Q2 + // Q1 Q3 + partition_contiguous_idx = ((lane_in_quad_pair >> 2) ^ (i >> 1)); + access_contiguous_idx = ((quad_quad + ((i & 1) << 1)) ^ lane_in_quad); + access_strided_idx = lane_in_quad_quad; + } else if (Policy::LdsmShape::kContiguous == 1) { + // Matrix multiply 16832.SP B + // Q0 + // Q1 + // Q2 + // Q3 + partition_contiguous_idx = ((lane_in_quad_pair >> 2) ^ (i >> 2)); + access_contiguous_idx = ((i & 3) ^ lane_in_quad); + access_strided_idx = lane_id; + } + + int access_contiguous = + partition_contiguous_idx * Layout::PartitionShape::kContiguous + + access_contiguous_idx; + + int access_strided = access_strided_idx; + + pointer_[i] = reinterpret_cast(ref.data()) + + access_contiguous + access_strided * stride_; + } + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + byte_offset_ += offset * sizeof(Element); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + int contiguous_offset = tile_offset.contiguous(); + if (Shape::kContiguous == + Layout::PartitionShape::kContiguous * Layout::kElementsPerAccess) { + if (tile_offset.contiguous() % 2) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPointerCount / 2; ++i) { + AccessType const *tmp_pointer = pointer_[i]; + pointer_[i] = pointer_[i + kPointerCount / 2]; + pointer_[i + kPointerCount / 2] = tmp_pointer; + } + } + contiguous_offset = (tile_offset.contiguous() >> 1) << 1; + } + + int offset = (tile_offset.strided() * InstructionShape::kStrided) * + stride_ * Layout::kElementsPerAccess + + contiguous_offset * Shape::kContiguous; + + add_pointer_offset(offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + add_tile_offset({0, 1}); + + if (kPartitionsK > 1) { + ++k_group_idx_; + // Jump to next stage + if (k_group_idx_ == Policy::kGroupsPerTile) { + k_group_idx_ = 0; + add_tile_offset( + {0, ((kPartitionsK - 1) * Policy::kGroupsPerTile)}); + } + } + + return *this; + } + + /// Advances the iterator along the opposite of the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + byte_offset_ -= stride_ * InstructionShape::kStrided * sizeof(Element) * + Layout::kElementsPerAccess; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_byte_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + Array *fetch_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsmIterations::kStrided; ++s) { + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsmIterations::kContiguous; ++c) { + + int access_idx = c + s * Policy::LdsmIterations::kContiguous; + + AccessType const *source_ptr = + pointer_[c % kPointerCount] + + Layout::TileShape::kContiguous * (c / kPointerCount) + + Policy::kLdsmOpInner * Policy::LdsmShape::kStrided * s * stride_; + + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_; + + mctlass::arch::ldsm( + fetch_ptr[access_idx], + source_byte_ptr + ); + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = + tile_offset.contiguous() * Shape::kContiguous / Layout::kElementsPerAccess + + tile_offset.strided() * InstructionShape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no op + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread MMA.TF32 NT TensorOps. It +/// uses LDS.32 to load from shared memory and therefore must be initialized +/// with a TensorRef to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::TensorOpMultiplicandCongruous<32, 32>, InstructionShape_, + OpDelta_, 32, PartitionsK_> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for " + "A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCongruous<32, 32>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + // Determine number of elements along outer dimension per individual 32bit + // shared memory load op. Every one warp of 32bit shared memory load loads + // 8x4 elements + static int const kLdsOpInner = Layout::TileShape::kStrided; + static int const kLdsOpOuter = kThreads / kLdsOpInner; + + static_assert(!(Shape::kContiguous % kLdsOpOuter), + "Shape of warp-level mma must be divisible by 32bit " + "fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsOpInner), + "Shape of warp-level mma must be divisible by 32bit " + "fundamental tile size."); + + /// Number of 32 bit shared memory load instructions needed by one MMA instruction + /// 1688 A 2x2 + /// 1688 B 1x2 + /// 16816 B 1x4 + static int const LdsShapeContiguous = + InstructionShape::kContiguous / kLdsOpOuter; + static int const LdsShapeStrided = InstructionShape::kStrided / kLdsOpInner; + using LdsShape = + layout::PitchLinearShape; + + /// Number and arrangement of LDS instructions + using LdsIterations = layout::PitchLinearShape< + Shape::kContiguous / LdsShapeContiguous / kLdsOpOuter, 1>; + + /// Number of groups for each tile + static int const kGroupsPerTile = + Shape::kStrided / InstructionShape::kStrided; + }; + + private: + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Number of internal pointers needed to reference shared memory + static int const kPointerCount = Layout::TileShape::kContiguous * + Layout::kElementsPerAccess / + Policy::kLdsOpOuter; + + /// Vectorized access is not used + static int const kElementsPerAccess = 1; + + /// Pointer type used for accesses + using AccessType = Element; + + /// Internal counter used to jump to next K partition + int k_group_idx_; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + + private: + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_[kPointerCount]; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() : stride_(0), byte_offset_(0) {} + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : stride_(ref.stride(0)), byte_offset_(0), k_group_idx_(0) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPointerCount; ++i) { + int access_strided = lane_id % Policy::kLdsOpInner; + int access_contiguous = (lane_id / Policy::kLdsOpInner) + + (access_strided ^ i) * Policy::kLdsOpOuter; + + pointer_[i] = reinterpret_cast(ref.data()) + + access_contiguous + access_strided * stride_; + } + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + byte_offset_ += offset * sizeof(Element); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + int contiguous_offset = tile_offset.contiguous(); + if (Shape::kContiguous == + Layout::TileShape::kContiguous * Layout::kElementsPerAccess / 2) { + if (tile_offset.contiguous() % 2) { + // Matrix multiply 1688 pointer_[0] <=> pointer_[4] pointer_[1] <=> pointer_[5] + // pointer_[2] <=> pointer_[6] pointer_[3] <=> pointer_[7] + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPointerCount / 2; ++i) { + AccessType const *tmp_pointer = pointer_[i]; + pointer_[i] = pointer_[i + kPointerCount / 2]; + pointer_[i + kPointerCount / 2] = tmp_pointer; + } + } + contiguous_offset = (tile_offset.contiguous() >> 1) << 1; + } + + int offset = (tile_offset.strided() * InstructionShape::kStrided) * stride_ + + contiguous_offset * Shape::kContiguous; + + add_pointer_offset(offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator++() { + add_tile_offset({0, 1}); + + if (kPartitionsK > 1) { + ++k_group_idx_; + // Jump to next stage + if (k_group_idx_ == Policy::kGroupsPerTile) { + k_group_idx_ = 0; + add_tile_offset( + {0, ((kPartitionsK - 1) * Policy::kGroupsPerTile)}); + } + } + + return *this; + } + + /// Advances the iterator along the opposite of the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator--() { + byte_offset_ -= stride_ * InstructionShape::kStrided * sizeof(Element) * + kElementsPerAccess; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { load_with_byte_offset(frag, 0); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + Element *fetch_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsIterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsIterations::kContiguous; ++c) { + MCTLASS_PRAGMA_UNROLL + for (int ss = 0; ss < Policy::LdsShape::kStrided; ++ss) { + MCTLASS_PRAGMA_UNROLL + for (int cc = 0; cc < Policy::LdsShape::kContiguous; ++cc) { + int access_idx = + cc + (ss + (c + s * Policy::LdsIterations::kContiguous) * + Policy::LdsShape::kStrided) * + Policy::LdsShape::kContiguous; + int access_idx_contiguous = cc + c * Policy::LdsShape::kContiguous; + int access_idx_strided = + (ss + s * Policy::LdsShape::kStrided) * Policy::kLdsOpInner; + + AccessType const *source_ptr = + pointer_[access_idx_contiguous % kPointerCount] + + Layout::TileShape::kContiguous * Layout::kElementsPerAccess * + (access_idx_contiguous / kPointerCount) + + access_idx_strided * stride_; + + char const *source_byte_ptr = + reinterpret_cast(source_ptr) + byte_offset + + byte_offset_; + + fetch_ptr[access_idx] = + *reinterpret_cast(source_byte_ptr); + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = + tile_offset.contiguous() * Shape::kContiguous / + Layout::kElementsPerAccess + + tile_offset.strided() * InstructionShape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no op + } +}; + +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, mctlass::tfloat32_t, + mctlass::layout::TensorOpMultiplicandCongruous<32, 32>, InstructionShape_, + OpDelta_, 32, PartitionsK_> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for " + "A or B operands to warp-level Mma."); + + /// Element type + using Element = mctlass::tfloat32_t; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCongruous<32, 32>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + // Determine number of elements along outer dimension per individual 32bit + // shared memory load op. Every one warp of 32bit shared memory load loads + // 8x4 elements + static int const kLdsOpInner = Layout::TileShape::kStrided; + static int const kLdsOpOuter = kThreads / kLdsOpInner; + + static_assert(!(Shape::kContiguous % kLdsOpOuter), + "Shape of warp-level mma must be divisible by 32bit " + "fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsOpInner), + "Shape of warp-level mma must be divisible by 32bit " + "fundamental tile size."); + + /// Number of 32 bit shared memory load instructions needed by one MMA instruction + /// 1688 A 2x2 + /// 1688 B 1x2 + /// 16816 B 1x4 + static int const LdsShapeContiguous = + InstructionShape::kContiguous / kLdsOpOuter; + static int const LdsShapeStrided = InstructionShape::kStrided / kLdsOpInner; + using LdsShape = + layout::PitchLinearShape; + + /// Number and arrangement of LDS instructions + using LdsIterations = layout::PitchLinearShape< + Shape::kContiguous / LdsShapeContiguous / kLdsOpOuter, 1>; + + /// Number of groups for each tile + static int const kGroupsPerTile = + Shape::kStrided / InstructionShape::kStrided; + }; + + private: + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Number of internal pointers needed to reference shared memory + static int const kPointerCount = Layout::TileShape::kContiguous * + Layout::kElementsPerAccess / + Policy::kLdsOpOuter; + + /// Vectorized access is not used + static int const kElementsPerAccess = 1; + + /// Pointer type used for accesses + using AccessType = Element; + + /// Internal counter used to jump to next K partition + int k_group_idx_; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + + private: + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_[kPointerCount]; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() : stride_(0), byte_offset_(0) {} + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : stride_(ref.stride(0)), byte_offset_(0), k_group_idx_(0) { + pointer_[0] = reinterpret_cast(ref.data()); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + byte_offset_ += offset * sizeof(Element); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + int contiguous_offset = tile_offset.contiguous(); + if (Shape::kContiguous == + Layout::TileShape::kContiguous * Layout::kElementsPerAccess / 2) { + if (tile_offset.contiguous() % 2) { + // Matrix multiply 1688 pointer_[0] <=> pointer_[4] pointer_[1] <=> pointer_[5] + // pointer_[2] <=> pointer_[6] pointer_[3] <=> pointer_[7] + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPointerCount / 2; ++i) { + AccessType const *tmp_pointer = pointer_[i]; + pointer_[i] = pointer_[i + kPointerCount / 2]; + pointer_[i + kPointerCount / 2] = tmp_pointer; + } + } + contiguous_offset = (tile_offset.contiguous() >> 1) << 1; + } + + int offset = (tile_offset.strided() * InstructionShape::kStrided) * stride_ + + contiguous_offset * Shape::kContiguous; + + add_pointer_offset(offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator++() { + add_tile_offset({0, 1}); + if (kPartitionsK > 1) { + ++k_group_idx_; + // Jump to next stage + if (k_group_idx_ == Policy::kGroupsPerTile) { + k_group_idx_ = 0; + add_tile_offset( + {0, ((kPartitionsK - 1) * Policy::kGroupsPerTile)}); + } + } + + return *this; + } + + /// Advances the iterator along the opposite of the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator--() { + byte_offset_ -= stride_ * InstructionShape::kStrided * sizeof(Element) * + kElementsPerAccess; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { load_with_byte_offset(frag, 0); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + Element *fetch_ptr = reinterpret_cast(&frag); + int lane_id = __lane_id(); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsIterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsIterations::kContiguous; ++c) { + if(kOperand == Operand::kA) { + if(InstructionShape::kStrided == 8) { + unsigned row = (lane_id & 0xf) + c * InstructionShape::kContiguous; + unsigned col = ((lane_id >> 4) << 1) ^ 0x7; + unsigned col_1 = col % 4; + unsigned row_1 = ((row / 8) ^ col_1) * 8 + (row % 8); + unsigned col_2 = (col - 1) % 4; + unsigned row_2 = ((row / 8) ^ col_2) * 8 + (row % 8); + + AccessType const *source_ptr = pointer_[0] + row_1 + col * stride_; + char const *source_byte_ptr = reinterpret_cast(source_ptr) + + byte_offset + byte_offset_; + + fetch_ptr[c * 4 + 0] = *reinterpret_cast(source_byte_ptr); + AccessType const *source_ptr1 = pointer_[0] + row_2 + (col - 1) * stride_; + char const *source_byte_ptr1 = reinterpret_cast(source_ptr1) + + byte_offset + byte_offset_; + + fetch_ptr[c * 4 + 1] = *reinterpret_cast(source_byte_ptr1); + } + else { + unsigned row = (lane_id & 0xf) + c * InstructionShape::kContiguous; + unsigned col = (lane_id >> 4) ^ 0x3; + unsigned col_1 = col % 4; + unsigned row_1 = ((row / 8) ^ col_1) * 8 + (row % 8); + + AccessType const *source_ptr = pointer_[0] + row_1 + col * stride_; + char const *source_byte_ptr = reinterpret_cast(source_ptr) + + byte_offset + byte_offset_; + fetch_ptr[c * 2 + 0] = *reinterpret_cast(source_byte_ptr); + fetch_ptr[c * 2 + 1] = static_cast(0); + } + } + else { + if(InstructionShape::kStrided == 8) { + unsigned row = ((lane_id >> 4) << 1) ^ 0x7; + unsigned col = (lane_id & 0x7) + c * InstructionShape::kStrided; + unsigned row_1 = row % 4; + unsigned col_1 = ((col / 8) ^ row_1) * 8 + (col % 8); + unsigned row_2 = (row - 1) % 4; + unsigned col_2 = ((col / 8) ^ row_2) * 8 + (col % 8); + + AccessType const *source_ptr = pointer_[0]; + char const *source_byte_ptr = reinterpret_cast(source_ptr) + + byte_offset + byte_offset_; + + fetch_ptr[c * 2 + 0] = *(reinterpret_cast(source_byte_ptr) + + col_1 + row * stride_); + + fetch_ptr[c * 2 + 1] = *(reinterpret_cast(source_byte_ptr) + + col_2 + (row - 1) * stride_); + } + else { + unsigned row = (lane_id >> 4) ^ 0x3; + unsigned col = (lane_id & 0x7) + c * InstructionShape::kContiguous; + unsigned row_1 = row % 4; + unsigned col_1 = ((col / 8) ^ row_1) * 8 + (col % 8); + + AccessType const *source_ptr = pointer_[0]; + char const *source_byte_ptr = reinterpret_cast(source_ptr) + + byte_offset + byte_offset_; + + fetch_ptr[c] = *(reinterpret_cast(source_byte_ptr) + + col_1 + row * stride_); + } + } + } + } + __syncthreads(); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = + tile_offset.contiguous() * Shape::kContiguous / + Layout::kElementsPerAccess + + tile_offset.strided() * InstructionShape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no op + } +}; + +/// +/// This special template for float is completely same with mctlass:tfloat32_t, +/// It's a very bad implement. We need to refactor them one day. +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, float, + mctlass::layout::TensorOpMultiplicandCongruous<32, 32>, InstructionShape_, + OpDelta_, 32, PartitionsK_> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for " + "A or B operands to warp-level Mma."); + + /// Element type + using Element = float; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCongruous<32, 32>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + // Determine number of elements along outer dimension per individual 32bit + // shared memory load op. Every one warp of 32bit shared memory load loads + // 8x4 elements + static int const kLdsOpInner = Layout::TileShape::kStrided; + static int const kLdsOpOuter = kThreads / kLdsOpInner; + + static_assert(!(Shape::kContiguous % kLdsOpOuter), + "Shape of warp-level mma must be divisible by 32bit " + "fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsOpInner), + "Shape of warp-level mma must be divisible by 32bit " + "fundamental tile size."); + + /// Number of 32 bit shared memory load instructions needed by one MMA instruction + /// 1688 A 2x2 + /// 1688 B 1x2 + /// 16816 B 1x4 + static int const LdsShapeContiguous = + InstructionShape::kContiguous / kLdsOpOuter; + static int const LdsShapeStrided = InstructionShape::kStrided / kLdsOpInner; + using LdsShape = + layout::PitchLinearShape; + + /// Number and arrangement of LDS instructions + using LdsIterations = layout::PitchLinearShape< + Shape::kContiguous / LdsShapeContiguous / kLdsOpOuter, 1>; + + /// Number of groups for each tile + static int const kGroupsPerTile = + Shape::kStrided / InstructionShape::kStrided; + }; + + private: + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Number of internal pointers needed to reference shared memory + static int const kPointerCount = Layout::TileShape::kContiguous * + Layout::kElementsPerAccess / + Policy::kLdsOpOuter; + + /// Vectorized access is not used + static int const kElementsPerAccess = 1; + + /// Pointer type used for accesses + using AccessType = Element; + + /// Internal counter used to jump to next K partition + int k_group_idx_; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + + private: + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_[kPointerCount]; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() : stride_(0), byte_offset_(0) {} + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : stride_(ref.stride(0)), byte_offset_(0), k_group_idx_(0) { + pointer_[0] = reinterpret_cast(ref.data()); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + byte_offset_ += offset * sizeof(Element); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + int contiguous_offset = tile_offset.contiguous(); + if (Shape::kContiguous == + Layout::TileShape::kContiguous * Layout::kElementsPerAccess / 2) { + if (tile_offset.contiguous() % 2) { + // Matrix multiply 1688 pointer_[0] <=> pointer_[4] pointer_[1] <=> pointer_[5] + // pointer_[2] <=> pointer_[6] pointer_[3] <=> pointer_[7] + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPointerCount / 2; ++i) { + AccessType const *tmp_pointer = pointer_[i]; + pointer_[i] = pointer_[i + kPointerCount / 2]; + pointer_[i + kPointerCount / 2] = tmp_pointer; + } + } + contiguous_offset = (tile_offset.contiguous() >> 1) << 1; + } + + int offset = (tile_offset.strided() * InstructionShape::kStrided) * stride_ + + contiguous_offset * Shape::kContiguous; + + add_pointer_offset(offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator++() { + add_tile_offset({0, 1}); + if (kPartitionsK > 1) { + ++k_group_idx_; + // Jump to next stage + if (k_group_idx_ == Policy::kGroupsPerTile) { + k_group_idx_ = 0; + add_tile_offset( + {0, ((kPartitionsK - 1) * Policy::kGroupsPerTile)}); + } + } + + return *this; + } + + /// Advances the iterator along the opposite of the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator--() { + byte_offset_ -= stride_ * InstructionShape::kStrided * sizeof(Element) * + kElementsPerAccess; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { load_with_byte_offset(frag, 0); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + Element *fetch_ptr = reinterpret_cast(&frag); + int lane_id = __lane_id(); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsIterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsIterations::kContiguous; ++c) { + if(kOperand == Operand::kA) { + + unsigned row = (lane_id & 0xf) + c * InstructionShape::kContiguous; + unsigned col = ((lane_id >> 4) << 1) ^ 0x7; + unsigned col_1 = col % 4; + unsigned row_1 = ((row / 8) ^ col_1) * 8 + (row % 8); + unsigned col_2 = (col - 1) % 4; + unsigned row_2 = ((row / 8) ^ col_2) * 8 + (row % 8); + + AccessType const *source_ptr = pointer_[0] + row_1 + col * stride_; + char const *source_byte_ptr = reinterpret_cast(source_ptr) + + byte_offset + byte_offset_; + + fetch_ptr[c * 4 + 0] = *reinterpret_cast(source_byte_ptr); + AccessType const *source_ptr1 = pointer_[0] + row_2 + (col - 1) * stride_; + char const *source_byte_ptr1 = reinterpret_cast(source_ptr1) + + byte_offset + byte_offset_; + + fetch_ptr[c * 4 + 1] = *reinterpret_cast(source_byte_ptr1); + } else { + unsigned row = ((lane_id >> 4) << 1) ^ 0x7; + unsigned col = (lane_id & 0x7) + c * InstructionShape::kStrided; + unsigned row_1 = row % 4; + unsigned col_1 = ((col / 8) ^ row_1) * 8 + (col % 8); + unsigned row_2 = (row - 1) % 4; + unsigned col_2 = ((col / 8) ^ row_2) * 8 + (col % 8); + + AccessType const *source_ptr = pointer_[0]; + char const *source_byte_ptr = reinterpret_cast(source_ptr) + + byte_offset + byte_offset_; + + fetch_ptr[c * 2 + 0] = *(reinterpret_cast(source_byte_ptr) + + col_1 + row * stride_); + + fetch_ptr[c * 2 + 1] = *(reinterpret_cast(source_byte_ptr) + + col_2 + (row - 1) * stride_); + } + } + } + __syncthreads(); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = + tile_offset.contiguous() * Shape::kContiguous / + Layout::kElementsPerAccess + + tile_offset.strided() * InstructionShape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no op + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to load from shared +/// memory and therefore must be initialized with a TensorRef to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element_))>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA, + "MmaTensorOpMultiplicandIterator for ColumnMajor Congruous may " + "only be instantiated for A operand to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element_))>; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element_))>, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): iterator_({ref.data(), ref.stride()}, lane_id) { + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.contiguous(), tile_offset.strided()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to load from shared +/// memory and therefore must be initialized with a TensorRef to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::RowMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element_))>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator for RowMajor Congruous may " + "only be instantiated for B operand to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element_))>; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element_))>, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): iterator_({ref.data(), ref.stride()}, lane_id) { + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.strided(), tile_offset.contiguous()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Element number when the layout crosses (in units of elements) + int Crosswise, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::TensorOpMultiplicandCrosswise::value, + Crosswise>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for " + "A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Element number when the layout crosses + static int const kCrosswise = Crosswise; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCrosswise< + sizeof_bits::value, kCrosswise>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + // Determine number of elements along outer dimension per individual LDSM op + static int const kLdsmOpOuter = Layout::kElementsPerAccess; + static int const kLdsmOpInner = 8; + + static_assert(!(Shape::kContiguous % kLdsmOpOuter), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsmOpInner), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + /// Shape of one individual LDSM instruction + static int const LdsmShapeContiguous = + InstructionShape::kContiguous / kLdsmOpOuter; + static int const LdsmShapeStrided = + ((4 / LdsmShapeContiguous * kLdsmOpInner) > Shape::kStrided) + ? (Shape::kStrided / kLdsmOpInner) + : (4 / LdsmShapeContiguous); + using LdsmShape = + layout::PitchLinearShape; + + /// Number and arrangement of LDSM instructions + using LdsmIterations = + layout::PitchLinearShape<1, Shape::kStrided / kLdsmOpInner / + LdsmShape::kStrided>; + + /// + static int const kGroupsPerTile = Layout::TileShape::kContiguous / + Layout::kFactor / LdsmShape::kContiguous; + }; + + private: + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = Array; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + + private: + + /// Total number of sections. The memory is divided into stages. One stage + /// can store one tile. Stage is divided into sections. Interleaved layout + /// can have multiple sections in a stage. The rest layout only has one section + /// in a stage. + int sections_; + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + /// Internal counter used to determine when to increment byte offset and when + /// to XOR it + int k_group_idx_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() + : pointer_(nullptr), + sections_(0), + stride_(0), + byte_offset_(0), + k_group_idx_(0) {} + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : pointer_(reinterpret_cast(ref.data())), + sections_(ref.stride(0) / kCrosswise), + // stride_ = kCrosswise x sections_ x kFactor + stride_(ref.stride(0) * Layout::kFactor / Layout::kElementsPerAccess), + byte_offset_(0), + k_group_idx_(0) { + // Warp level iterator at most use double buffer to hide latency. If there + // are more than 2 sections, every stage should have more than 1 section. + + // Turing silicon requires all 32 threads in a warp provide valid addresses + // even for LDSM.1 and LDSM.2 +//#if defined(__MACA_ARCH__) +#if 0 + lane_id = lane_id % (Policy::LdsmShape::kCount * Policy::kLdsmOpInner); +#endif + + int quad_quad = (lane_id >> 4); + int quad_pair = (lane_id >> 3); + int lane_in_pair = (lane_id & 1); + int lane_in_quad = (lane_id & 3); + int lane_in_quad_pair = (lane_id & 7); + int lane_in_quad_quad = (lane_id & 15); + + int partition_contiguous_idx = -1; + int access_contiguous_idx = -1; + int access_strided_idx = -1; + + if (Layout::kFactor == 4) { + // Super Integer matrix multiply Interleaved-32 + + int factor_in_partition = + (Layout::PartitionShape::kContiguous * Layout::kFactor / + Layout::TileShape::kContiguous); + + if (Policy::LdsmShape::kStrided == Policy::LdsmShape::kCount) { + // Integer matrix multiply 8816 A/B + partition_contiguous_idx = lane_in_quad / factor_in_partition; + access_contiguous_idx = ((lane_in_pair * factor_in_partition) ^ + (lane_in_quad_quad / Layout::kFactor)); + access_strided_idx = lane_id / Layout::kFactor; + } + else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kA) { + // Integer matrix multiply 16832 A + partition_contiguous_idx = lane_in_quad / factor_in_partition; + access_strided_idx = lane_in_quad_quad / Layout::kFactor; + access_contiguous_idx = + ((lane_in_pair * factor_in_partition + quad_quad) ^ + access_strided_idx); + } + else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kB) { + // Integer matrix multiply 16832 B + partition_contiguous_idx = lane_in_quad / factor_in_partition; + access_strided_idx = lane_in_quad_pair / Layout::kFactor + quad_quad * 2; + access_contiguous_idx = + ((lane_in_pair * factor_in_partition + ((lane_id & 8) >> 3)) ^ + access_strided_idx); + } + } else if (Layout::kFactor == 2) { + // Super Matrix multiply kBlock = 32 + if (Policy::LdsmShape::kStrided == Policy::LdsmShape::kCount) { + // Matrix multiply 1688 A/B + // (Q stands for 1 8x128bit block). + // Q0 + // Q1 + // Q2 + // Q3 + // Four blocks are next to each other in the strided dimension. + partition_contiguous_idx = (lane_id % Layout::kFactor); + access_contiguous_idx = (lane_in_quad_pair / Layout::kFactor); + access_strided_idx = lane_id / Layout::kFactor; + } + else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kA) { + // Matrix multiply 16816|1688.TF32 A + // Q0 Q2 + // Q1 Q3 + partition_contiguous_idx = (lane_id % Layout::kFactor); + access_contiguous_idx = + (quad_quad ^ (lane_in_quad_pair / Layout::kFactor)); + access_strided_idx = (lane_in_quad_quad / Layout::kFactor); + } else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kB) { + // Matrix multiply 16816|1688.TF32 B + // Q0 Q1 + // Q2 Q3 + partition_contiguous_idx = (lane_id % Layout::kFactor); + access_contiguous_idx = + ((quad_pair & 1) ^ (lane_in_quad_pair / Layout::kFactor)); + access_strided_idx = + (lane_in_quad_pair + (lane_id >> 4 << 3)) / Layout::kFactor; + } + else if (Policy::LdsmShape::kContiguous == Policy::LdsmShape::kCount) { + // Matrix multiply 16832.SP B + // Q0 Q1 Q2 Q3 + partition_contiguous_idx = (lane_id % Layout::kFactor); + access_contiguous_idx = + (quad_pair ^ (lane_in_quad_pair / Layout::kFactor)); + access_strided_idx = lane_in_quad_pair / Layout::kFactor; + } + } else if (Layout::kFactor == 1) { + // Super Matrix multiply kBlock = 64 + if (Policy::LdsmShape::kStrided == Policy::LdsmShape::kCount) { + // Q0 + // Q1 + // Q2 + // Q3 + partition_contiguous_idx = (lane_in_quad_pair >> 2); + access_contiguous_idx = lane_in_quad; + access_strided_idx = lane_id; + } + else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kA) { + // Matrix multiply 16816|1688.TF32 A + // Q0 Q2 + // Q1 Q3 + partition_contiguous_idx = (lane_in_quad_pair >> 2); + access_contiguous_idx = (quad_quad ^ lane_in_quad); + access_strided_idx = lane_in_quad_quad; + } else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kB) { + // Matrix multiply 16816|1688.TF32 B + // Q0 Q1 + // Q2 Q3 + partition_contiguous_idx = (lane_in_quad_pair >> 2); + access_contiguous_idx = ((quad_pair & 1) ^ lane_in_quad); + access_strided_idx = lane_in_quad_pair + (lane_id >> 4 << 3); + } + else if (Policy::LdsmShape::kContiguous == Policy::LdsmShape::kCount) { + // Matrix multiply 16832.SP B + // Q0 Q1 Q2 Q3 + partition_contiguous_idx = (lane_in_quad_pair >> 2); + access_contiguous_idx = (quad_pair ^ lane_in_quad); + access_strided_idx = lane_in_quad_pair; + } + } + + int access_contiguous = + partition_contiguous_idx * Layout::PartitionShape::kContiguous + + access_contiguous_idx; + + int access_strided = access_strided_idx; + + byte_offset_ = (access_contiguous + access_strided * stride_) * + sizeof_bits::value * Layout::kElementsPerAccess / 8; + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + byte_offset_ += offset * sizeof_bits::value / 8; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + + byte_offset_ ^= k_groups_delta * sizeof_bits::value * + Layout::kElementsPerAccess * + Policy::LdsmShape::kContiguous / 8; + pointer_ += + tile_offset.strided() * stride_ * Shape::kStrided / Layout::kFactor + + whole_tiles * stride_ / sections_; + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative( + TensorCoord const &tile_offset) { + + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + if (k_groups_delta < 0) { + whole_tiles -= 1; + k_groups_delta += Policy::kGroupsPerTile; + } + + if ((Policy::kGroupsPerTile / kPartitionsK) >= 2) { + byte_offset_ ^= (k_groups_delta & 1) * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + if ((Policy::kGroupsPerTile / kPartitionsK) >= 4) { + byte_offset_ ^= ((k_groups_delta + (k_group_idx_ & 1)) & 2) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + if ((Policy::kGroupsPerTile / kPartitionsK) == 8) { + byte_offset_ ^= ((k_groups_delta + (k_group_idx_ & 3)) & 4) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + + k_group_idx_ += k_groups_delta; + whole_tiles += k_group_idx_ / (Policy::kGroupsPerTile / kPartitionsK); + k_group_idx_ = k_group_idx_ % (Policy::kGroupsPerTile / kPartitionsK); + + pointer_ += + tile_offset.strided() * stride_ * Shape::kStrided / Layout::kFactor + + whole_tiles * stride_ / sections_; + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator++() { + + // Integer matrix multiply 16832 Interleaved-32 + // NONE + // Integer matrix multiply 16816 Interleaved-32 || Integer matrix multiply 16816 kblock=32 + + // Integer matrix multiply 8816 Interleaved-32 + // ^1 ^1 + // Matrix multiply 1684.TF32 kblock=16 || Integer matrix multiply 16816 kblock=64 + // Matrix multiply 1688 kblock=32 || Integer matrix multiply 8816 kblock=64 + // ^1 ^3 ^1 ^3 + // Matrix multiply 1688 kblock=64 + // ^1 ^3 ^1 ^7 ^1 ^3 ^1 ^7 + + // Matrix multiply 16816 kblock=32 | 1688.TF32 kblock=16 || Integer matrix multiply 16832 kblock=64 + // ^2 ^2 + // Matrix multiply 16816 kblock=64 | 1688.TF32 kblock=32 || Integer matrix multiply 16832 kblock=128 + // ^2 ^6 ^2 ^6 + + if ((Policy::kGroupsPerTile / kPartitionsK) > 1) { + int mask = ((Policy::kGroupsPerTile / kPartitionsK) == 8) + ? 3 + : (((Policy::kGroupsPerTile / kPartitionsK) == 4) ? 1 : 0); + + if (((k_group_idx_ & mask) % 2) == 0) + byte_offset_ ^= 1 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + else if ((k_group_idx_ & mask) == 1) + byte_offset_ ^= 3 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + else if ((k_group_idx_ & mask) == 3) + byte_offset_ ^= 7 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + + k_group_idx_++; + + if (k_group_idx_ == (Policy::kGroupsPerTile / kPartitionsK)) { + k_group_idx_ = 0; + add_tile_offset({Policy::kGroupsPerTile, 0}); + } + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator--() { assert(0); } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { load_with_byte_offset(frag, 0); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + Array *fetch_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsmIterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsmIterations::kContiguous; ++c) { + int access_idx = c + s * Policy::LdsmIterations::kContiguous; + + AccessType const *source_ptr = + pointer_ + Policy::LdsmShape::kContiguous * c + + Policy::kLdsmOpInner / Layout::kFactor * + Policy::LdsmShape::kStrided * s * stride_; + + char const *source_byte_ptr = + reinterpret_cast(source_ptr) + byte_offset + + byte_offset_; + + mctlass::arch::ldsm( + fetch_ptr[access_idx], source_byte_ptr); + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = tile_offset.contiguous() * + InstructionShape::kContiguous / + Layout::kElementsPerAccess + + tile_offset.strided() * Shape::kStrided * stride_; + + byte_offset += sizeof_bits::value * pointer_offset / 8; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + k_group_idx_ = k_group % (Policy::kGroupsPerTile / kPartitionsK); + } +}; + +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Element number when the layout crosses (in units of elements) + int Crosswise, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, mctlass::tfloat32_t, + mctlass::layout::TensorOpMultiplicandCrosswise::value, + Crosswise>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for " + "A or B operands to warp-level Mma."); + + /// Element type + using Element = mctlass::tfloat32_t; + + /// Element number when the layout crosses + static int const kCrosswise = Crosswise; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCrosswise< + sizeof_bits::value, kCrosswise>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + // Determine number of elements along outer dimension per individual LDSM op + static int const kLdsmOpOuter = Layout::kElementsPerAccess; + static int const kLdsmOpInner = 8; + + static_assert(!(Shape::kContiguous % kLdsmOpOuter), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsmOpInner), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + /// Shape of one individual LDSM instruction + static int const LdsmShapeContiguous = + InstructionShape::kContiguous / kLdsmOpOuter; + static int const LdsmShapeStrided = + ((4 / LdsmShapeContiguous * kLdsmOpInner) > Shape::kStrided) + ? (Shape::kStrided / kLdsmOpInner) + : (4 / LdsmShapeContiguous); + using LdsmShape = + layout::PitchLinearShape; + + /// Number and arrangement of LDSM instructions + using LdsmIterations = + layout::PitchLinearShape<1, Shape::kStrided / kLdsmOpInner / + LdsmShape::kStrided>; + + /// + static int const kGroupsPerTile = Layout::TileShape::kContiguous / + Layout::kFactor / LdsmShape::kContiguous; + }; + + private: + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = Array; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + + private: + + /// Total number of sections. The memory is divided into stages. One stage + /// can store one tile. Stage is divided into sections. Interleaved layout + /// can have multiple sections in a stage. The rest layout only has one section + /// in a stage. + int sections_; + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + /// Internal counter used to determine when to increment byte offset and when + /// to XOR it + int k_group_idx_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() + : pointer_(nullptr), + sections_(0), + stride_(0), + byte_offset_(0), + k_group_idx_(0) {} + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : pointer_(reinterpret_cast(ref.data())), + sections_(ref.stride(0) / kCrosswise), + // stride_ = kCrosswise x sections_ x kFactor + stride_(ref.stride(0) * Layout::kFactor / Layout::kElementsPerAccess), + byte_offset_(0), + k_group_idx_(0) { + // Warp level iterator at most use double buffer to hide latency. If there + // are more than 2 sections, every stage should have more than 1 section. + + // Turing silicon requires all 32 threads in a warp provide valid addresses + // even for LDSM.1 and LDSM.2 +//#if defined(__MACA_ARCH__) +#if 0 + lane_id = lane_id % (Policy::LdsmShape::kCount * Policy::kLdsmOpInner); +#endif + + int quad_quad = (lane_id >> 4); + int quad_pair = (lane_id >> 3); + int lane_in_pair = (lane_id & 1); + int lane_in_quad = (lane_id & 3); + int lane_in_quad_pair = (lane_id & 7); + int lane_in_quad_quad = (lane_id & 15); + + int partition_contiguous_idx = -1; + int access_contiguous_idx = -1; + int access_strided_idx = -1; + + if (Layout::kFactor == 4) { + // Super Integer matrix multiply Interleaved-32 + + int factor_in_partition = + (Layout::PartitionShape::kContiguous * Layout::kFactor / + Layout::TileShape::kContiguous); + + if (Policy::LdsmShape::kStrided == Policy::LdsmShape::kCount) { + // Integer matrix multiply 8816 A/B + partition_contiguous_idx = lane_in_quad / factor_in_partition; + access_contiguous_idx = ((lane_in_pair * factor_in_partition) ^ + (lane_in_quad_quad / Layout::kFactor)); + access_strided_idx = lane_id / Layout::kFactor; + } + else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kA) { + // Integer matrix multiply 16832 A + partition_contiguous_idx = lane_in_quad / factor_in_partition; + access_strided_idx = lane_in_quad_quad / Layout::kFactor; + access_contiguous_idx = + ((lane_in_pair * factor_in_partition + quad_quad) ^ + access_strided_idx); + } + else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kB) { + // Integer matrix multiply 16832 B + partition_contiguous_idx = lane_in_quad / factor_in_partition; + access_strided_idx = lane_in_quad_pair / Layout::kFactor + quad_quad * 2; + access_contiguous_idx = + ((lane_in_pair * factor_in_partition + ((lane_id & 8) >> 3)) ^ + access_strided_idx); + } + } else if (Layout::kFactor == 2) { + // Super Matrix multiply kBlock = 32 + if (Policy::LdsmShape::kStrided == Policy::LdsmShape::kCount) { + // Matrix multiply 1688 A/B + // (Q stands for 1 8x128bit block). + // Q0 + // Q1 + // Q2 + // Q3 + // Four blocks are next to each other in the strided dimension. + partition_contiguous_idx = (lane_id % Layout::kFactor); + access_contiguous_idx = (lane_in_quad_pair / Layout::kFactor); + access_strided_idx = lane_id / Layout::kFactor; + } + else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kA) { + // Matrix multiply 16816|1688.TF32 A + // Q0 Q2 + // Q1 Q3 + partition_contiguous_idx = (lane_id % Layout::kFactor); + access_contiguous_idx = + (quad_quad ^ (lane_in_quad_pair / Layout::kFactor)); + access_strided_idx = (lane_in_quad_quad / Layout::kFactor); + } else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kB) { + // Matrix multiply 16816|1688.TF32 B + // Q0 Q1 + // Q2 Q3 + partition_contiguous_idx = (lane_id % Layout::kFactor); + access_contiguous_idx = + ((quad_pair & 1) ^ (lane_in_quad_pair / Layout::kFactor)); + access_strided_idx = + (lane_in_quad_pair + (lane_id >> 4 << 3)) / Layout::kFactor; + } + else if (Policy::LdsmShape::kContiguous == Policy::LdsmShape::kCount) { + // Matrix multiply 16832.SP B + // Q0 Q1 Q2 Q3 + partition_contiguous_idx = (lane_id % Layout::kFactor); + access_contiguous_idx = + (quad_pair ^ (lane_in_quad_pair / Layout::kFactor)); + access_strided_idx = lane_in_quad_pair / Layout::kFactor; + } + } else if (Layout::kFactor == 1) { + // Super Matrix multiply kBlock = 64 + if (Policy::LdsmShape::kStrided == Policy::LdsmShape::kCount) { + // Q0 + // Q1 + // Q2 + // Q3 + partition_contiguous_idx = (lane_in_quad_pair >> 2); + access_contiguous_idx = lane_in_quad; + access_strided_idx = lane_id; + } + else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kA) { + // Matrix multiply 16816|1688.TF32 A + // Q0 Q2 + // Q1 Q3 + partition_contiguous_idx = (lane_in_quad_pair >> 2); + access_contiguous_idx = (quad_quad ^ lane_in_quad); + access_strided_idx = lane_in_quad_quad; + } else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kB) { + // Matrix multiply 16816|1688.TF32 B + // Q0 Q1 + // Q2 Q3 + partition_contiguous_idx = (lane_in_quad_pair >> 2); + access_contiguous_idx = ((quad_pair & 1) ^ lane_in_quad); + access_strided_idx = lane_in_quad_pair + (lane_id >> 4 << 3); + } + else if (Policy::LdsmShape::kContiguous == Policy::LdsmShape::kCount) { + // Matrix multiply 16832.SP B + // Q0 Q1 Q2 Q3 + partition_contiguous_idx = (lane_in_quad_pair >> 2); + access_contiguous_idx = (quad_pair ^ lane_in_quad); + access_strided_idx = lane_in_quad_pair; + } + } + + int access_contiguous = partition_contiguous_idx * Layout::PartitionShape::kContiguous + + access_contiguous_idx; + + int access_strided = access_strided_idx; + int byte_offset_o = (access_contiguous + access_strided * stride_) * + sizeof_bits::value * Layout::kElementsPerAccess / 8; + + int row, col; + if (kOperand == Operand::kA) { + if(InstructionShape::kContiguous == 8) { + row = lane_id & 0xf; + col = ((lane_id >> 4) << 1) ^ 0x7; + } + else { + row = lane_id & 0xf; + col = (lane_id >> 4) ^ 0x3; + } + } + else { + if(InstructionShape::kContiguous == 8) { + row = lane_id & 0x7; + col = ((lane_id >> 4) << 1) ^ 0x7; + } + else { + row = lane_id & 0x7; + col = (lane_id >> 4) ^ 0x3; + } + } + + TensorCoord coord = make_Coord(col, row); + Layout lo(ref.stride(0)); + byte_offset_ = lo(coord) * (sizeof_bits::value / 8); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + byte_offset_ += offset * sizeof_bits::value / 8; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + + byte_offset_ ^= k_groups_delta * sizeof_bits::value * + Layout::kElementsPerAccess * + Policy::LdsmShape::kContiguous / 8; + pointer_ += + tile_offset.strided() * stride_ * Shape::kStrided / Layout::kFactor + + whole_tiles * stride_ / sections_; + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative( + TensorCoord const &tile_offset) { + + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + if (k_groups_delta < 0) { + whole_tiles -= 1; + k_groups_delta += Policy::kGroupsPerTile; + } + + if ((Policy::kGroupsPerTile / kPartitionsK) >= 2) { + byte_offset_ ^= (k_groups_delta & 1) * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + if ((Policy::kGroupsPerTile / kPartitionsK) >= 4) { + byte_offset_ ^= ((k_groups_delta + (k_group_idx_ & 1)) & 2) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + if ((Policy::kGroupsPerTile / kPartitionsK) == 8) { + byte_offset_ ^= ((k_groups_delta + (k_group_idx_ & 3)) & 4) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + + k_group_idx_ += k_groups_delta; + whole_tiles += k_group_idx_ / (Policy::kGroupsPerTile / kPartitionsK); + k_group_idx_ = k_group_idx_ % (Policy::kGroupsPerTile / kPartitionsK); + + pointer_ += + tile_offset.strided() * stride_ * Shape::kStrided / Layout::kFactor + + whole_tiles * stride_ / sections_; + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator++() { + + if ((Policy::kGroupsPerTile / kPartitionsK) > 1) { + int mask = ((Policy::kGroupsPerTile / kPartitionsK) == 8) + ? 3 + : (((Policy::kGroupsPerTile / kPartitionsK) == 4) ? 1 : 0); + + if (((k_group_idx_ & mask) % 2) == 0) + byte_offset_ ^= 1 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + else if ((k_group_idx_ & mask) == 1) + byte_offset_ ^= 3 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + else if ((k_group_idx_ & mask) == 3) + byte_offset_ ^= 7 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + + k_group_idx_++; + + if (k_group_idx_ == (Policy::kGroupsPerTile / kPartitionsK)) { + k_group_idx_ = 0; + add_tile_offset({Policy::kGroupsPerTile, 0}); + } + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator--() { assert(0); } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_byte_offset(frag, 0); + + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + Array *fetch_ptr = + reinterpret_cast *>(&frag); + int ldm = stride_ * Layout::kElementsPerAccess / Layout::kFactor; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsmIterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsmIterations::kContiguous; ++c) { + int access_idx = c + s * Policy::LdsmIterations::kContiguous; + + AccessType const *source_ptr = + pointer_ + Policy::LdsmShape::kContiguous * c + + Policy::kLdsmOpInner / Layout::kFactor * + Policy::LdsmShape::kStrided * s * stride_; + + char const *source_byte_ptr = + reinterpret_cast(source_ptr) + byte_offset + + byte_offset_; + if(kOperand == Operand::kA) { + if(InstructionShape::kContiguous == 8) { + mctlass::arch::ldsmAtf32( + fetch_ptr[access_idx], source_byte_ptr); + } + else { + int const *p = reinterpret_cast(source_byte_ptr); + float const *pf = reinterpret_cast(source_byte_ptr); + if(4 == Policy::LdsmShape::kCount) { + int x, y, z, w; + x = p[0]; + z = p[ldm * 16]; + reinterpret_cast(fetch_ptr[access_idx]) = make_int4(x, y, z, w); + } + else { + int x, y; + x = p[0]; + reinterpret_cast(fetch_ptr[access_idx]) = make_int2(x, y); + } + } + } + else { + if(InstructionShape::kContiguous == 8) { + mctlass::arch::ldsmBtf32( + fetch_ptr[access_idx], source_byte_ptr, ldm); + } + else { + int const *p = reinterpret_cast(source_byte_ptr); + float const *fp = reinterpret_cast(source_byte_ptr); + int x, y, z, w; + if(4 == Policy::LdsmShape::kCount) { + x = p[0]; + y = p[8 * ldm]; + z = p[16 * ldm]; + w = p[24 * ldm]; + reinterpret_cast(fetch_ptr[access_idx]) = make_int4(x, y, z, w); + } + else { + x = p[0]; + y = p[8 * ldm]; + reinterpret_cast(fetch_ptr[access_idx]) = make_int2(x, y); + } + } + } + } + } + __syncthreads(); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = tile_offset.contiguous() * + InstructionShape::kContiguous / + Layout::kElementsPerAccess + + tile_offset.strided() * Shape::kStrided * stride_; + + byte_offset += sizeof_bits::value * pointer_offset / 8; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + k_group_idx_ = k_group % (Policy::kGroupsPerTile / kPartitionsK); + } +}; + +/// +/// This special template for float is completely same with mctlass:tfloat32_t, +/// It's a very bad implement. We need to refactor them one day. +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Element number when the layout crosses (in units of elements) + int Crosswise, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, float, + mctlass::layout::TensorOpMultiplicandCrosswise::value, + Crosswise>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for " + "A or B operands to warp-level Mma."); + + /// Element type + using Element =float; + + /// Element number when the layout crosses + static int const kCrosswise = Crosswise; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCrosswise< + sizeof_bits::value, kCrosswise>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + // Determine number of elements along outer dimension per individual LDSM op + static int const kLdsmOpOuter = Layout::kElementsPerAccess; + static int const kLdsmOpInner = 8; + + static_assert(!(Shape::kContiguous % kLdsmOpOuter), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsmOpInner), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + /// Shape of one individual LDSM instruction + static int const LdsmShapeContiguous = + InstructionShape::kContiguous / kLdsmOpOuter; + static int const LdsmShapeStrided = + ((4 / LdsmShapeContiguous * kLdsmOpInner) > Shape::kStrided) + ? (Shape::kStrided / kLdsmOpInner) + : (4 / LdsmShapeContiguous); + using LdsmShape = + layout::PitchLinearShape; + + /// Number and arrangement of LDSM instructions + using LdsmIterations = + layout::PitchLinearShape<1, Shape::kStrided / kLdsmOpInner / + LdsmShape::kStrided>; + + /// + static int const kGroupsPerTile = Layout::TileShape::kContiguous / + Layout::kFactor / LdsmShape::kContiguous; + }; + + private: + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = Array; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + + private: + + /// Total number of sections. The memory is divided into stages. One stage + /// can store one tile. Stage is divided into sections. Interleaved layout + /// can have multiple sections in a stage. The rest layout only has one section + /// in a stage. + int sections_; + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + /// Internal counter used to determine when to increment byte offset and when + /// to XOR it + int k_group_idx_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() + : pointer_(nullptr), + sections_(0), + stride_(0), + byte_offset_(0), + k_group_idx_(0) {} + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : pointer_(reinterpret_cast(ref.data())), + sections_(ref.stride(0) / kCrosswise), + // stride_ = kCrosswise x sections_ x kFactor + stride_(ref.stride(0) * Layout::kFactor / Layout::kElementsPerAccess), + byte_offset_(0), + k_group_idx_(0) { + // Warp level iterator at most use double buffer to hide latency. If there + // are more than 2 sections, every stage should have more than 1 section. + + // Turing silicon requires all 32 threads in a warp provide valid addresses + // even for LDSM.1 and LDSM.2 +//#if defined(__MACA_ARCH__) +#if 0 + lane_id = lane_id % (Policy::LdsmShape::kCount * Policy::kLdsmOpInner); +#endif + + int row, col; + if (kOperand == Operand::kA) { + row = lane_id & 0xf; + col = ((lane_id >> 4) << 1) ^ 0x7; + } else { + row = lane_id & 0x7; + col = ((lane_id >> 4) << 1) ^ 0x7; + } + + TensorCoord coord = make_Coord(col, row); + Layout lo(ref.stride(0)); + byte_offset_ = lo(coord) * (sizeof_bits::value / 8); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + byte_offset_ += offset * sizeof_bits::value / 8; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + + byte_offset_ ^= k_groups_delta * sizeof_bits::value * + Layout::kElementsPerAccess * + Policy::LdsmShape::kContiguous / 8; + pointer_ += + tile_offset.strided() * stride_ * Shape::kStrided / Layout::kFactor + + whole_tiles * stride_ / sections_; + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative( + TensorCoord const &tile_offset) { + + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + if (k_groups_delta < 0) { + whole_tiles -= 1; + k_groups_delta += Policy::kGroupsPerTile; + } + + if ((Policy::kGroupsPerTile / kPartitionsK) >= 2) { + byte_offset_ ^= (k_groups_delta & 1) * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + if ((Policy::kGroupsPerTile / kPartitionsK) >= 4) { + byte_offset_ ^= ((k_groups_delta + (k_group_idx_ & 1)) & 2) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + if ((Policy::kGroupsPerTile / kPartitionsK) == 8) { + byte_offset_ ^= ((k_groups_delta + (k_group_idx_ & 3)) & 4) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + + k_group_idx_ += k_groups_delta; + whole_tiles += k_group_idx_ / (Policy::kGroupsPerTile / kPartitionsK); + k_group_idx_ = k_group_idx_ % (Policy::kGroupsPerTile / kPartitionsK); + + pointer_ += + tile_offset.strided() * stride_ * Shape::kStrided / Layout::kFactor + + whole_tiles * stride_ / sections_; + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator++() { + + if ((Policy::kGroupsPerTile / kPartitionsK) > 1) { + int mask = ((Policy::kGroupsPerTile / kPartitionsK) == 8) + ? 3 + : (((Policy::kGroupsPerTile / kPartitionsK) == 4) ? 1 : 0); + + if (((k_group_idx_ & mask) % 2) == 0) + byte_offset_ ^= 1 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + else if ((k_group_idx_ & mask) == 1) + byte_offset_ ^= 3 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + else if ((k_group_idx_ & mask) == 3) + byte_offset_ ^= 7 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + + k_group_idx_++; + + if (k_group_idx_ == (Policy::kGroupsPerTile / kPartitionsK)) { + k_group_idx_ = 0; + add_tile_offset({Policy::kGroupsPerTile, 0}); + } + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator--() { assert(0); } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_byte_offset(frag, 0); + + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + Array *fetch_ptr = + reinterpret_cast *>(&frag); + int ldm = stride_ * Layout::kElementsPerAccess / Layout::kFactor; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsmIterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsmIterations::kContiguous; ++c) { + int access_idx = c + s * Policy::LdsmIterations::kContiguous; + + AccessType const *source_ptr = + pointer_ + Policy::LdsmShape::kContiguous * c + + Policy::kLdsmOpInner / Layout::kFactor * + Policy::LdsmShape::kStrided * s * stride_; + + char const *source_byte_ptr = + reinterpret_cast(source_ptr) + byte_offset + + byte_offset_; + if(kOperand == Operand::kA) { + mctlass::arch::ldsmAtf32( + fetch_ptr[access_idx], source_byte_ptr); + } else { + mctlass::arch::ldsmBtf32( + fetch_ptr[access_idx], source_byte_ptr, ldm); + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = tile_offset.contiguous() * + InstructionShape::kContiguous / + Layout::kElementsPerAccess + + tile_offset.strided() * Shape::kStrided * stride_; + + byte_offset += sizeof_bits::value * pointer_offset / 8; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + k_group_idx_ = k_group % (Policy::kGroupsPerTile / kPartitionsK); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Element number when the layout crosses (in units of elements) + int Crosswise, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, signed char, + mctlass::layout::TensorOpMultiplicandCrosswise::value, + Crosswise>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for " + "A or B operands to warp-level Mma."); + + /// Element type + using Element = signed char; + + /// Element number when the layout crosses + static int const kCrosswise = Crosswise; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCrosswise< + sizeof_bits::value, kCrosswise>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + // Determine number of elements along outer dimension per individual LDSM op + static int const kLdsmOpOuter = Layout::kElementsPerAccess; + static int const kLdsmOpInner = 8; + + static_assert(!(Shape::kContiguous % kLdsmOpOuter), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsmOpInner), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + // Check supported InstructionShape + static_assert( + InstructionShape::kContiguous == 16|| + InstructionShape::kContiguous == 32, + "Supported list of mma operator shape for int8_t multiplicands are: 16x8x16, 16x8x32"); + + /// Shape of one individual LDSM instruction + static int const LdsmShapeContiguous = + InstructionShape::kContiguous / kLdsmOpOuter; + static int const LdsmShapeStrided = + ((4 / LdsmShapeContiguous * kLdsmOpInner) > Shape::kStrided) + ? (Shape::kStrided / kLdsmOpInner) + : (4 / LdsmShapeContiguous); + + using LdsmShape = + layout::PitchLinearShape; + + /// Number and arrangement of LDSM instructions + using LdsmIterations = + layout::PitchLinearShape<1, Shape::kStrided / kLdsmOpInner / + LdsmShape::kStrided>; + + /// + static int const kGroupsPerTile = Layout::TileShape::kContiguous / + Layout::kFactor / LdsmShape::kContiguous; + }; + + private: + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = Array; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + + private: + + /// Total number of sections. The memory is divided into stages. One stage + /// can store one tile. Stage is divided into sections. Interleaved layout + /// can have multiple sections in a stage. The rest layout only has one section + /// in a stage. + int sections_; + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + int distance_matrix[3]; + + /// Internal counter used to determine when to increment byte offset and when + /// to XOR it + int k_group_idx_; + Layout layout_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() + : pointer_(nullptr), + sections_(0), + stride_(0), + byte_offset_(0), + k_group_idx_(0) {} + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : pointer_(reinterpret_cast(ref.data())), + sections_(ref.stride(0) / kCrosswise), + // stride_ = kCrosswise x sections_ x kFactor + stride_(ref.stride(0) * Layout::kFactor / Layout::kElementsPerAccess), + byte_offset_(0), + k_group_idx_(0) { + // Warp level iterator at most use double buffer to hide latency. If there + // are more than 2 sections, every stage should have more than 1 section. + + // Turing silicon requires all 32 threads in a warp provide valid addresses + // even for LDSM.1 and LDSM.2 +//#if defined(__MACA_ARCH__) +#if 0 + lane_id = lane_id % (Policy::LdsmShape::kCount * Policy::kLdsmOpInner); +#endif + + int quad_quad = (lane_id >> 4); + int quad_pair = (lane_id >> 3); + int lane_in_pair = (lane_id & 1); + int lane_in_quad = (lane_id & 3); + int lane_in_quad_pair = (lane_id & 7); + int lane_in_quad_quad = (lane_id & 15); + + int partition_contiguous_idx = -1; + int access_contiguous_idx = -1; + int access_strided_idx = -1; + + if (Layout::kFactor == 4) { + // Super Integer matrix multiply Interleaved-32 + + int factor_in_partition = + (Layout::PartitionShape::kContiguous * Layout::kFactor / + Layout::TileShape::kContiguous); + + if (Policy::LdsmShape::kStrided == Policy::LdsmShape::kCount) { + // Integer matrix multiply 8816 A/B + partition_contiguous_idx = lane_in_quad / factor_in_partition; + access_contiguous_idx = ((lane_in_pair * factor_in_partition) ^ + (lane_in_quad_quad / Layout::kFactor)); + access_strided_idx = lane_id / Layout::kFactor; + } + else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kA) { + // Integer matrix multiply 16832 A + partition_contiguous_idx = lane_in_quad / factor_in_partition; + access_strided_idx = lane_in_quad_quad / Layout::kFactor; + access_contiguous_idx = + ((lane_in_pair * factor_in_partition + quad_quad) ^ + access_strided_idx); + } + else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kB) { + // Integer matrix multiply 16832 B + partition_contiguous_idx = lane_in_quad / factor_in_partition; + access_strided_idx = lane_in_quad_pair / Layout::kFactor + quad_quad * 2; + access_contiguous_idx = + ((lane_in_pair * factor_in_partition + ((lane_id & 8) >> 3)) ^ + access_strided_idx); + } + } else if (Layout::kFactor == 2) { + // Super Matrix multiply kBlock = 32 + if (Policy::LdsmShape::kStrided == Policy::LdsmShape::kCount) { + // Matrix multiply 1688 A/B + // (Q stands for 1 8x128bit block). + // Q0 + // Q1 + // Q2 + // Q3 + // Four blocks are next to each other in the strided dimension. + partition_contiguous_idx = (lane_id % Layout::kFactor); + access_contiguous_idx = (lane_in_quad_pair / Layout::kFactor); + access_strided_idx = lane_id / Layout::kFactor; + } + else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kA) { + // Matrix multiply 16816|1688.TF32 A + // Q0 Q2 + // Q1 Q3 + partition_contiguous_idx = (lane_id % Layout::kFactor); + access_contiguous_idx = + (quad_quad ^ (lane_in_quad_pair / Layout::kFactor)); + access_strided_idx = (lane_in_quad_quad / Layout::kFactor); + } else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kB) { + // Matrix multiply 16816|1688.TF32 B + // Q0 Q1 + // Q2 Q3 + partition_contiguous_idx = (lane_id % Layout::kFactor); + access_contiguous_idx = + ((quad_pair & 1) ^ (lane_in_quad_pair / Layout::kFactor)); + access_strided_idx = + (lane_in_quad_pair + (lane_id >> 4 << 3)) / Layout::kFactor; + } + else if (Policy::LdsmShape::kContiguous == Policy::LdsmShape::kCount) { + // Matrix multiply 16832.SP B + // Q0 Q1 Q2 Q3 + partition_contiguous_idx = (lane_id % Layout::kFactor); + access_contiguous_idx = + (quad_pair ^ (lane_in_quad_pair / Layout::kFactor)); + access_strided_idx = lane_in_quad_pair / Layout::kFactor; + } + } else if (Layout::kFactor == 1) { + // Super Matrix multiply kBlock = 64 + if (Policy::LdsmShape::kStrided == Policy::LdsmShape::kCount) { + // Q0 + // Q1 + // Q2 + // Q3 + partition_contiguous_idx = (lane_in_quad_pair >> 2); + access_contiguous_idx = lane_in_quad; + access_strided_idx = lane_id; + } + else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kA) { + // Matrix multiply 16816|1688.TF32 A + // Q0 Q2 + // Q1 Q3 + partition_contiguous_idx = (lane_in_quad_pair >> 2); + access_contiguous_idx = (quad_quad ^ lane_in_quad); + access_strided_idx = lane_in_quad_quad; + } else if (Policy::LdsmShape::kStrided == + (Policy::LdsmShape::kCount / 2) && + kOperand == Operand::kB) { + // Matrix multiply 16816|1688.TF32 B + // Q0 Q1 + // Q2 Q3 + partition_contiguous_idx = (lane_in_quad_pair >> 2); + access_contiguous_idx = ((quad_pair & 1) ^ lane_in_quad); + access_strided_idx = lane_in_quad_pair + (lane_id >> 4 << 3); + } + else if (Policy::LdsmShape::kContiguous == Policy::LdsmShape::kCount) { + // Matrix multiply 16832.SP B + // Q0 Q1 Q2 Q3 + partition_contiguous_idx = (lane_in_quad_pair >> 2); + access_contiguous_idx = (quad_pair ^ lane_in_quad); + access_strided_idx = lane_in_quad_pair; + } + } + + int access_contiguous = + partition_contiguous_idx * Layout::PartitionShape::kContiguous + + access_contiguous_idx; + + int access_strided = access_strided_idx; + + int byte_offset_o = (access_contiguous + access_strided * stride_) * + sizeof_bits::value * Layout::kElementsPerAccess / 8; + + // For InstructionShape 16x8x16 and 16x8x32 + if (InstructionShape::kContiguous == 16) { + int x, y, x1, y1, x2, y2, x3, y3; + x = (((lane_id >> 4) << 2) ^ 0xf) - 3; + //for matrix A InstructionShape::kStrided=16, for matrix B InstructionShape::kStrided=8 + y = lane_id & (InstructionShape::kStrided - 1); + + if (kOperand == Operand::kA) { + x1 = x + 16; + y1 = y; + x2 = x; + y2 = y + 16; + x3 = x + 16; + y3 = y + 16; + } + else { + x1 = x; + y1 = y + 8; + x2 = x; + y2 = y + 16; + x3 = x; + y3 = y + 24; + } + + TensorCoord coord0 = make_Coord(x, y); + Layout lo = ref.layout(); + byte_offset_ = lo(coord0) * (sizeof_bits::value / 8); + + TensorCoord coord1 = make_Coord(x1, y1); + Index byte_offset_1 = lo(coord1) * (sizeof_bits::value / 8); + TensorCoord coord2 = make_Coord(x2, y2); + Index byte_offset_2 = lo(coord2) * (sizeof_bits::value / 8); + TensorCoord coord3 = make_Coord(x3, y3); + Index byte_offset_3 = lo(coord3) * (sizeof_bits::value / 8); + + distance_matrix[0] = (byte_offset_1 - byte_offset_) / 4; + distance_matrix[1] = (byte_offset_2 - byte_offset_) / 4; + distance_matrix[2] = (byte_offset_3 - byte_offset_) / 4; + } + else if (InstructionShape::kContiguous == 32) { + layout_ = ref.layout(); + byte_offset_ = 0; + } + + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + byte_offset_ += offset * sizeof_bits::value / 8; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + + byte_offset_ ^= k_groups_delta * sizeof_bits::value * + Layout::kElementsPerAccess * + Policy::LdsmShape::kContiguous / 8; + pointer_ += + tile_offset.strided() * stride_ * Shape::kStrided / Layout::kFactor + + whole_tiles * stride_ / sections_; + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative( + TensorCoord const &tile_offset) { + + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + if (k_groups_delta < 0) { + whole_tiles -= 1; + k_groups_delta += Policy::kGroupsPerTile; + } + + if ((Policy::kGroupsPerTile / kPartitionsK) >= 2) { + byte_offset_ ^= (k_groups_delta & 1) * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + if ((Policy::kGroupsPerTile / kPartitionsK) >= 4) { + byte_offset_ ^= ((k_groups_delta + (k_group_idx_ & 1)) & 2) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + if ((Policy::kGroupsPerTile / kPartitionsK) == 8) { + byte_offset_ ^= ((k_groups_delta + (k_group_idx_ & 3)) & 4) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + + k_group_idx_ += k_groups_delta; + whole_tiles += k_group_idx_ / (Policy::kGroupsPerTile / kPartitionsK); + k_group_idx_ = k_group_idx_ % (Policy::kGroupsPerTile / kPartitionsK); + + pointer_ += + tile_offset.strided() * stride_ * Shape::kStrided / Layout::kFactor + + whole_tiles * stride_ / sections_; + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator++() { + + // Integer matrix multiply 16832 Interleaved-32 + // NONE + // Integer matrix multiply 16816 Interleaved-32 || Integer matrix multiply 16816 kblock=32 + + // Integer matrix multiply 8816 Interleaved-32 + // ^1 ^1 + // Matrix multiply 1684.TF32 kblock=16 || Integer matrix multiply 16816 kblock=64 + // Matrix multiply 1688 kblock=32 || Integer matrix multiply 8816 kblock=64 + // ^1 ^3 ^1 ^3 + // Matrix multiply 1688 kblock=64 + // ^1 ^3 ^1 ^7 ^1 ^3 ^1 ^7 + + // Matrix multiply 16816 kblock=32 | 1688.TF32 kblock=16 || Integer matrix multiply 16832 kblock=64 + // ^2 ^2 + // Matrix multiply 16816 kblock=64 | 1688.TF32 kblock=32 || Integer matrix multiply 16832 kblock=128 + // ^2 ^6 ^2 ^6 + if (InstructionShape::kContiguous == 32) { + byte_offset_ += InstructionShape::kContiguous; + } + else { + if ((Policy::kGroupsPerTile / kPartitionsK) > 1) { + int mask = ((Policy::kGroupsPerTile / kPartitionsK) == 8) + ? 3 + : (((Policy::kGroupsPerTile / kPartitionsK) == 4) ? 1 : 0); + + if (((k_group_idx_ & mask) % 2) == 0) + byte_offset_ ^= 1 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + else if ((k_group_idx_ & mask) == 1) + byte_offset_ ^= 3 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + else if ((k_group_idx_ & mask) == 3) + byte_offset_ ^= 7 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + + k_group_idx_++; + + if (k_group_idx_ == (Policy::kGroupsPerTile / kPartitionsK)) { + k_group_idx_ = 0; + add_tile_offset({Policy::kGroupsPerTile, 0}); + } + } + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator--() { assert(0); } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { load_with_byte_offset(frag, 0); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + if (InstructionShape::kContiguous == 32) { + + Array *fetch_element_ptr = + reinterpret_cast *>(&frag); + const int lane_id = __lane_id(); + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsmIterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsmIterations::kContiguous; ++c) { + + Element const *ptr = reinterpret_cast(pointer_) + byte_offset; + Index row = (((lane_id >> 4) << 2) ^ 0xf) + byte_offset_; + Index col = (lane_id & (InstructionShape::kStrided - 1)) + (s * 16); + int access_idx = c + s * Policy::LdsmIterations::kContiguous * 4; + + for (int i = 0; i < Policy::LdsmShape::kCount; ++i) { + + Index source_idx = layout_({row - i, col}); + Index source_idx1 = layout_({row - i + 16, col}); + Index source_idx2 = layout_({row - i, col + 8}); + Index source_idx3 = layout_({row - i + 16, col + 8}); + Index dst_idx = Policy::LdsmShape::kCount - i - 1; + + fetch_element_ptr[access_idx][dst_idx] = ptr[source_idx]; + fetch_element_ptr[(access_idx + 1)][dst_idx] = ptr[source_idx1]; + fetch_element_ptr[(access_idx + 2)][dst_idx] = ptr[source_idx2]; + fetch_element_ptr[(access_idx + 3)][dst_idx] = ptr[source_idx3]; + } + } + } + } + else { + + Array *fetch_ptr = + reinterpret_cast *>(&frag); + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsmIterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsmIterations::kContiguous; ++c) { + int access_idx = c + s * Policy::LdsmIterations::kContiguous; + AccessType const *source_ptr = + pointer_ + Policy::LdsmShape::kContiguous * c + + Policy::kLdsmOpInner / Layout::kFactor * + Policy::LdsmShape::kStrided * s * stride_; + + char const *source_byte_ptr = + reinterpret_cast(source_ptr) + byte_offset + + byte_offset_; + + mctlass::arch::ldsmi8( + fetch_ptr[access_idx], source_byte_ptr, distance_matrix); + } + } + + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = tile_offset.contiguous() * + InstructionShape::kContiguous / + Layout::kElementsPerAccess + + tile_offset.strided() * Shape::kStrided * stride_; + + byte_offset += sizeof_bits::value * pointer_offset / 8; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + k_group_idx_ = k_group % (Policy::kGroupsPerTile / kPartitionsK); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Element number when the layout crosses (in units of elements) + int Crosswise, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Crosswise>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator for ColumnMajor Crosswise may " + "only be instantiated for B operand to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// KBlock size + static int const kCrosswise = Crosswise; + + /// Layout of source tile + using Layout = mctlass::layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, kCrosswise>; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCrosswise::value, + kCrosswise>, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + + private: + /// Underlying tile iterator + Base iterator_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() {} + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : iterator_({ref.data(), ref.stride()}, lane_id) {} + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset_negative({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator++() { + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator--() { + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { iterator_.load(frag); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, {tile_offset.contiguous(), tile_offset.strided()}, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Element number when the layout crosses (in units of elements) + int Crosswise, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::RowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Crosswise>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA, + "MmaTensorOpMultiplicandIterator for RowMajor Crosswise may " + "only be instantiated for A operand to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Element number when the layout crosses + static int const kCrosswise = Crosswise; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, kCrosswise>; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCrosswise::value, + kCrosswise>, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + + private: + /// Underlying tile iterator + Base iterator_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() {} + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : iterator_({ref.data(), ref.stride()}, lane_id) {} + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset_negative({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator++() { + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator--() { + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { iterator_.load(frag); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, {tile_offset.strided(), tile_offset.contiguous()}, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Layout of operand in memory + typename Layout_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_> +class MmaTensorOpAccumulatorTileIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It is used to load or store +/// accumulators from memory and is agnostic to layout. It could be faster if it assumed row-major +/// accumulator layout. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept | +/// WriteableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_> +class MmaTensorOpAccumulatorTileIterator< + Shape_, Element_, mctlass::layout::RowMajor, InstructionShape_, OpDelta_> { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + using OpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static bool const kDivisible = + !(Shape::kRow % InstructionShape::kM) && + !(Shape::kColumn % InstructionShape::kN); + + static_assert(platform::is_same::value, + "Layouts must be defined for logical MatrixCoord coordinate space."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + (Shape::kRow + InstructionShape::kM - 1) / InstructionShape::kM, + (Shape::kColumn + InstructionShape::kN - 1) / InstructionShape::kN + >; + }; + +private: + + // Assume accumulator tile is an arrangement of 8-by-8 tiles replicated over the entire + // shape, with each quad mapped to one row and each thread mapped to 1/4 of the elements + // of that row. The accumulators within one row are assumed to be consecutive. + static int const kElementsPerAccess = InstructionShape::kN / 4; + static int const kRowsPerTile = 8; + static int const kAccumulatorRows = InstructionShape::kM / kRowsPerTile; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array< + Element, + Policy::MmaIterations::kCount * InstructionShape::kMN / kThreads>; + +private: + + /// Reference to output tensor + TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator( + TensorRef const &ref, + int lane_id + ): + ref_(ref) { + + // int quad = (lane_id >> 2); + // int lane_in_quad = (lane_id & 3); + // MatrixCoord lane_offset(quad, lane_in_quad * kElementsPerAccess); + int lane_in_quad = lane_id & 0x7; + int quad = (lane_id >> 4) & 0x3; + + MatrixCoord lane_offset(quad, lane_in_quad); + + ref_.add_coord_offset(lane_offset); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + ref_.add_coord_offset(tile_offset * make_Coord(Shape::kRow, Shape::kColumn)); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator++() { + // deliberate no-op + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator--() { + // deliberate no-op + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to load from the tensor + Index pointer_offset) const { ///< loads a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + + frag[mma_accum_start + row * kElementsPerAccess + col] = offset_ref.at({accum_m, accum_n}); + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + Fragment &frag, ///< fragment to load from the tensor + Index byte_offset) const { ///< loads a tile with a linear offset + + load_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset) const { ///< loads a tile with a logical offset in units of whole tiles + + load(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset, ///< loads a tile with a logical offset in units of whole tiles + Index pointer_offset) const { ///< loads a tile with a logical offset AND a pointer offset + + load_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } + + /// Stores a fragment to memory + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index pointer_offset) const { ///< store a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + // int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + // row * kRowsPerTile; + // int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + row * kRowsPerTile + + col * (InstructionShape::kM / kElementsPerAccess); + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn; + int idx = mma_accum_start + row * kElementsPerAccess + col; + + offset_ref.at({accum_m, accum_n}) = frag[idx]; + } + } + } + } + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_byte_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index byte_offset) const { ///< store a tile with a linear offset + + store_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Stores a fragment to memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + Fragment &frag, ///< fragment to store to the tensor + TensorCoord const &tile_offset) const { ///< stores a tile with a logical offset in units of whole tiles + + store(frag, tile_offset, 0); + } + + /// Stores a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + /// fragment to store to the tensor + Fragment const &frag, + /// stores a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// stores a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + store_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } +}; + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_> +class MmaTensorOpAccumulatorTileIterator< + Shape_, int, mctlass::layout::RowMajor, InstructionShape_, OpDelta_> { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = int; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + using OpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static bool const kDivisible = + !(Shape::kRow % InstructionShape::kM) && + !(Shape::kColumn % InstructionShape::kN); + + static_assert(platform::is_same::value, + "Layouts must be defined for logical MatrixCoord coordinate space."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + (Shape::kRow + InstructionShape::kM - 1) / InstructionShape::kM, + (Shape::kColumn + InstructionShape::kN - 1) / InstructionShape::kN + >; + }; + +private: + + // Assume accumulator tile is an arrangement of 8-by-8 tiles replicated over the entire + // shape, with each quad mapped to one row and each thread mapped to 1/4 of the elements + // of that row. The accumulators within one row are assumed to be consecutive. + static int const kElementsPerAccess = InstructionShape::kN / 4; + static int const kRowsPerTile = 8; + static int const kAccumulatorRows = InstructionShape::kM / kRowsPerTile; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array< + Element, + Policy::MmaIterations::kCount * InstructionShape::kMN / kThreads>; + +private: + + /// Reference to output tensor + TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator( + TensorRef const &ref, + int lane_id + ): + ref_(ref) { + + // int quad = (lane_id >> 2); + // int lane_in_quad = (lane_id & 3); + // MatrixCoord lane_offset(quad, lane_in_quad * kElementsPerAccess); + int quad, lane_in_quad; + + quad = ((lane_id >> 4) << 2); + lane_in_quad = lane_id & 0x7; + + MatrixCoord lane_offset(quad, lane_in_quad); + + ref_.add_coord_offset(lane_offset); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + ref_.add_coord_offset(tile_offset * make_Coord(Shape::kRow, Shape::kColumn)); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator++() { + // deliberate no-op + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator--() { + // deliberate no-op + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator-=(TensorCoord const &tile_offset) { + // add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to load from the tensor + Index pointer_offset) const { ///< loads a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + + frag[mma_accum_start + row * kElementsPerAccess + col] = offset_ref.at({accum_m, accum_n}); + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + Fragment &frag, ///< fragment to load from the tensor + Index byte_offset) const { ///< loads a tile with a linear offset + + load_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset) const { ///< loads a tile with a logical offset in units of whole tiles + + load(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset, ///< loads a tile with a logical offset in units of whole tiles + Index pointer_offset) const { ///< loads a tile with a logical offset AND a pointer offset + + load_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } + + /// Stores a fragment to memory + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index pointer_offset) const { ///< store a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + // Policy::MmaIterations::kColumn, Policy::MmaIterations::kRow, kAccumulatorRows, kElementsPerAccess); + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + row * kAccumulatorRows + + col; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn; + int idx = mma_accum_start + row * kElementsPerAccess + col; + offset_ref.at({accum_m, accum_n}) = frag[idx]; + } + } + } + } + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_byte_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index byte_offset) const { ///< store a tile with a linear offset + + store_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Stores a fragment to memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + Fragment &frag, ///< fragment to store to the tensor + TensorCoord const &tile_offset) const { ///< stores a tile with a logical offset in units of whole tiles + + store(frag, tile_offset, 0); + } + + /// Stores a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + /// fragment to store to the tensor + Fragment const &frag, + /// stores a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// stores a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + store_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } +}; + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_> +class MmaTensorOpAccumulatorTileIterator< + Shape_, float, mctlass::layout::RowMajor, InstructionShape_, OpDelta_> { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = float; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + using OpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static bool const kDivisible = + !(Shape::kRow % InstructionShape::kM) && + !(Shape::kColumn % InstructionShape::kN); + + static_assert(platform::is_same::value, + "Layouts must be defined for logical MatrixCoord coordinate space."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + (Shape::kRow + InstructionShape::kM - 1) / InstructionShape::kM, + (Shape::kColumn + InstructionShape::kN - 1) / InstructionShape::kN + >; + }; + +private: + + // Assume accumulator tile is an arrangement of 8-by-8 tiles replicated over the entire + // shape, with each quad mapped to one row and each thread mapped to 1/4 of the elements + // of that row. The accumulators within one row are assumed to be consecutive. + static int const kElementsPerAccess = InstructionShape::kN / 4; + static int const kRowsPerTile = 8; + static int const kAccumulatorRows = InstructionShape::kM / kRowsPerTile; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array< + Element, + Policy::MmaIterations::kCount * InstructionShape::kMN / kThreads>; + +private: + + /// Reference to output tensor + TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator( + TensorRef const &ref, + int lane_id + ): + ref_(ref) { + int quad = (lane_id >> 4) << 2; + int lane_in_quad = lane_id & 0x7; + MatrixCoord lane_offset(quad, lane_in_quad); + ref_.add_coord_offset(lane_offset); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + ref_.add_coord_offset(tile_offset * make_Coord(Shape::kRow, Shape::kColumn)); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator++() { + // deliberate no-op + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator--() { + // deliberate no-op + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator-=(TensorCoord const &tile_offset) { + //add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to load from the tensor + Index pointer_offset) const { ///< loads a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + + frag[mma_accum_start + row * kElementsPerAccess + col] = offset_ref.at({accum_m, accum_n}); + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + Fragment &frag, ///< fragment to load from the tensor + Index byte_offset) const { ///< loads a tile with a linear offset + + load_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset) const { ///< loads a tile with a logical offset in units of whole tiles + + load(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset, ///< loads a tile with a logical offset in units of whole tiles + Index pointer_offset) const { ///< loads a tile with a logical offset AND a pointer offset + + load_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } + + /// Stores a fragment to memory + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index pointer_offset) const { ///< store a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow); + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_m = mma_m * InstructionShape::kM; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn; + int t_idx = row * kElementsPerAccess + col; + int idx = mma_accum_start + t_idx + mma_m * 4; + offset_ref.at({accum_m + t_idx, accum_n}) = frag[idx]; + } + } + } + } + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_byte_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index byte_offset) const { ///< store a tile with a linear offset + + store_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Stores a fragment to memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + Fragment &frag, ///< fragment to store to the tensor + TensorCoord const &tile_offset) const { ///< stores a tile with a logical offset in units of whole tiles + + store(frag, tile_offset, 0); + } + + /// Stores a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + /// fragment to store to the tensor + Fragment const &frag, + /// stores a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// stores a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + store_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It is used to load or store +/// accumulators from memory and is agnostic to layout. +/// +/// This iterator is not tested. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept | +/// WriteableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_> +class MmaTensorOpAccumulatorTileIterator< + Shape_, Element_, mctlass::layout::AffineRankN<2>, InstructionShape_, OpDelta_> { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + using OpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static bool const kDivisible = + !(Shape::kRow % InstructionShape::kM) && + !(Shape::kColumn % InstructionShape::kN); + + static_assert(platform::is_same::value, + "Layouts must be defined for logical MatrixCoord coordinate space."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + (Shape::kRow + InstructionShape::kM - 1) / InstructionShape::kM, + (Shape::kColumn + InstructionShape::kN - 1) / InstructionShape::kN + >; + }; + +private: + + // Assume accumulator tile is an arrangement of 8-by-8 tiles replicated over the entire + // shape, with each quad mapped to one row and each thread mapped to 1/4 of the elements + // of that row. The accumulators within one row are assumed to be consecutive. + static int const kElementsPerAccess = InstructionShape::kN / 4; + static int const kRowsPerTile = 8; + static int const kAccumulatorRows = InstructionShape::kM / kRowsPerTile; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array< + Element, + Policy::MmaIterations::kCount * InstructionShape::kMN / kThreads>; + +private: + + /// Reference to output tensor + TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator( + TensorRef const &ref, + int lane_id + ): + ref_(ref) { + + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + + MatrixCoord lane_offset(quad, lane_in_quad * kElementsPerAccess); + + ref_.add_coord_offset(lane_offset); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + ref_.add_coord_offset(tile_offset * make_Coord(Shape::kRow, Shape::kColumn)); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator++() { + // deliberate no-op + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator--() { + // deliberate no-op + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to load from the tensor + Index pointer_offset) const { ///< loads a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + + frag[mma_accum_start + row * kElementsPerAccess + col] = offset_ref.at({accum_m, accum_n}); + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + Fragment &frag, ///< fragment to load from the tensor + Index byte_offset) const { ///< loads a tile with a linear offset + + load_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset) const { ///< loads a tile with a logical offset in units of whole tiles + + load(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset, ///< loads a tile with a logical offset in units of whole tiles + Index pointer_offset) const { ///< loads a tile with a logical offset AND a pointer offset + + load_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } + + /// Stores a fragment to memory + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index pointer_offset) const { ///< store a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + int idx = mma_accum_start + row * kElementsPerAccess + col; + + offset_ref.at({accum_m, accum_n}) = frag[idx]; + } + } + } + } + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_byte_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index byte_offset) const { ///< store a tile with a linear offset + + store_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Stores a fragment to memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + Fragment &frag, ///< fragment to store to the tensor + TensorCoord const &tile_offset) const { ///< stores a tile with a logical offset in units of whole tiles + + store(frag, tile_offset, 0); + } + + /// Stores a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + /// fragment to store to the tensor + Fragment const &frag, + /// stores a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// stores a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + store_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It is used to load or store +/// accumulators from memory and is agnostic to layout. It could be faster if it assumed row-major +/// accumulator layout. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept | +/// WriteableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_> +class MmaTensorOpAccumulatorTileIterator { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::ColumnMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + using OpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static bool const kDivisible = + !(Shape::kRow % InstructionShape::kM) && + !(Shape::kColumn % InstructionShape::kN); + + static_assert(platform::is_same::value, + "Layouts must be defined for logical MatrixCoord coordinate space."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + (Shape::kRow + InstructionShape::kM - 1) / InstructionShape::kM, + (Shape::kColumn + InstructionShape::kN - 1) / InstructionShape::kN + >; + }; + +private: + + // Assume accumulator tile is an arrangement of 8-by-8 tiles replicated over the entire + // shape, with each quad mapped to one row and each thread mapped to 1/4 of the elements + // of that row. The accumulators within one row are assumed to be consecutive. + static int const kElementsPerAccess = InstructionShape::kN / 4; + static int const kRowsPerTile = 8; + static int const kAccumulatorRows = InstructionShape::kM / kRowsPerTile; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + +private: + + /// Reference to output tensor + TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator( + TensorRef const &ref, + int lane_id + ): + ref_(ref) { + + if (platform::is_same::value == true) { + if (InstructionShape::kM == 16 && InstructionShape::kN == 8 && (InstructionShape::kK == 16 || InstructionShape::kK == 4 || InstructionShape::kK == 8)) { + + const int quad = (lane_id >> 4) << 2; + const int lane_in_quad = lane_id & 0x7; + MatrixCoord lane_offset(quad, lane_in_quad); + ref_.add_coord_offset(lane_offset); + + } + else { + printf("Not Impl Other InstructionShape for float.\n"); + } + } + else if(platform::is_same::value == true) { + if (InstructionShape::kM == 16 && InstructionShape::kN == 8 && InstructionShape::kK == 32) { + + const int quad = ((lane_id >> 4) << 2); + const int lane_in_quad = lane_id & 0x7; + MatrixCoord lane_offset(quad, lane_in_quad); + ref_.add_coord_offset(lane_offset); + + } + else { + printf("Not Impl Other InstructionShape for int.\n"); + } + } + else { + + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + + MatrixCoord lane_offset(quad, lane_in_quad * kElementsPerAccess); + + ref_.add_coord_offset(lane_offset); + } + + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + ref_.add_coord_offset(tile_offset * make_Coord(Shape::kRow, Shape::kColumn)); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator++() { + // deliberate no-op + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator--() { + // deliberate no-op + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to load from the tensor + Index pointer_offset) const { ///< loads a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + int idx = mma_accum_start + row * kElementsPerAccess + col; + + frag[idx] = offset_ref.at({accum_m, accum_n}); + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + Fragment &frag, ///< fragment to load from the tensor + Index byte_offset) const { ///< loads a tile with a linear offset + + load_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset) const { ///< loads a tile with a logical offset in units of whole tiles + + load(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset, ///< loads a tile with a logical offset in units of whole tiles + Index pointer_offset) const { ///< loads a tile with a logical offset AND a pointer offset + + load_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } + + /// Stores a fragment to memory + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index pointer_offset) const { ///< store a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + if (InstructionShape::kM == 16 && InstructionShape::kN == 8 && InstructionShape::kK == 4) { + mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow); + } + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + if (platform::is_same::value == true) { + if (InstructionShape::kM == 16 && InstructionShape::kN == 8 && (InstructionShape::kK == 16 || InstructionShape::kK == 8)) { + const int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + row * kElementsPerAccess + col; + const int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn; + const int idx = mma_accum_start + row * kElementsPerAccess + col; + offset_ref.at({accum_m, accum_n}) = frag[idx]; + } + else if (InstructionShape::kM == 16 && InstructionShape::kN == 8 && InstructionShape::kK == 4) { + int accum_m = mma_m * InstructionShape::kM; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn; + int t_idx = row * kElementsPerAccess + col; + int idx = mma_accum_start + t_idx + mma_m * 4; + offset_ref.at({accum_m + t_idx, accum_n}) = frag[idx]; + } + else { + printf("Not Impl other InstructionShape with float Now.\n"); + } + } + else if (platform::is_same::value == true) { + if (InstructionShape::kM == 16 && InstructionShape::kN == 8 && InstructionShape::kK == 32) { + const int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + row * kElementsPerAccess + col; + const int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn; + const int idx = mma_accum_start + row * kElementsPerAccess + col; + offset_ref.at({accum_m, accum_n}) = frag[idx]; + } + else { + printf("Not Impl other InstructionShape with int Now.\n"); + } + } + else { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + int idx = mma_accum_start + row * kElementsPerAccess + col; + + offset_ref.at({accum_m, accum_n}) = frag[idx]; + } + } + } + } + } + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_byte_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index byte_offset) const { ///< store a tile with a linear offset + + store_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Stores a fragment to memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + Fragment &frag, ///< fragment to store to the tensor + TensorCoord const &tile_offset) const { ///< stores a tile with a logical offset in units of whole tiles + + store(frag, tile_offset, 0); + } + + /// Stores a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + /// fragment to store to the tensor + Fragment const &frag, + /// stores a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// stores a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + store_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It is used to load or store +/// accumulators from memory and is agnostic to layout. It could be faster if it assumed row-major +/// accumulator layout. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept | +/// WriteableRandomAccessContiguousTileIteratorConcept +/// + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Element typ + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_, + /// Interleaved N + int InterleavedN> +class MmaTensorOpAccumulatorTileIterator< + Shape_, Element_, mctlass::layout::ColumnMajorInterleaved, + InstructionShape_, OpDelta_> { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::ColumnMajorInterleaved; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + using OpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kRow % InstructionShape::kM) && + !(Shape::kColumn % InstructionShape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + + static_assert(platform::is_same::value, + "Layouts must be defined for logical MatrixCoord coordinate space."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape; + }; + +private: + + static int const kElementsPerAccess = 2; + +public: + + // + // Derived quantities + // + + using AccessType = Array; + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + +private: + + /// Reference to output tensor + TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator( + TensorRef const &ref, + int lane_id + ): + ref_(ref) { + + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + + MatrixCoord lane_offset(quad, lane_in_quad * kElementsPerAccess); + + ref_.add_coord_offset(lane_offset); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + ref_.add_coord_offset(tile_offset * make_Coord(Shape::kRow, Shape::kColumn)); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator++() { + // deliberate no-op + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator--() { + // deliberate no-op + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to load from the tensor + Index pointer_offset) const { ///< loads a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + AccessType* frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + int accum_m = mma_m * InstructionShape::kM; + int accum_n = mma_n * InstructionShape::kN; + + int idx = mma_m + mma_n * Policy::MmaIterations::kRow; + + AccessType* access_ptr = reinterpret_cast(offset_ref.data() + + offset_ref.offset(TensorCoord(accum_m, accum_n))); + + frag_ptr[idx] = access_ptr[0]; + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + Fragment &frag, ///< fragment to load from the tensor + Index byte_offset) const { ///< loads a tile with a linear offset + + load_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset) const { ///< loads a tile with a logical offset in units of whole tiles + + load(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset, ///< loads a tile with a logical offset in units of whole tiles + Index pointer_offset) const { ///< loads a tile with a logical offset AND a pointer offset + + load_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } + + /// Stores a fragment to memory + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index pointer_offset) const { ///< store a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + int accum_m = mma_m * InstructionShape::kM; + int accum_n = mma_n * InstructionShape::kN; + + int idx = mma_m + mma_n * Policy::MmaIterations::kRow; + + AccessType* access_ptr = reinterpret_cast(offset_ref.data() + + offset_ref.offset(TensorCoord(accum_m, accum_n))); + + access_ptr[0] = frag_ptr[idx]; + } + } + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_byte_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index byte_offset) const { ///< store a tile with a linear offset + + store_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Stores a fragment to memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + Fragment &frag, ///< fragment to store to the tensor + TensorCoord const &tile_offset) const { ///< stores a tile with a logical offset in units of whole tiles + + store(frag, tile_offset, 0); + } + + /// Stores a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + /// fragment to store to the tensor + Fragment const &frag, + /// stores a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// stores a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + store_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It is used to load or store +/// accumulators from memory and is agnostic to layout. It could be faster if it assumed row-major +/// accumulator layout. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept | +/// WriteableRandomAccessContiguousTileIteratorConcept +/// + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Element typ + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_, + /// Interleaved N + int InterleavedN> +class MmaTensorOpAccumulatorTileIterator< + Shape_, Element_, mctlass::layout::TensorNCxHWx, + InstructionShape_, OpDelta_> { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = int8_t; + + /// Layout of source tile + using Layout = mctlass::layout::TensorNCxHWx; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + using OpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kRow % InstructionShape::kM) && + !(Shape::kColumn % InstructionShape::kN), + "Shape of warp-level Mma must be divisible by operator shape."); + + /// Number of elements in strided dimension that each STG writes + static int const kStridedPerSTG = 8; + + /// Factor to calculate reorder index to pack accumulator. + static int const kPackedFactor = Shape::kColumn / 32; + + /// Number of mma operations performed + using MmaIterations = MatrixShape; + }; + +private: + + static int const kElementsPerAccess = InterleavedN / 4; + +public: + + // + // Derived quantities + // + + struct alignas((kElementsPerAccess * sizeof_bits::value / 8)) AccessType { + Array storage; + }; + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + +private: + + /// Reference to output tensor + TensorRef ref_; + + /// Row offset index globally + LongIndex global_offset_row_; + + /// Column offset index globally + LongIndex global_offset_col_; + + /// Output tensor size + TensorCoord extent_; + + /// Alpha + float alpha_; + + /// Beta + float beta_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator( + TensorRef const &ref, + int const lane_id, + TensorCoord extent, + float alpha = 1.0f, + float beta = 0.0f + ): + ref_(ref), + extent_(extent), + alpha_(alpha), + beta_(beta) { + + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + + global_offset_row_ = quad; + + global_offset_col_ = lane_in_quad * kElementsPerAccess; + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_tile_offset(MatrixCoord const &tile_offset) { + + global_offset_row_ += tile_offset.row() * Shape::kRow; + + global_offset_col_ += tile_offset.column() * Shape::kColumn; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator++() { + // deliberate no-op + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator--() { + // deliberate no-op + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to load from the tensor + Index pointer_offset) const { ///< loads a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + AccessType* frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kN; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kM; ++mma_m) { + int accum_m = mma_m * InstructionShape::kM; + int accum_n = mma_n * InstructionShape::kN; + + int idx = mma_m + mma_n * Policy::MmaIterations::kM; + + AccessType* access_ptr = reinterpret_cast(offset_ref.data() + + accum_m * offset_ref.stride(0) + accum_n); + + frag_ptr[idx] = access_ptr[0]; + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + Fragment &frag, ///< fragment to load from the tensor + Index byte_offset) const { ///< loads a tile with a linear offset + + load_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset) const { ///< loads a tile with a logical offset in units of whole tiles + + load(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset, ///< loads a tile with a logical offset in units of whole tiles + Index pointer_offset) const { ///< loads a tile with a logical offset AND a pointer offset + + load_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } + + /// Stores a fragment to memory + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index pointer_offset) const { ///< store a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + Array output_frag_f; + Array output_frag; + + LongIndex pq = extent_.h() * extent_.w(); + + LongIndex extent_row = extent_.n() * pq; + LongIndex extent_col = extent_.c(); + + LongIndex k_major = (global_offset_col_ / InterleavedN) * pq; + Index k_minor = global_offset_col_ % InterleavedN; + LongIndex k_offset = k_major * InterleavedN + k_minor; + LongIndex k_offset_delta = pq * InterleavedN; + + LongIndex stride_n = pq * extent_.c(); + + Index n; + LongIndex pq_rem; + + unsigned int pq_mul, pq_shr; + find_divisor(pq_mul, pq_shr, pq); + + if(beta_ == 0.0f) { + MCTLASS_PRAGMA_UNROLL + for(int i = 0; i < frag.size(); ++i) { + output_frag_f[i] = frag[i]; + } + + if(InstructionShape::kM == Policy::kStridedPerSTG) { + MCTLASS_PRAGMA_UNROLL + for(int i = 0; i < frag.size(); ++i) { + output_frag[i] = (Element)(output_frag_f[i] * alpha_); + } + } else { + MCTLASS_PRAGMA_UNROLL + for(int i = 0; i < frag.size(); ++i) { + int map_i = (i / (16 * Policy::kPackedFactor)) * (16 * Policy::kPackedFactor) + + (i % (8 * Policy::kPackedFactor)) / 2 * 4 + + (i % (8 * Policy::kPackedFactor)) % 2 + + (i / (8 * Policy::kPackedFactor)) % 2 * 2; + output_frag[i] = (Element)(output_frag_f[map_i] * alpha_); + } + } + + AccessType const *frag_ptr = reinterpret_cast(&output_frag); + + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + int accum_m = mma_m * Policy::kStridedPerSTG; + + fast_divmod(n, pq_rem, global_offset_row_ + accum_m, pq, pq_mul, pq_shr); + LongIndex offset_m = n * stride_n + k_offset + pq_rem * InterleavedN; + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + + int accum_n = mma_n * InterleavedN; + + int idx = mma_n + mma_m * Policy::MmaIterations::kColumn; + + if((global_offset_row_ + accum_m < extent_row) && (global_offset_col_ + accum_n < extent_col)) { + AccessType* access_ptr = reinterpret_cast(offset_ref.data() + + offset_m + mma_n * k_offset_delta); + + access_ptr[0] = frag_ptr[idx]; + } + } + } + } else { + if(InstructionShape::kM == Policy::kStridedPerSTG) { + MCTLASS_PRAGMA_UNROLL + for(int i = 0; i < frag.size(); ++i) { + output_frag_f[i] = frag[i]; + } + } else { + MCTLASS_PRAGMA_UNROLL + for(int i = 0; i < frag.size(); ++i) { + int map_i = (i / (16 * Policy::kPackedFactor)) * (16 * Policy::kPackedFactor) + + (i % (8 * Policy::kPackedFactor)) / 2 * 4 + + (i % (8 * Policy::kPackedFactor)) % 2 + + (i / (8 * Policy::kPackedFactor)) % 2 * 2; + output_frag_f[i] = frag[map_i]; + } + } + + AccessType const *frag_ptr = reinterpret_cast(&output_frag); + + Array ref_frag; + AccessType *ref_frag_ptr = reinterpret_cast(&ref_frag); + + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + int accum_m = mma_m * Policy::kStridedPerSTG; + + fast_divmod(n, pq_rem, global_offset_row_ + accum_m, pq, pq_mul, pq_shr); + LongIndex offset_m = n * stride_n + k_offset + pq_rem * InterleavedN; + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + + int accum_n = mma_n * InterleavedN; + + int idx = mma_n + mma_m * Policy::MmaIterations::kColumn; + + if((global_offset_row_ + accum_m < extent_row) && (global_offset_col_ + accum_n < extent_col)) { + AccessType* access_ptr = reinterpret_cast(offset_ref.data() + + offset_m + mma_n * k_offset_delta); + + ref_frag_ptr[0] = access_ptr[0]; + + MCTLASS_PRAGMA_UNROLL + for(int i = 0; i < kElementsPerAccess; ++i) { + output_frag[idx * kElementsPerAccess + i] = Element(alpha_ * output_frag_f[idx * kElementsPerAccess + i] + + beta_ * ref_frag[i]); + } + + access_ptr[0] = frag_ptr[idx]; + } + } + } + } + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_byte_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index byte_offset) const { ///< store a tile with a linear offset + + store_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Stores a fragment to memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + Fragment &frag, ///< fragment to store to the tensor + TensorCoord const &tile_offset) const { ///< stores a tile with a logical offset in units of whole tiles + + store(frag, tile_offset, 0); + } + + /// Stores a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + /// fragment to store to the tensor + Fragment const &frag, + /// stores a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// stores a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + store_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } +}; + +/// This tile iterator is specialized data type with half_t for MmaTensorOpMultiplicandTileIterator +/// And it now just tested with InstructionShape_=<16, 8, 16> in warp_level +// and <16, 8, 8> in conv/device +/// Other InstructionShape_ maybe cannot work correctly. And so for threadblock and device-level. +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, half_t, + mctlass::layout::TensorOpMultiplicandCongruous::value, + 64>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = half_t; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCongruous< + sizeof_bits::value, 64>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + // Determine number of elements along outer dimension per individual LDSM op + static int const kLdsmOpOuter_16x8x16 = Layout::kElementsPerAccess; + //static int const kLdsmOpInner = 8; + static int const kLdsmOpInner_16x8x16 = kOperand == Operand::kA ? 8 : 4; + + static int const kLdsmOpInner_16x8x8 = Layout::TileShape::kStrided; + static int const kLdsmOpOuter_16x8x8 = kThreads / kLdsmOpInner_16x8x8; + + static_assert(!(Shape::kContiguous % kLdsmOpOuter_16x8x16), + "Shape of warp-level mma must be divisible by LDSM's fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsmOpInner_16x8x16), + "Shape of warp-level mma must be divisible by LDSM's fundamental tile size."); + + static_assert(!(Shape::kContiguous % kLdsmOpOuter_16x8x8), + "Shape of warp-level mma must be divisible by LDSM's fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsmOpInner_16x8x8), + "Shape of warp-level mma must be divisible by LDSM's fundamental tile size."); + + /// Shape of one individual LDSM instruction + static int const LdsmShapeStrided_16x8x16 = + InstructionShape::kStrided / kLdsmOpInner_16x8x16; + static int const LdsmShapeContiguous_16x8x16 = 4 / LdsmShapeStrided_16x8x16; + static int const LdsmShapeContiguous_16x8x8 = + InstructionShape::kContiguous / kLdsmOpOuter_16x8x8; + static int const LdsmShapeStrided_16x8x8 = InstructionShape::kStrided / kLdsmOpInner_16x8x8; + + using LdsmShape = typename platform::conditional, + layout::PitchLinearShape>::type; + + using LdsmIterations = typename platform::conditional, + layout::PitchLinearShape>::type; + + /// Number of groups for each tile + static int const kGroupsPerTile = + Shape::kStrided / InstructionShape::kStrided; + }; + +private: + + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Number of internal pointers needed to reference shared memory + static int const kPointerCount = Layout::TileShape::kContiguous * + Layout::kElementsPerAccess / + Policy::kLdsmOpOuter_16x8x8; + + /// Pointer type used for accesses + using AccessType = typename platform::conditional>::type; + /// Vectorized access is not used + static int const kElementsPerAccess = 1; + + /// Internal counter used to jump to next K partition + int k_group_idx_; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + +private: + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_[kPointerCount]; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + Layout layout_; + Index tile_row_offset; + Index tile_col_offset; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(): stride_(0), byte_offset_(0), tile_row_offset(0), tile_col_offset(0) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): + stride_(ref.stride(0)), byte_offset_(0), + k_group_idx_(0) { + + layout_ = ref.layout(); + pointer_[0] = reinterpret_cast(ref.data()); + tile_row_offset = 0; + tile_col_offset = 0; + if (InstructionShape::kStrided != 8) { + stride_ /= Layout::kElementsPerAccess; + } + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + byte_offset_ += offset * sizeof(Element); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + if (InstructionShape::kStrided == 8) { + int contiguous_offset = tile_offset.contiguous(); + if (Shape::kContiguous == + Layout::TileShape::kContiguous * Layout::kElementsPerAccess / 2) { + if (tile_offset.contiguous() % 2) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPointerCount / 2; ++i) { + AccessType const *tmp_pointer = pointer_[i]; + pointer_[i] = pointer_[i + kPointerCount / 2]; + pointer_[i + kPointerCount / 2] = tmp_pointer; + } + } + contiguous_offset = (tile_offset.contiguous() >> 1) << 1; + } + int offset = (tile_offset.strided() * InstructionShape::kStrided) * stride_ + + contiguous_offset * Shape::kContiguous; + add_pointer_offset(offset); + } + else { + tile_row_offset += tile_offset.contiguous() * Shape::kContiguous; + tile_col_offset += tile_offset.strided() * InstructionShape::kStrided; + } + return *this; + + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + if (InstructionShape::kStrided == 8) { + add_tile_offset({0, 1}); + if (kPartitionsK > 1) { + ++k_group_idx_; + // Jump to next stage + if (k_group_idx_ == Policy::kGroupsPerTile) { + k_group_idx_ = 0; + add_tile_offset( + {0, ((kPartitionsK - 1) * Policy::kGroupsPerTile)}); + } + } + } + else { + //For RowMajorXXX, the Matrices was transposed,so for Operand::kA and Operand::kB can use same variable. + byte_offset_ += InstructionShape::kStrided; + } + return *this; + } + + /// Advances the iterator along the opposite of the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + if (InstructionShape::kStrided == 8) { + byte_offset_ -= stride_ * InstructionShape::kStrided * sizeof(Element) * + kElementsPerAccess; + } + else { + byte_offset_ -= stride_ * InstructionShape::kStrided * sizeof(Element) * + Layout::kElementsPerAccess; + } + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_byte_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + // Array *fetch_ptr = + // reinterpret_cast *>(&frag); + Element *fetch_ptr = reinterpret_cast(&frag); + const size_t fetch_size = kOperand == Operand::kA ? Policy::LdsmShape::kCount * 2 : Policy::LdsmShape::kCount; + Array *fetch_element_ptr = reinterpret_cast *>(&frag); + + const int lane_id = __lane_id(); + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsmIterations::kStrided; ++s) { + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsmIterations::kContiguous; ++c) { + + int access_idx = c + s * Policy::LdsmIterations::kContiguous; + + /* AccessType const *source_ptr = + pointer_[c % kPointerCount] + + Layout::TileShape::kContiguous * (c / kPointerCount) + + Policy::kLdsmOpInner * Policy::LdsmShape::kStrided * s * stride_; + + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_; + mctlass::arch::ldsm( + fetch_ptr[access_idx], + source_byte_ptr + ); */ + if (InstructionShape::kStrided == 8) { + if(kOperand == Operand::kA) { + int row = (lane_id & 0xf) + c * InstructionShape::kContiguous; + int col = ((lane_id >> 4) << 1) ^ 0x7; + for (int i=0; i < 2; i++) { + Index idx = layout_(MatrixCoord(row, col - i)); + AccessType const *source_ptr = pointer_[0] + idx; + char const *source_byte_ptr = reinterpret_cast(source_ptr) + + byte_offset + byte_offset_; + fetch_ptr[c * 4 + i] = *reinterpret_cast(source_byte_ptr); + } + } + else { + int col = ((lane_id >> 4) << 1) ^ 0x7; + int row = (lane_id & 0x7) + c * InstructionShape::kStrided; + for (int i=0; i < 2; i++) { + Index idx = layout_(MatrixCoord(row, col - i)); + AccessType const *source_ptr = pointer_[0] + idx; + char const *source_byte_ptr = reinterpret_cast(source_ptr) + + byte_offset + byte_offset_; + fetch_ptr[c * 2 + i] = *reinterpret_cast(source_byte_ptr); + } + } + } + else { + int row = (lane_id & (InstructionShape::kContiguous - 1)) + c * InstructionShape::kContiguous + tile_row_offset; + int col = (((lane_id >> 4) << 2) ^ (InstructionShape::kStrided - 1)) + s * InstructionShape::kStrided + byte_offset_ + tile_col_offset; + Element const *ptr = reinterpret_cast(pointer_[0]) + byte_offset; + for (int i = 0; i < Policy::LdsmShape::kCount; ++i) { + Index idx = layout_(MatrixCoord(row, col - i)); + fetch_element_ptr[access_idx][i] = ptr[idx]; + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + //load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + load_with_byte_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = + tile_offset.contiguous() * Shape::kContiguous / Layout::kElementsPerAccess + + tile_offset.strided() * InstructionShape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no op + } +}; + +//////////////////////////////////////////////////////////////////////////////// +/// This tile iterator is specialized data type with half_t for MmaTensorOpMultiplicandTileIterator +/// And it now just tested with InstructionShape_=<16, 8, 16> in warp_level +// and <16, 8, 8> in conv/device +/// Other InstructionShape_ maybe cannot work correctly. And so for threadblock and device-level. +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Element number when the layout crosses (in units of elements) + int Crosswise, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, half_t, + mctlass::layout::TensorOpMultiplicandCrosswise::value, + Crosswise>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for " + "A or B operands to warp-level Mma."); + + /// Element type + using Element = half_t; + + /// Element number when the layout crosses + static int const kCrosswise = Crosswise; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCrosswise< + sizeof_bits::value, kCrosswise>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + // Determine number of elements along outer dimension per individual LDSM op + static int const kLdsmOpOuter = Layout::kElementsPerAccess; + static int const kLdsmOpInner_16x8x16 = kOperand == Operand::kA ? 8 : 4; + static int const kLdsmOpInner_16x8x8 = 4; + + static_assert(!(Shape::kContiguous % kLdsmOpOuter), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsmOpInner_16x8x16), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsmOpInner_16x8x8), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + /// Shape of one individual LDSM instruction + static int const LdsmShapeContiguous = + InstructionShape::kContiguous / kLdsmOpOuter; + static int const LdsmShapeStrided_16x8x16 = + ((4 / LdsmShapeContiguous * kLdsmOpInner_16x8x16) > Shape::kStrided) + ? (Shape::kStrided / kLdsmOpInner_16x8x16) + : (4 / LdsmShapeContiguous); + static int const LdsmShapeStrided_16x8x8 = + ((4 / LdsmShapeContiguous * kLdsmOpInner_16x8x8) > Shape::kStrided) + ? (Shape::kStrided / kLdsmOpInner_16x8x8) + : (4 / LdsmShapeContiguous); + using LdsmShape = typename platform::conditional, + layout::PitchLinearShape>::type; + /// Number and arrangement of LDSM instructions + using LdsmIterations = typename platform::conditional, + layout::PitchLinearShape<1, Shape::kStrided / kLdsmOpInner_16x8x16 / LdsmShape::kStrided>>::type; + + /// + static int const kGroupsPerTile = Layout::TileShape::kContiguous / + Layout::kFactor / LdsmShape::kContiguous; + }; + + private: + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = Array; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + + private: + + /// Total number of sections. The memory is divided into stages. One stage + /// can store one tile. Stage is divided into sections. Interleaved layout + /// can have multiple sections in a stage. The rest layout only has one section + /// in a stage. + int sections_; + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + /// Internal counter used to determine when to increment byte offset and when + /// to XOR it + int k_group_idx_; + + Layout layout_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() + : pointer_(nullptr), + sections_(0), + stride_(0), + byte_offset_(0), + k_group_idx_(0) {} + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : pointer_(reinterpret_cast(ref.data())), + sections_(ref.stride(0) / kCrosswise), + // stride_ = kCrosswise x sections_ x kFactor + stride_(ref.stride(0) * Layout::kFactor / Layout::kElementsPerAccess), + byte_offset_(0), + k_group_idx_(0) { + + layout_ = ref.layout(); + byte_offset_ = 0; + if (InstructionShape::kContiguous == 8) { + int row, col; + if (kOperand == Operand::kA) { + row = lane_id & 0xf; + col = ((lane_id >> 4) << 1) ^ 0x7; + } + else { + row = lane_id & 0x7; + col = ((lane_id >> 4) << 1) ^ 0x7; + } + TensorCoord coord = make_Coord(col, row); + Layout lo(ref.stride(0)); + byte_offset_ = lo(coord) * sizeof_bits::value / 8; + } + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + byte_offset_ += offset * sizeof_bits::value / 8; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + + byte_offset_ ^= k_groups_delta * sizeof_bits::value * + Layout::kElementsPerAccess * + Policy::LdsmShape::kContiguous / 8; + pointer_ += + tile_offset.strided() * stride_ * Shape::kStrided / Layout::kFactor + + whole_tiles * stride_ / sections_; + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative( + TensorCoord const &tile_offset) { + + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + if (k_groups_delta < 0) { + whole_tiles -= 1; + k_groups_delta += Policy::kGroupsPerTile; + } + + if ((Policy::kGroupsPerTile / kPartitionsK) >= 2) { + byte_offset_ ^= (k_groups_delta & 1) * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + if ((Policy::kGroupsPerTile / kPartitionsK) >= 4) { + byte_offset_ ^= ((k_groups_delta + (k_group_idx_ & 1)) & 2) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + if ((Policy::kGroupsPerTile / kPartitionsK) == 8) { + byte_offset_ ^= ((k_groups_delta + (k_group_idx_ & 3)) & 4) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + + k_group_idx_ += k_groups_delta; + whole_tiles += k_group_idx_ / (Policy::kGroupsPerTile / kPartitionsK); + k_group_idx_ = k_group_idx_ % (Policy::kGroupsPerTile / kPartitionsK); + + pointer_ += + tile_offset.strided() * stride_ * Shape::kStrided / Layout::kFactor + + whole_tiles * stride_ / sections_; + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator++() { + if (InstructionShape::kContiguous == 8) { + if ((Policy::kGroupsPerTile / kPartitionsK) > 1) { + int mask = ((Policy::kGroupsPerTile / kPartitionsK) == 8) + ? 3 + : (((Policy::kGroupsPerTile / kPartitionsK) == 4) ? 1 : 0); + + if (((k_group_idx_ & mask) % 2) == 0) + byte_offset_ ^= 1 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + else if ((k_group_idx_ & mask) == 1) + byte_offset_ ^= 3 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + else if ((k_group_idx_ & mask) == 3) + byte_offset_ ^= 7 * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + k_group_idx_++; + + if (k_group_idx_ == (Policy::kGroupsPerTile / kPartitionsK)) { + k_group_idx_ = 0; + add_tile_offset({Policy::kGroupsPerTile, 0}); + } + } + else { + byte_offset_ += InstructionShape::kContiguous; + } + return *this; + + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator--() { assert(0); } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + if (InstructionShape::kContiguous == 8) { + add_tile_offset(-tile_offset); + } + else { + printf("Not Impl Now.\n"); + //add_tile_offset(-tile_offset); + } + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { load_with_byte_offset(frag, 0); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + // Array *fetch_ptr = + // reinterpret_cast *>(&frag); + Array *fetch_ptr = + reinterpret_cast *>(&frag); + const size_t frag_fetch_size = kOperand == Operand::kA ? Policy::LdsmShape::kCount * 2 : Policy::LdsmShape::kCount; + Array *fetch_element_ptr = + reinterpret_cast *>(&frag); + int ldm = stride_ * Layout::kElementsPerAccess / Layout::kFactor; + + const int lane_id = __lane_id(); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsmIterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsmIterations::kContiguous; ++c) { + int access_idx = c + s * Policy::LdsmIterations::kContiguous; + int kLdsmOpInner = 0; + if (InstructionShape::kContiguous == 8) { + kLdsmOpInner = Policy::kLdsmOpInner_16x8x8; + } + else { + kLdsmOpInner = Policy::kLdsmOpInner_16x8x16; + } + AccessType const *source_ptr = + pointer_ + Policy::LdsmShape::kContiguous * c + + kLdsmOpInner / Layout::kFactor * + Policy::LdsmShape::kStrided * s * stride_; + + // char const *source_byte_ptr = + // reinterpret_cast(source_ptr) + byte_offset + + // byte_offset_; + + // mctlass::arch::ldsm( + // fetch_ptr[access_idx], source_byte_ptr); + if (InstructionShape::kContiguous == 8) { + char const *source_byte_ptr = + reinterpret_cast(source_ptr) + byte_offset + + byte_offset_; + if(kOperand == Operand::kA){ + fetch_ptr[access_idx][0] = *reinterpret_cast(source_byte_ptr); + fetch_ptr[access_idx][1] = *(reinterpret_cast(source_byte_ptr) - 1); + } + else { + Element const *p = reinterpret_cast(source_byte_ptr); + Element const *p0 = p - 1; + Element const *p1 = p0 + 8 * ldm; + fetch_ptr[access_idx][0] = p[0]; + fetch_ptr[access_idx][1] = p0[0]; + fetch_ptr[access_idx][2] = p1[1]; + fetch_ptr[access_idx][3] = p1[0]; + } + } + else { + Element const *ptr = reinterpret_cast(source_ptr) + byte_offset / sizeof(Element); + Index row = ((lane_id >> 4) << 2) ^ (InstructionShape::kContiguous - 1) + byte_offset_; + Index col = lane_id & (InstructionShape::kStrided - 1); + for (int i = 0; i < Policy::LdsmShape::kCount; ++i) { + Index idx = layout_(MatrixCoord(row - i, col)); + fetch_element_ptr[access_idx][i] = ptr[idx]; + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = tile_offset.contiguous() * + InstructionShape::kContiguous / + Layout::kElementsPerAccess + + tile_offset.strided() * Shape::kStrided * stride_; + + byte_offset += sizeof_bits::value * pointer_offset / 8; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + + k_group_idx_ = k_group % (Policy::kGroupsPerTile / kPartitionsK); + } +}; + +/// This tile iterator is specialized data type with float and InstructionShape_=<16,8,16> for MmaTensorOpMultiplicandTileIterator +/// And it now just tested with InstructionShape_=<16, 8, 16> in warp_level. +/// Threadblock-level and device-level maybe cannot work correctly. +/// Date: 06/12/2022 +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_> +class MmaTensorOpAccumulatorTileIterator< + Shape_, float, mctlass::layout::RowMajor, gemm::GemmShape<16, 8, 16>, OpDelta_> { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = float; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = gemm::GemmShape<16, 8, 16>; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + using OpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static bool const kDivisible = + !(Shape::kRow % InstructionShape::kM) && + !(Shape::kColumn % InstructionShape::kN); + + static_assert(platform::is_same::value, + "Layouts must be defined for logical MatrixCoord coordinate space."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + (Shape::kRow + InstructionShape::kM - 1) / InstructionShape::kM, + (Shape::kColumn + InstructionShape::kN - 1) / InstructionShape::kN + >; + }; + +private: + + // Assume accumulator tile is an arrangement of 8-by-8 tiles replicated over the entire + // shape, with each quad mapped to one row and each thread mapped to 1/4 of the elements + // of that row. The accumulators within one row are assumed to be consecutive. + static int const kElementsPerAccess = InstructionShape::kN / 4; + static int const kRowsPerTile = 8; + static int const kAccumulatorRows = InstructionShape::kM / kRowsPerTile; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array< + Element, + Policy::MmaIterations::kCount * InstructionShape::kMN / kThreads>; + +private: + + /// Reference to output tensor + TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator( + TensorRef const &ref, + int lane_id + ): + ref_(ref) { + + int quad = (lane_id >> 4) << 2; + int lane_in_quad = lane_id & 0x7; + MatrixCoord lane_offset(quad, lane_in_quad); + ref_.add_coord_offset(lane_offset); + + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + ref_.add_coord_offset(tile_offset * make_Coord(Shape::kRow, Shape::kColumn)); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator++() { + // deliberate no-op + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator--() { + // deliberate no-op + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator-=(TensorCoord const &tile_offset) { + printf("Not Impl Now.\n"); + //add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to load from the tensor + Index pointer_offset) const { ///< loads a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + + frag[mma_accum_start + row * kElementsPerAccess + col] = offset_ref.at({accum_m, accum_n}); + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + Fragment &frag, ///< fragment to load from the tensor + Index byte_offset) const { ///< loads a tile with a linear offset + + load_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset) const { ///< loads a tile with a logical offset in units of whole tiles + + load(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset, ///< loads a tile with a logical offset in units of whole tiles + Index pointer_offset) const { ///< loads a tile with a logical offset AND a pointer offset + + load_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } + + /// Stores a fragment to memory + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index pointer_offset) const { ///< store a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + row * kElementsPerAccess + col; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn; + + int idx = mma_accum_start + row * kElementsPerAccess + col; + + offset_ref.at({accum_m, accum_n}) = frag[idx]; + + } + } + } + + } + + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_byte_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index byte_offset) const { ///< store a tile with a linear offset + + store_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Stores a fragment to memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + Fragment &frag, ///< fragment to store to the tensor + TensorCoord const &tile_offset) const { ///< stores a tile with a logical offset in units of whole tiles + + store(frag, tile_offset, 0); + } + + /// Stores a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + /// fragment to store to the tensor + Fragment const &frag, + /// stores a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// stores a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + store_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } +}; + + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_> +class MmaTensorOpAccumulatorTileIterator< + Shape_, float, mctlass::layout::RowMajor, gemm::GemmShape<16, 16, 16>, OpDelta_> { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = float; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = gemm::GemmShape<16, 16, 16>; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + using OpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 64; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static bool const kDivisible = + !(Shape::kRow % InstructionShape::kM) && + !(Shape::kColumn % InstructionShape::kN); + + static_assert(platform::is_same::value, + "Layouts must be defined for logical MatrixCoord coordinate space."); + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + (Shape::kRow + InstructionShape::kM - 1) / InstructionShape::kM, + (Shape::kColumn + InstructionShape::kN - 1) / InstructionShape::kN + >; + }; + +private: + + // Assume accumulator tile is an arrangement of 8-by-8 tiles replicated over the entire + // shape, with each quad mapped to one row and each thread mapped to 1/4 of the elements + // of that row. The accumulators within one row are assumed to be consecutive. + static int const kElementsPerAccess = InstructionShape::kN / 4; + static int const kRowsPerTile = 16; + static int const kAccumulatorRows = InstructionShape::kM / kRowsPerTile; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array< + Element, + Policy::MmaIterations::kCount * InstructionShape::kMN / kThreads>; + +private: + + /// Reference to output tensor + TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator( + TensorRef const &ref, + int lane_id + ): + ref_(ref) { + + int quad = (lane_id >> 4) << 2; + int lane_in_quad = lane_id & 0xf; + MatrixCoord lane_offset(quad, lane_in_quad); + ref_.add_coord_offset(lane_offset); + + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + ref_.add_coord_offset(tile_offset * make_Coord(Shape::kRow, Shape::kColumn)); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator++() { + // deliberate no-op + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpAccumulatorTileIterator & operator--() { + // deliberate no-op + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpAccumulatorTileIterator & operator-=(TensorCoord const &tile_offset) { + printf("Not Impl Now.\n"); + //add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to load from the tensor + Index pointer_offset) const { ///< loads a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + + row * kRowsPerTile; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn + col; + + frag[mma_accum_start + row * kElementsPerAccess + col] = offset_ref.at({accum_m, accum_n}); + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + Fragment &frag, ///< fragment to load from the tensor + Index byte_offset) const { ///< loads a tile with a linear offset + + load_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset) const { ///< loads a tile with a logical offset in units of whole tiles + + load(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset, ///< loads a tile with a logical offset in units of whole tiles + Index pointer_offset) const { ///< loads a tile with a logical offset AND a pointer offset + + load_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } + + /// Stores a fragment to memory + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index pointer_offset) const { ///< store a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = kAccumulatorRows * kElementsPerAccess * + (mma_n * Policy::MmaIterations::kRow + mma_m); + + MCTLASS_PRAGMA_UNROLL + for (int row = 0; row < kAccumulatorRows; ++row) { + MCTLASS_PRAGMA_UNROLL + for (int col = 0; col < kElementsPerAccess; ++col) { + + int accum_m = mma_m * InstructionShape::kM * OpDelta::kRow + row * kElementsPerAccess + col; + int accum_n = mma_n * InstructionShape::kN * OpDelta::kColumn; + + int idx = mma_accum_start + row * kElementsPerAccess + col; + + offset_ref.at({accum_m, accum_n}) = frag[idx]; + + } + } + } + + } + + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_DEVICE + void store_with_byte_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index byte_offset) const { ///< store a tile with a linear offset + + store_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Stores a fragment to memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + Fragment &frag, ///< fragment to store to the tensor + TensorCoord const &tile_offset) const { ///< stores a tile with a logical offset in units of whole tiles + + store(frag, tile_offset, 0); + } + + /// Stores a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void store( + /// fragment to store to the tensor + Fragment const &frag, + /// stores a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// stores a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + store_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 64-thread TensorOps. It's used to load from shared +/// memory and therefore must be initialized with a TensorRef to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::TensorOpMultiplicandCongruous::value, + 64>, + InstructionShape_, OpDelta_, 64, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCongruous< + sizeof_bits::value, 64>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 64; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + // Determine number of elements along outer dimension per individual LDSM op + static int const kLdsmOpOuter = Layout::kElementsPerAccess; + static int const kLdsmOpInner = 8; + + static_assert(!(Shape::kContiguous % kLdsmOpOuter), + "Shape of warp-level mma must be divisible by LDSM's fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsmOpInner), + "Shape of warp-level mma must be divisible by LDSM's fundamental tile size."); + + /// Shape of one individual LDSM instruction + static int const LdsmShapeStrided = + InstructionShape::kStrided / kLdsmOpInner; + static int const LdsmShapeContiguous = 4 / LdsmShapeStrided; + using LdsmShape = + layout::PitchLinearShape; + + /// Number and arrangement of LDSM instructions + using LdsmIterations = layout::PitchLinearShape< + Shape::kContiguous / Layout::kElementsPerAccess / LdsmShapeContiguous, + 1>; + + /// Number of groups for each tile + static int const kGroupsPerTile = + Shape::kStrided / InstructionShape::kStrided; + }; + +private: + + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Number of internal pointers needed to reference shared memory + static int const kPointerCount = + Layout::TileShape::kContiguous / Policy::LdsmShape::kContiguous; + + /// Pointer type used for accesses + using AccessType = Array; + + /// Internal counter used to jump to next K partition + int k_group_idx_; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + +private: + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_[kPointerCount]; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + Layout layout_; + Index tile_row_offset; + Index tile_col_offset; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(): stride_(0), byte_offset_(0) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): + stride_(ref.stride(0) / Layout::kElementsPerAccess), byte_offset_(0), + k_group_idx_(0) { + + layout_ = ref.layout(); + pointer_[0] = reinterpret_cast(ref.data()); + tile_row_offset = 0; + tile_col_offset = 0; + + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + byte_offset_ += offset * sizeof(Element); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + tile_row_offset += tile_offset.contiguous() * Shape::kContiguous; + tile_col_offset += tile_offset.strided() * InstructionShape::kStrided; + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + byte_offset_ += InstructionShape::kStrided; + return *this; + } + + /// Advances the iterator along the opposite of the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + byte_offset_ -= stride_ * InstructionShape::kStrided * sizeof(Element) * + Layout::kElementsPerAccess; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_byte_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + Array *fetch_element_ptr = + reinterpret_cast *>(&frag); + + const int lane_id = mctlass::arch::LaneId(); + Element const *ptr = reinterpret_cast(pointer_[0]) + byte_offset; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsmIterations::kStrided; ++s) { + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsmIterations::kContiguous; ++c) { + + int access_idx = c + s * Policy::LdsmIterations::kContiguous; + + int row = (lane_id & (InstructionShape::kContiguous - 1)) + + c * InstructionShape::kContiguous + tile_row_offset; + int col = (((lane_id >> 4) << 2) ^ (InstructionShape::kStrided - 1)) + + s * InstructionShape::kStrided + byte_offset_ + tile_col_offset; + + for (int i = 0; i < Policy::LdsmShape::kCount; ++i) { + Index idx = layout_(MatrixCoord(row, col - i)); + fetch_element_ptr[access_idx][i] = ptr[idx]; + } + + } + } + + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = + tile_offset.contiguous() * Shape::kContiguous / Layout::kElementsPerAccess + + tile_offset.strided() * InstructionShape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no op + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 64-thread TensorOps. It's used to load from shared +/// memory and therefore must be initialized with a TensorRef to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::MacaColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element_))>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA, + "MmaTensorOpMultiplicandIterator for ColumnMajor Congruous may " + "only be instantiated for A operand to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::MacaColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element_))>; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 64; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element_))>, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): iterator_({ref.data(), ref.stride()}, lane_id) { + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.contiguous(), tile_offset.strided()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 64-thread TensorOps. It's used to load from shared +/// memory and therefore must be initialized with a TensorRef to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::MacaRowMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element_))>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator for RowMajor Congruous may " + "only be instantiated for B operand to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::MacaRowMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element_))>; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 64; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element_))>, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): iterator_({ref.data(), ref.stride()}, lane_id) { + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.strided(), tile_offset.contiguous()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 64-thread TensorOps. It's used to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Element number when the layout crosses (in units of elements) + int Crosswise, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::TensorOpMultiplicandCrosswise::value, + Crosswise>, + InstructionShape_, OpDelta_, 64, PartitionsK_> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for " + "A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Element number when the layout crosses + static int const kCrosswise = Crosswise; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCrosswise< + sizeof_bits::value, kCrosswise>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 64; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + // Determine number of elements along outer dimension per individual LDSM op + static int const kLdsmOpOuter = Layout::kElementsPerAccess; + static int const kLdsmOpInner = 8; + + static_assert(!(Shape::kContiguous % kLdsmOpOuter), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + static_assert(!(Shape::kStrided % kLdsmOpInner), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + /// Shape of one individual LDSM instruction + static int const LdsmShapeContiguous = + InstructionShape::kContiguous / kLdsmOpOuter; + static int const LdsmShapeStrided = + ((4 / LdsmShapeContiguous * kLdsmOpInner) > Shape::kStrided) + ? (Shape::kStrided / kLdsmOpInner) + : (4 / LdsmShapeContiguous); + using LdsmShape = + layout::PitchLinearShape; + + /// Number and arrangement of LDSM instructions + using LdsmIterations = + layout::PitchLinearShape<1, Shape::kStrided / kLdsmOpInner / + LdsmShape::kStrided>; + + /// + static int const kGroupsPerTile = Layout::TileShape::kContiguous / + Layout::kFactor / LdsmShape::kContiguous; + }; + + private: + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = Array; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + + private: + + /// Total number of sections. The memory is divided into stages. One stage + /// can store one tile. Stage is divided into sections. Interleaved layout + /// can have multiple sections in a stage. The rest layout only has one section + /// in a stage. + int sections_; + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + /// Internal counter used to determine when to increment byte offset and when + /// to XOR it + int k_group_idx_; + Layout layout_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() + : pointer_(nullptr), + sections_(0), + stride_(0), + byte_offset_(0), + k_group_idx_(0) {} + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : pointer_(reinterpret_cast(ref.data())), + sections_(ref.stride(0) / kCrosswise), + // stride_ = kCrosswise x sections_ x kFactor + stride_(ref.stride(0) * Layout::kFactor / Layout::kElementsPerAccess), + byte_offset_(0), + k_group_idx_(0) { + layout_ = ref.layout(); + byte_offset_ = 0; + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + byte_offset_ += offset * sizeof_bits::value / 8; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + + byte_offset_ ^= k_groups_delta * sizeof_bits::value * + Layout::kElementsPerAccess * + Policy::LdsmShape::kContiguous / 8; + pointer_ += + tile_offset.strided() * stride_ * Shape::kStrided / Layout::kFactor + + whole_tiles * stride_ / sections_; + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative( + TensorCoord const &tile_offset) { + + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + if (k_groups_delta < 0) { + whole_tiles -= 1; + k_groups_delta += Policy::kGroupsPerTile; + } + + if ((Policy::kGroupsPerTile / kPartitionsK) >= 2) { + byte_offset_ ^= (k_groups_delta & 1) * Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + if ((Policy::kGroupsPerTile / kPartitionsK) >= 4) { + byte_offset_ ^= ((k_groups_delta + (k_group_idx_ & 1)) & 2) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + if ((Policy::kGroupsPerTile / kPartitionsK) == 8) { + byte_offset_ ^= ((k_groups_delta + (k_group_idx_ & 3)) & 4) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * + Layout::kElementsPerAccess / 8; + } + + k_group_idx_ += k_groups_delta; + whole_tiles += k_group_idx_ / (Policy::kGroupsPerTile / kPartitionsK); + k_group_idx_ = k_group_idx_ % (Policy::kGroupsPerTile / kPartitionsK); + + pointer_ += + tile_offset.strided() * stride_ * Shape::kStrided / Layout::kFactor + + whole_tiles * stride_ / sections_; + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator++() { + + byte_offset_ += InstructionShape::kContiguous; + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator--() { assert(0); } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { load_with_byte_offset(frag, 0); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + Array *fetch_element_ptr = + reinterpret_cast *>(&frag); + + const int lane_id = mctlass::arch::LaneId(); + + Index row = ((lane_id >> 4) << 2) ^ (InstructionShape::kContiguous - 1) + byte_offset_; + Index col = lane_id & (InstructionShape::kStrided - 1); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsmIterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsmIterations::kContiguous; ++c) { + int access_idx = c + s * Policy::LdsmIterations::kContiguous; + AccessType const *source_ptr = + pointer_ + Policy::LdsmShape::kContiguous * c + + Policy::kLdsmOpInner / Layout::kFactor * + Policy::LdsmShape::kStrided * s * stride_; + Element const *ptr = reinterpret_cast(source_ptr); + + for (int i = 0; i < Policy::LdsmShape::kCount; ++i) { + Index idx = layout_(MatrixCoord(row - i, col)); + fetch_element_ptr[access_idx][i] = ptr[idx]; + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = tile_offset.contiguous() * + InstructionShape::kContiguous / + Layout::kElementsPerAccess + + tile_offset.strided() * Shape::kStrided * stride_; + + byte_offset += sizeof_bits::value * pointer_offset / 8; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + k_group_idx_ = k_group % (Policy::kGroupsPerTile / kPartitionsK); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 64-thread TensorOps. It's used to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Element number when the layout crosses (in units of elements) + int Crosswise, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::MacaColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Crosswise>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator for ColumnMajor Crosswise may " + "only be instantiated for B operand to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// KBlock size + static int const kCrosswise = Crosswise; + + /// Layout of source tile + using Layout = mctlass::layout::MacaColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, kCrosswise>; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 64; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCrosswise::value, + kCrosswise>, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + + private: + /// Underlying tile iterator + Base iterator_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() {} + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : iterator_({ref.data(), ref.stride()}, lane_id) {} + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset_negative({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator++() { + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator--() { + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { iterator_.load(frag); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, {tile_offset.contiguous(), tile_offset.strided()}, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Element number when the layout crosses (in units of elements) + int Crosswise, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::MacaRowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Crosswise>, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA, + "MmaTensorOpMultiplicandIterator for RowMajor Crosswise may " + "only be instantiated for A operand to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Element number when the layout crosses + static int const kCrosswise = Crosswise; + + /// Layout of source tile + using Layout = mctlass::layout::MacaRowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, kCrosswise>; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 64; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCrosswise::value, + kCrosswise>, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + + private: + /// Underlying tile iterator + Base iterator_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() {} + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : iterator_({ref.data(), ref.stride()}, lane_id) {} + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset_negative({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator++() { + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &operator--() { + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { iterator_.load(frag); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, {tile_offset.strided(), tile_offset.contiguous()}, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm70.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm70.h new file mode 100644 index 0000000..91abe5e --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm70.h @@ -0,0 +1,3094 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines iterators used by warp-level matrix multiply operations targeting Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/gemm.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor_op_multiplicand_sm70.h" + +#include "mctlass/platform/platform.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Operand identity + Operand Operand, + /// Data type of A elements + typename Element_, + /// Layout of operand + typename Layout_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + int OpDelta_, + /// Number of threads participating in one matrix operation + int Threads> +class MmaVoltaTensorOpMultiplicandTileIterator; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_> +class MmaVoltaTensorOpMultiplicandTileIterator< + Shape_, Operand::kA, Element_, + mctlass::layout::VoltaTensorOpMultiplicandCongruous< + sizeof_bits::value>, + InstructionShape_, OpDelta_, 32> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kA; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::VoltaTensorOpMultiplicandCongruous::value>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + using LdsShape = layout::PitchLinearShape< + 32, + 4 + >; + + // LdsShapes are arranged in the strided direction in SMEM + using LdsIterations = layout::PitchLinearShape< + InstructionShape::kStrided / LdsShape::kStrided, + Shape::kContiguous / LdsShape::kContiguous + >; + }; + +private: + + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Number of internal pointers needed to reference shared memory + static int const kPointerCount = 2; + + /// Pointer type used for accesses + using AccessType = AlignedArray; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + +private: + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_[kPointerCount]; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator(): stride_(0), byte_offset_(0) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): + stride_(ref.stride(0) / Layout::kElementsPerAccess), byte_offset_(0) { + // swizzle patterns for operandA LDS are + // 1. (tid[4] << 3) | (tid[2:0] ^ tid[4]) + // 2. (tid[4] << 3) | (tid[2:0] ^ tid[4] ^ 0b10010) + + int vec_row = (lane_id >> 4); // tid[4] + int vec_col = ((lane_id & 4) >> 2); // tid[2] + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPointerCount; ++i) { + + if(i == 1) { + vec_row |= 2; + } + int access_contiguous_idx = (vec_col << 2) | ((lane_id & 3) ^ vec_row); + int access_contiguous = access_contiguous_idx; + + int access_strided = vec_row; + pointer_[i] = reinterpret_cast(ref.data()) + + access_contiguous + access_strided * stride_; + } + + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + byte_offset_ += offset * sizeof(Element); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + int contiguous_offset = tile_offset.contiguous(); + int strided_offset = tile_offset.strided(); + + // To support 32x32 tile size + if (Shape::kContiguous == Policy::LdsShape::kContiguous) { + if (contiguous_offset % 2) { + AccessType const *tmp_pointer = pointer_[0]; + pointer_[0] = pointer_[1]; + pointer_[1] = tmp_pointer; + } + contiguous_offset = contiguous_offset / 2 * 2; + } + + int offset = (strided_offset * InstructionShape::kStrided) * stride_ * + Layout::kElementsPerAccess + + contiguous_offset * Shape::kContiguous; + + add_pointer_offset(offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator++() { + byte_offset_ += stride_ * InstructionShape::kStrided * sizeof(Element) * + Layout::kElementsPerAccess; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator--() { + byte_offset_ -= stride_ * InstructionShape::kStrided * sizeof(Element) * + Layout::kElementsPerAccess; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_byte_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + AccessType * fetch_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsIterations::kStrided; ++s) { + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsIterations::kContiguous; ++c) { + + int access_idx = c + s * Policy::LdsIterations::kContiguous; + + AccessType const *source_ptr = pointer_[s & 1] + + Policy::LdsShape::kContiguous * c + + Policy::LdsShape::kStrided * (s / 2) * stride_; + + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_; + fetch_ptr[access_idx] = *(reinterpret_cast (source_byte_ptr)); + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = + tile_offset.contiguous() * Shape::kContiguous / + Layout::kElementsPerAccess + + tile_offset.strided() * InstructionShape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation here + } +}; + +////////////////////////////////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_> + +class MmaVoltaTensorOpMultiplicandTileIterator< + Shape_, Operand::kB, Element_, + mctlass::layout::VoltaTensorOpMultiplicandBCongruous< + sizeof_bits::value>, + InstructionShape_, OpDelta_, 32> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kB; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::VoltaTensorOpMultiplicandBCongruous::value>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kContiguous % InstructionShape::kContiguous), + "Shape of warp-level Mma must be divisible by operator shape."); + + using LdsShape = layout::PitchLinearShape< + 32, + 4 + >; + + using LdsIterations = layout::PitchLinearShape< + Shape::kContiguous / LdsShape::kContiguous, + InstructionShape::kStrided / LdsShape::kStrided + >; + }; + +private: + + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = AlignedArray; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile, needs on more time number of registers + using Fragment = Array; + +private: + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator(): stride_(0), byte_offset_(0) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): + stride_(ref.stride(0) / Layout::kElementsPerAccess), byte_offset_(0) { + + // swizzle pattern is (tid & (3 << 3) | (tid[1:0] ^ tid[4:3])) + int access_strided = (lane_id >> 3) & 0x3; + int access_contiguous = ((lane_id ^ (lane_id >> 3)) & 0x3); + + pointer_ = reinterpret_cast(ref.data()) + + access_contiguous + access_strided * stride_; + + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + byte_offset_ += offset * sizeof(Element); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + int contiguous_offset = tile_offset.contiguous(); + int strided_offset = tile_offset.strided(); + + int offset = (strided_offset * InstructionShape::kStrided) * stride_ * + Layout::kElementsPerAccess + + contiguous_offset * Shape::kContiguous; + + add_pointer_offset(offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator++() { + byte_offset_ += stride_ * InstructionShape::kStrided * sizeof(Element) * + Layout::kElementsPerAccess; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator--() { + byte_offset_ += stride_ * InstructionShape::kStrided * sizeof(Element) * + Layout::kElementsPerAccess; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_byte_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + AccessType * fetch_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsIterations::kStrided; ++s) { + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsIterations::kContiguous; ++c) { + + int access_idx = c + s * Policy::LdsIterations::kContiguous; + + AccessType const *source_ptr = pointer_ + + Policy::LdsShape::kContiguous / Layout::kElementsPerAccess * c + + Policy::LdsShape::kStrided * s * stride_; + + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_; + fetch_ptr[access_idx] = *(reinterpret_cast (source_byte_ptr)); + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = + tile_offset.contiguous() * Shape::kContiguous / + Layout::kElementsPerAccess + + tile_offset.strided() * InstructionShape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation here + } +}; + +////////////////////////////////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to load from shared +/// memory and therefore must be initialized with a TensorRef to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_> +class MmaVoltaTensorOpMultiplicandTileIterator< + Shape_, Operand::kA, Element_, + mctlass::layout::ColumnMajorVoltaTensorOpMultiplicandCongruous< + sizeof_bits::value>, + InstructionShape_, OpDelta_, 32> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kA; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::ColumnMajorVoltaTensorOpMultiplicandCongruous::value>; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaVoltaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::VoltaTensorOpMultiplicandCongruous::value>, + layout::PitchLinearShape, + kOpDelta, kThreads>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): iterator_({ref.data(), ref.stride()}, lane_id) { + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.contiguous(), tile_offset.strided()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to load from shared +/// memory and therefore must be initialized with a TensorRef to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_> +class MmaVoltaTensorOpMultiplicandTileIterator< + Shape_, Operand::kB, Element_, + mctlass::layout::RowMajorVoltaTensorOpMultiplicandBCongruous< + sizeof_bits::value>, + InstructionShape_, OpDelta_, 32> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kB; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajorVoltaTensorOpMultiplicandBCongruous::value>; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaVoltaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::VoltaTensorOpMultiplicandBCongruous::value>, + layout::PitchLinearShape, + kOpDelta, kThreads>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): iterator_({ref.data(), ref.stride()}, lane_id) { + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.strided(), tile_offset.contiguous()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It is used to load or store +/// accumulators from memory and is agnostic to layout. It could be faster if it assumed row-major +/// accumulator layout. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept | +/// WriteableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Layout of operand in memory + typename Layout_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions, concept: MatrixShape) + typename OpDelta_> +class MmaVoltaTensorOpAccumulatorTileIterator { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kC; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = Layout_; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + using OpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + + /// Volta Tensor Op uses 32x32 interleaved tile + using InterleavedTile = MatrixShape<32, 32>; + + static_assert(!(Shape::kRow % InterleavedTile::kRow) && !(Shape::kColumn % InterleavedTile::kColumn), + "Shape of warp-level Mma must be divisible by operator shape."); + + static_assert(platform::is_same::value, + "Layouts must be defined for logical MatrixCoord coordinate space."); + + /// Number of mma operations performed + using TileIterations = MatrixShape< + Shape::kRow / InterleavedTile::kRow, + Shape::kColumn / InterleavedTile::kColumn + >; + + using MmaIterations = + MatrixShape; + }; + +private: + + // Assume accumulator tile is multipile interleaved 32x32 tile. + static int const kElementsPerPartial = 4; + using EleShapePerPatial = typename platform::conditional< + platform::is_same::value, + MatrixShape<2, 2>, + MatrixShape<1, 4> >::type; + static int const kElementsPerMma = 8; + static int const kAccumulatorPatials = 2; + using QuadShapePerPatialMma = MatrixShape<4, 4>; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + +private: + + /// Reference to output tensor + TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpAccumulatorTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpAccumulatorTileIterator( + TensorRef const &ref, + int lane_id + ): + ref_(ref) { + + int quad = (lane_id >> 2); + int lane_in_quad = (lane_id & 3); + int accum_m, accum_n; + + if (platform::is_same::value) { + // (quad[2],quad[0])+lane_in_quad[0] + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + (lane_in_quad & 1); + // (quad[1])+lane_in_quad[1] + accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials + + (lane_in_quad & 2); + } else { + accum_m = (((quad & 0x4) >> 1) + (quad & 0x1)) * 8 + lane_in_quad; // (quad[2],quad[0]) + accum_n = ((quad >> 1) & 0x1) * kElementsPerPartial * kAccumulatorPatials; + } + MatrixCoord lane_offset(accum_m, accum_n); + + ref_.add_coord_offset(lane_offset); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpAccumulatorTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpAccumulatorTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + ref_.add_coord_offset(tile_offset * make_Coord(Shape::kRow, Shape::kColumn)); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpAccumulatorTileIterator & operator++() { + // deliberate no-op + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpAccumulatorTileIterator & operator--() { + // deliberate no-op + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpAccumulatorTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpAccumulatorTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_HOST_DEVICE + void load_with_pointer_offset( + Fragment &frag, ///< fragment to load from the tensor + Index pointer_offset) const { ///< loads a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; ++tile_n) { + MCTLASS_PRAGMA_UNROLL + for (int tile_m = 0; tile_m < Policy::TileIterations::kRow; ++tile_m) { + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = + (((tile_n * Policy::TileIterations::kRow + tile_m) * + Policy::MmaIterations::kColumn + mma_n) * + Policy::MmaIterations::kRow + mma_m) * + kElementsPerMma; + + MCTLASS_PRAGMA_UNROLL + for (int p = 0; p < kAccumulatorPatials; ++p) { + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < EleShapePerPatial::kRow; ++m) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < EleShapePerPatial::kColumn; ++n) { + int accum_m = tile_m * Policy::InterleavedTile::kRow + + mma_m * QuadShapePerPatialMma::kRow + m * 2; + int accum_n = tile_n * Policy::InterleavedTile::kColumn + + mma_n * QuadShapePerPatialMma::kColumn + + p * Policy::InterleavedTile::kColumn/2 + n; + int idx = mma_accum_start + p * kElementsPerPartial + + m * EleShapePerPatial::kColumn + n; + frag[idx] = offset_ref.at({accum_m, accum_n}); + } + } + } + } + } + } + } + } + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + Fragment &frag, ///< fragment to load from the tensor + Index byte_offset) const { ///< loads a tile with a linear offset + + load_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_HOST_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset) const { ///< loads a tile with a logical offset in units of whole tiles + + load(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_HOST_DEVICE + void load( + Fragment &frag, ///< fragment to load from the tensor + TensorCoord const &tile_offset, ///< loads a tile with a logical offset in units of whole tiles + Index pointer_offset) const { ///< loads a tile with a logical offset AND a pointer offset + + load_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } + + /// Stores a fragment to memory + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_HOST_DEVICE + void store_with_pointer_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index pointer_offset) const { ///< store a tile with a linear offset + + TensorRef offset_ref(ref_); + offset_ref.add_pointer_offset(pointer_offset); + + MCTLASS_PRAGMA_UNROLL + for (int tile_n = 0; tile_n < Policy::TileIterations::kColumn; ++tile_n) { + MCTLASS_PRAGMA_UNROLL + for (int tile_m = 0; tile_m < Policy::TileIterations::kRow; ++tile_m) { + MCTLASS_PRAGMA_UNROLL + for (int mma_n = 0; mma_n < Policy::MmaIterations::kColumn; ++mma_n) { + MCTLASS_PRAGMA_UNROLL + for (int mma_m = 0; mma_m < Policy::MmaIterations::kRow; ++mma_m) { + + int mma_accum_start = + (((tile_n * Policy::TileIterations::kRow + tile_m) * + Policy::MmaIterations::kColumn + mma_n) * + Policy::MmaIterations::kRow + mma_m) * + kElementsPerMma; + + MCTLASS_PRAGMA_UNROLL + for (int p = 0; p < kAccumulatorPatials; ++p) { + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < EleShapePerPatial::kRow; ++m) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < EleShapePerPatial::kColumn; ++n) { + int accum_m = tile_m * Policy::InterleavedTile::kRow + + mma_m * QuadShapePerPatialMma::kRow + m * 2; + int accum_n = tile_n * Policy::InterleavedTile::kColumn + + mma_n * QuadShapePerPatialMma::kColumn + + p * Policy::InterleavedTile::kColumn/2 + n; + int idx = mma_accum_start + p * kElementsPerPartial + + m * EleShapePerPatial::kColumn + n; + offset_ref.at({accum_m, accum_n}) = frag[idx]; + } + } + } + } + } + } + } + } + + /// Stores a fragment to memory with additional pointer offset + MCTLASS_HOST_DEVICE + void store_with_byte_offset( + Fragment const &frag, ///< fragment to store from the tensor + Index byte_offset) const { ///< store a tile with a linear offset + + store_with_pointer_offset(byte_offset / sizeof(Element)); + } + + /// Stores a fragment to memory with logical offset in units of whole tiles. + MCTLASS_HOST_DEVICE + void store( + Fragment &frag, ///< fragment to store to the tensor + TensorCoord const &tile_offset) const { ///< stores a tile with a logical offset in units of whole tiles + + store(frag, tile_offset, 0); + } + + /// Stores a fragment from memory with logical offset in units of whole tiles. + MCTLASS_HOST_DEVICE + void store( + /// fragment to store to the tensor + Fragment const &frag, + /// stores a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// stores a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + store_with_pointer_offset(frag, ref_.offset(tile_offset) + pointer_offset); + } +}; + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDS to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// KBlock size (in units of elements) + int KBlock> +class MmaVoltaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::VoltaTensorOpMultiplicandCrosswise< + sizeof_bits::value, KBlock>, + InstructionShape_, OpDelta_, 32> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand == Operand::kB, + "MmaVoltaTensorOpMultiplicandIterator may only be instantiated for " + "A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// KBlock size + static int const kKBlock = KBlock; + + /// Layout of source tile + using Layout = mctlass::layout::VoltaTensorOpMultiplicandCrosswise< + sizeof_bits::value, kKBlock>; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + + /// Shape of one individual LDS instruction + using LdsShape = layout::PitchLinearShape<1, 32>; + + /// Number and arrangement of LDSM instructions + using LdsIterations = layout::PitchLinearShape<1, Shape::kStrided / 32>; + + /// Using LDS.128 + static int const kElementsPerAccess = 8; + + /// Contiguous elements per line + static int const kContiguousElementsPerLine = 4; + }; + + private: + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = AlignedArray; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + + private: + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + /// Crosswised elements are arranged in a SMEM line + /// in units of AccessType + Index line_size; + + /// Internal counter used to determine load addr offset + /// and when to swap higher 64bit with lower 64bit + int k_group_idx_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator() + : pointer_(nullptr), + stride_(0), + line_size(0), + byte_offset_(0), + k_group_idx_(0) {} + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : pointer_(reinterpret_cast(ref.data())), + stride_(ref.stride(0) * Policy::kElementsPerAccess), + line_size((ref.stride(0) * Policy::kContiguousElementsPerLine) / + Policy::kElementsPerAccess), + k_group_idx_(0), + byte_offset_(0) { + + int quad = (lane_id / 4); + int lane_in_quad = (lane_id % 4); + int access_contiguous; + + if(kOperand == Operand::kA) { + + // swizzle id: tid[4]|tid[1:0]|(tid[2]^tid[4]) + access_contiguous = ((quad & 0x4) << 1) + ((lane_in_quad) << 1) + + ((quad & 0x1) ^ ((quad & 0x4) >> 2)); + } else { + + // swizzle id: tid[4]|tid[1:0]|tid[3] + access_contiguous = ((quad & 0x4) << 1) + (lane_in_quad << 1) + + ((quad & 0x2) >> 1 ^ ((quad & 0x4) >> 2)); + } + + byte_offset_ = access_contiguous * + sizeof(Element) * Policy::kElementsPerAccess; + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + byte_offset_ += offset * sizeof(Element); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + + int contiguous_offset = tile_offset.contiguous(); + int strided_offset = tile_offset.strided(); + k_group_idx_ = 0; + + pointer_ += contiguous_offset * + (InstructionShape::kContiguous / + Policy::kContiguousElementsPerLine) * + line_size + + strided_offset * Shape::kStrided / 2; + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &operator++() { + k_group_idx_ = (k_group_idx_ + 1) % 8; + + if (k_group_idx_ == 4 || k_group_idx_ == 0) { + byte_offset_ ^= 1 * sizeof(Element) * Policy::kElementsPerAccess; + } + + pointer_ += line_size; + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &operator--() { assert(0); } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { load_with_byte_offset(frag, 0); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + AccessType * fetch_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsIterations::kStrided; ++s) { + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsIterations::kContiguous; ++c) { + + int access_idx = c + s * Policy::LdsIterations::kContiguous; + + AccessType const *source_ptr = pointer_ + + Policy::LdsShape::kContiguous * c * line_size + + Policy::LdsShape::kStrided * s / 2; + + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_; + fetch_ptr[access_idx] = *(reinterpret_cast (source_byte_ptr)); + + // swap higher 64bit and lower 64bit + if (k_group_idx_ & 0x2) { + uint64_t *low = reinterpret_cast(&frag) + access_idx * 2; + uint64_t *high = reinterpret_cast(&frag) + access_idx * 2 + 1; + uint64_t tmp = *low; + *low = *high; + *high = tmp; + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = tile_offset.contiguous() * + InstructionShape::kContiguous / + Policy::kElementsPerAccess + + tile_offset.strided() * Shape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + k_group_idx_ = k_group; + } +}; + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDS to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// KBlock size (in units of elements) + int KBlock> +class MmaVoltaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::ColumnMajorVoltaTensorOpMultiplicandCrosswise< + sizeof_bits::value, KBlock>, + InstructionShape_, OpDelta_, 32> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for " + "A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// KBlock size + static int const kKBlock = KBlock; + + + /// Layout of source tile + using Layout = mctlass::layout::ColumnMajorVoltaTensorOpMultiplicandCrosswise< + sizeof_bits::value, kKBlock>; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaVoltaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::VoltaTensorOpMultiplicandCrosswise::value, + kKBlock>, + layout::PitchLinearShape, + kOpDelta, kThreads>; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + + private: + /// Underlying tile iterator + Base iterator_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator() {} + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : iterator_({ref.data(), ref.stride()}, lane_id) {} + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &operator++() { + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &operator--() { + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { iterator_.load(frag); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, {tile_offset.contiguous(), tile_offset.strided()}, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDS to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// KBlock size (in units of elements) + int KBlock> +class MmaVoltaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::RowMajorVoltaTensorOpMultiplicandCrosswise< + sizeof_bits::value, KBlock>, + InstructionShape_, OpDelta_, 32> { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand == Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for " + "A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// KBlock size + static int const kKBlock = KBlock; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajorVoltaTensorOpMultiplicandCrosswise< + sizeof_bits::value, kKBlock>; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaVoltaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::VoltaTensorOpMultiplicandCrosswise::value, + kKBlock>, + layout::PitchLinearShape, + kOpDelta, kThreads>; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + + private: + /// Underlying tile iterator + Base iterator_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator() {} + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) + : iterator_({ref.data(), ref.stride()}, lane_id) {} + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &operator++() { + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &operator--() { + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { iterator_.load(frag); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, {tile_offset.strided(), tile_offset.contiguous()}, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for 'TN' arrangement +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Operand identity + Operand Operand_, + /// Data type of A elements + typename Element_, + /// Layout of matrix operand + typename Layout_, + /// Shape of one matrix production operation (concept: MatrixShape) + typename InstructionShape_, + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + int OpDelta_, + /// Number of threads participating in one matrix operation + int Threads = 32, + /// Number of partitions along K dimension + int PartitionsK_ = 1> +class MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + /// Basic check + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaVoltaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = Layout_; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Number of elements accessed per Shared Memory load + static int const kElementsPerAccess = 4; + +private: + + static int const kInterleavedTileRows = 32; + static int const kInterleavedTileColumns = 32; + static int const kInstructionsPerTile = 2; + + /// Rounded up instruction counts + using TileCount = MatrixShape< + Shape::kRow / kInterleavedTileRows, + Shape::kColumn / kInterleavedTileColumns + >; + + using FragmentCount = MatrixShape< + TileCount::kRow * kInstructionsPerTile, + TileCount::kColumn * kInstructionsPerTile + >; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array< + Element, + (kOperand == Operand::kA ? FragmentCount::kRow : FragmentCount::kColumn) * kElementsPerAccess + >; + + /// Memory access type + using AccessType = AlignedArray; + +private: + + /// Underlying tensor reference + TensorRef ref_; + + /// Extent of tensor + MatrixCoord extent_; + + /// Origin + MatrixCoord origin_; + + /// Used to conditionally enable extents checking + bool divisible_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner(): divisible_(true) { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner( + TensorRef const &ref, + int lane_id + ): + ref_(ref), extent_(Shape::kRow, Shape::kColumn), divisible_(true) { + + int quad_id = lane_id / 4; + int lane_in_quad = (lane_id % 4); + + if (kOperand == Operand::kA) { + + int row_idx = ((quad_id & 1) + ((quad_id & 4) / 2)) * 4 * kInstructionsPerTile + lane_in_quad; + int col_idx = 0; + + origin_ = MatrixCoord(row_idx, col_idx); + } + else { + + int row_idx = 0; + int col_idx = (quad_id / 2) * 4 * kInstructionsPerTile + lane_in_quad; + + origin_ = MatrixCoord(row_idx, col_idx); + } + + ref_.add_coord_offset(origin_); + } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner( + TensorRef const &ref, + TensorCoord extent, + int lane_id + ): ref_(ref), extent_(extent), divisible_(false) { + + int quad_id = lane_id / 4; + int lane_in_quad = (lane_id % 4); + + if (kOperand == Operand::kA) { + + int row_idx = ((quad_id & 1) + ((quad_id & 4) / 2)) * 4 * kInstructionsPerTile + lane_in_quad; + int col_idx = 0; + + origin_ = MatrixCoord(row_idx, col_idx); + } + else { + + int row_idx = 0; + int col_idx = (quad_id / 2) * 4 * kInstructionsPerTile + lane_in_quad; + + origin_ = MatrixCoord(row_idx, col_idx); + } + + #if defined(__MACA_ARCH__) + __syncthreads(); + #endif + + ref_.add_coord_offset(origin_); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner &add_pointer_offset(LongIndex offset) { + + ref_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner &add_tile_offset(TensorCoord const &tile_offset) { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + origin_ += coord_offset; + + ref_.add_coord_offset(coord_offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner & operator++() { + + if (kOperand == Operand::kA) { + add_tile_offset({0, 1}); + } + else { + add_tile_offset({1, 0}); + } + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner & operator--() { + + if (kOperand == Operand::kA) { + add_tile_offset({0, -1}); + } + else { + add_tile_offset({-1, 0}); + } + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + AccessType const *access_ptr = reinterpret_cast(ref_.data()); + int ldm = ref_.stride()[0]; + + if (kOperand == Operand::kA) { + + MCTLASS_PRAGMA_UNROLL + for (int idx = 0; idx < FragmentCount::kRow; ++idx) { + + int tile_idx = idx / 2; + int quad_idx = idx % 2; + + int row_offset = tile_idx * kInterleavedTileRows + quad_idx * 4; + frag_ptr[idx] = access_ptr[row_offset * ldm / kElementsPerAccess]; + } + } + else { + MCTLASS_PRAGMA_UNROLL + for (int idx = 0; idx < FragmentCount::kColumn; ++idx) { + + int tile_idx = idx / 2; + int quad_idx = idx % 2; + + int col_offset = tile_idx * kInterleavedTileColumns + quad_idx * 4; + frag_ptr[idx] = access_ptr[col_offset * ldm / kElementsPerAccess]; + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + + load_with_pointer_offset(frag, byte_offset * 8 / sizeof_bits::value); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + + load_with_pointer_offset(frag, ref_.offset(coord_offset)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + + load_with_pointer_offset(frag, ref_.offset(coord_offset) + pointer_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + + load_with_pointer_offset(frag, ref_.offset(coord_offset) + byte_offset * 8 / sizeof_bits::value); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation + } +}; + + +/// Tile iterator specialized for 'NT' arrangement +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Operand identity + Operand Operand_, + /// Data type of A elements + typename Element_, + /// Layout of matrix operand + typename Layout_, + /// Shape of one matrix production operation (concept: MatrixShape) + typename InstructionShape_, + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + int OpDelta_, + /// Number of threads participating in one matrix operation + int Threads = 32, + /// Number of partitions along K dimension + int PartitionsK_ = 1> +class MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + /// Basic check + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaVoltaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = Layout_; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Number of elements accessed per Shared Memory load + static int const kElementsPerAccess = 4; + +private: + + static int const kInterleavedTileRows = 32; + static int const kInterleavedTileColumns = 32; + static int const kInstructionsPerTile = 2; + + /// Rounded up instruction counts + using TileCount = MatrixShape< + Shape::kRow / kInterleavedTileRows, + Shape::kColumn / kInterleavedTileColumns + >; + + using FragmentCount = MatrixShape< + TileCount::kRow * kInstructionsPerTile, + TileCount::kColumn * kInstructionsPerTile + >; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array< + Element, + (kOperand == Operand::kA ? FragmentCount::kRow : FragmentCount::kColumn) * kElementsPerAccess + >; + + /// Memory access type + using AccessType = AlignedArray; + +private: + + /// Underlying tensor reference + TensorRef ref_; + + /// Extent of tensor + MatrixCoord extent_; + + /// Origin + MatrixCoord origin_; + + /// Used to conditionally enable extents checking + bool divisible_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter(): divisible_(true) { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter( + TensorRef const &ref, + int lane_id + ): + ref_(ref), extent_(Shape::kRow, Shape::kColumn), divisible_(true) { + + int quad_id = lane_id / 4; + int lane_in_quad = (lane_id % 4); + + if (kOperand == Operand::kA) { + + int row_idx = ((quad_id & 1) + ((quad_id & 4) / 2)) * 4 * kInstructionsPerTile; + int col_idx = lane_in_quad; + + origin_ = MatrixCoord(row_idx, col_idx); + } + else { + + int row_idx = lane_in_quad; + int col_idx = (quad_id / 2) * 4 * kInstructionsPerTile; + + origin_ = MatrixCoord(row_idx, col_idx); + } + + ref_.add_coord_offset(origin_); + } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter( + TensorRef const &ref, + TensorCoord extent, + int lane_id + ): ref_(ref), extent_(extent), divisible_(false) { + + int quad_id = lane_id / 4; + int lane_in_quad = (lane_id % 4); + + if (kOperand == Operand::kA) { + + int row_idx = ((quad_id & 1) + ((quad_id & 4) / 2)) * 4 * kInstructionsPerTile; + int col_idx = lane_in_quad; + + origin_ = MatrixCoord(row_idx, col_idx); + } + else { + + int row_idx = lane_in_quad; + int col_idx = (quad_id / 2) * 4 * kInstructionsPerTile; + + origin_ = MatrixCoord(row_idx, col_idx); + } + + #if defined(__MACA_ARCH__) + __syncthreads(); + #endif + + ref_.add_coord_offset(origin_); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter &add_pointer_offset(LongIndex offset) { + + ref_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter &add_tile_offset(TensorCoord const &tile_offset) { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + origin_ += coord_offset; + + ref_.add_coord_offset(coord_offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter & operator++() { + + if (kOperand == Operand::kA) { + add_tile_offset({0, 1}); + } + else { + add_tile_offset({1, 0}); + } + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter & operator--() { + + if (kOperand == Operand::kA) { + add_tile_offset({0, -1}); + } + else { + add_tile_offset({-1, 0}); + } + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + AccessType const *access_ptr = reinterpret_cast(ref_.data()); + int ldm = ref_.stride()[0]; + + if (kOperand == Operand::kA) { + + MCTLASS_PRAGMA_UNROLL + for (int idx = 0; idx < FragmentCount::kRow; ++idx) { + + int tile_idx = idx / 2; + int quad_idx = idx % 2; + + int row_offset = tile_idx * kInterleavedTileRows; + frag_ptr[idx] = access_ptr[row_offset / kElementsPerAccess + quad_idx]; + } + } + else { + MCTLASS_PRAGMA_UNROLL + for (int idx = 0; idx < FragmentCount::kColumn; ++idx) { + + int tile_idx = idx / 2; + int quad_idx = idx % 2; + + int col_offset = tile_idx * kInterleavedTileColumns; + frag_ptr[idx] = access_ptr[col_offset / kElementsPerAccess + quad_idx]; + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + + load_with_pointer_offset(frag, byte_offset * 8 / sizeof_bits::value); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + + load_with_pointer_offset(frag, ref_.offset(coord_offset)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + + load_with_pointer_offset(frag, ref_.offset(coord_offset) + pointer_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + + load_with_pointer_offset(frag, ref_.offset(coord_offset) + byte_offset * 8 / sizeof_bits::value); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_> +class MmaVoltaTensorOpMultiplicandTileIterator< + Shape_, + Operand::kA, + Element_, + mctlass::layout::RowMajor, + InstructionShape_, + OpDelta_, + 32 +> : public MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner< + Shape_, Operand::kA, Element_, mctlass::layout::RowMajor, InstructionShape_, OpDelta_> { + +public: + using Base = MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner< + Shape_, Operand::kA, Element_, mctlass::layout::RowMajor, InstructionShape_, OpDelta_> ; + + using TensorRef = typename Base::TensorRef; + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): Base(ref, lane_id) { } + +}; + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_> +class MmaVoltaTensorOpMultiplicandTileIterator< + Shape_, + Operand::kA, + Element_, + mctlass::layout::ColumnMajor, + InstructionShape_, + OpDelta_, + 32 +> : public MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter< + Shape_, Operand::kA, Element_, mctlass::layout::ColumnMajor, InstructionShape_, OpDelta_> { + +public: + using Base = MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter< + Shape_, Operand::kA, Element_, mctlass::layout::ColumnMajor, InstructionShape_, OpDelta_> ; + + using TensorRef = typename Base::TensorRef; + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): Base(ref, lane_id) { } + +}; + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_> +class MmaVoltaTensorOpMultiplicandTileIterator< + Shape_, Operand::kB, Element_, + mctlass::layout::ColumnMajor, + InstructionShape_, OpDelta_, 32 +> : public MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner< + Shape_, Operand::kB, Element_, mctlass::layout::ColumnMajor, InstructionShape_, OpDelta_> { + +public: + using Base = MmaVoltaTensorOpMultiplicandTileIteratorCanonicalInner< + Shape_, Operand::kB, Element_, mctlass::layout::ColumnMajor, InstructionShape_, OpDelta_>; + + using TensorRef = typename Base::TensorRef; + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): Base(ref, lane_id) { } +}; + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_> +class MmaVoltaTensorOpMultiplicandTileIterator< + Shape_, Operand::kB, Element_, + mctlass::layout::RowMajor, + InstructionShape_, OpDelta_, 32 +> : public MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter< + Shape_, Operand::kB, Element_, mctlass::layout::RowMajor, InstructionShape_, OpDelta_> { + +public: + using Base = MmaVoltaTensorOpMultiplicandTileIteratorCanonicalOuter< + Shape_, Operand::kB, Element_, mctlass::layout::RowMajor, InstructionShape_, OpDelta_>; + + using TensorRef = typename Base::TensorRef; + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaVoltaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): Base(ref, lane_id) { } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h new file mode 100644 index 0000000..0779652 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h @@ -0,0 +1,2946 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines iterators used by warp-level matrix multiply operations targeting Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/gemm/gemm.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor_op_multiplicand_sm80.h" + +#include "mctlass/platform/platform.h" +#include "mctlass/fast_math.h" + +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for loading 128b vectors of 64b elements. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::TensorOpMultiplicandCongruous64b, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + static_assert(!(Shape::kContiguous % 16) && !(Shape::kStrided % 4), "Divisibility."); + + static_assert(sizeof_bits::value == 64, "This is specialized for 64b accesses."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicandCongruous64b; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Load two elements per access + static int const kElementsPerAccess = 2; + + /// Policy defining internal details of tile iterator + struct Policy { + + /// Shape of one access + using Delta = layout::PitchLinearShape<8, 4>; + + /// Number of iterations to load + using Iterations = layout::PitchLinearShape< + Shape::kContiguous / kElementsPerAccess / Delta::kContiguous, + InstructionShape::kStrided / Delta::kStrided + >; + + }; + +private: + + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = AlignedArray; + + /// Internal counter used to jump to next K partition + int k_group_idx_; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + +private: + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(): stride_(0), byte_offset_(0) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): + stride_(ref.stride(0) / kElementsPerAccess), byte_offset_(0), + k_group_idx_(0) { + + int access_strided = lane_id / Policy::Delta::kContiguous; + int access_contiguous = (lane_id % Policy::Delta::kContiguous) ^ access_strided; + + pointer_= reinterpret_cast(ref.data()) + + access_contiguous + access_strided * stride_; + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + byte_offset_ += offset * sizeof(Element); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + int offset = + (tile_offset.strided() * InstructionShape::kStrided) * stride_ * kElementsPerAccess + + tile_offset.contiguous() * Shape::kContiguous; + + add_pointer_offset(offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + add_tile_offset({0, 1}); + + return *this; + } + + /// Advances the iterator along the opposite of the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + add_tile_offset({0, -1}); + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_byte_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + AccessType *fetch_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::Iterations::kStrided; ++s) { + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::Iterations::kContiguous; ++c) { + + int access_idx = c + s * Policy::Iterations::kContiguous; + + AccessType const *source_ptr = pointer_ + + Policy::Delta::kContiguous * c + + Policy::Delta::kStrided * s * stride_; + + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_; + + AccessType const *source = reinterpret_cast(source_byte_ptr); + + fetch_ptr[access_idx] = *source; + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + + Index pointer_offset = + tile_offset.contiguous() * Shape::kContiguous / Layout::kElementsPerAccess + + tile_offset.strided() * InstructionShape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::RowMajorTensorOpMultiplicandCongruous64b, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajorTensorOpMultiplicandCongruous64b; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCongruous64b, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) { + int t_id = lane_id; + int row = t_id & 0x7; + int col = (t_id >> 4) & 0x3; + t_id = row * 4 + col; + iterator_ = Base({ref.data(), ref.stride()}, t_id); + } + // MCTLASS_HOST_DEVICE + // MmaTensorOpMultiplicandTileIterator( + // TensorRef const &ref, + // int lane_id + // ): iterator_({ref.data(), ref.stride()}, lane_id) { + // } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.strided(), tile_offset.contiguous()}, + byte_offset); + } + + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to load from shared +/// memory and therefore must be initialized with a TensorRef to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::ColumnMajorTensorOpMultiplicandCongruous64b, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::ColumnMajorTensorOpMultiplicandCongruous64b; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicandCongruous64b, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) { + int t_id = lane_id; + int row = t_id & 0x7; + int col = (t_id >> 4) & 0x3; + t_id = row * 4 + col; + iterator_ = Base({ref.data(), ref.stride()}, t_id); + } + // MCTLASS_HOST_DEVICE + // MmaTensorOpMultiplicandTileIterator( + // TensorRef const &ref, + // int lane_id + // ): iterator_({ref.data(), ref.stride()}, lane_id) { + // } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.contiguous(), tile_offset.strided()}, + byte_offset); + } + + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for loading 128b vectors of 64b elements. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::TensorOpMultiplicand64bCrosswise, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + static_assert(!(Shape::kContiguous % 4) && !(Shape::kStrided % 16), "Divisibility."); + + static_assert(sizeof_bits::value == 64, "This is specialized for 64b accesses."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::TensorOpMultiplicand64bCrosswise; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Long Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Load two elements per access + static int const kElementsPerAccess = 2; + + /// Policy defining internal details of tile iterator + struct Policy { + + /// Shape of one access + using Delta = layout::PitchLinearShape<4, 16>; + + /// Number of iterations to load + using Iterations = layout::PitchLinearShape< + InstructionShape::kContiguous / Delta::kContiguous, + Shape::kStrided / Delta::kStrided + >; + + }; + +private: + + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = AlignedArray; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + +private: + + /// Layout object storing stride values + StrideIndex stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + /// Internal counter for tracking K-group + Index k_group_idx_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(): stride_(0), byte_offset_(0) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): + stride_(ref.stride(0) / kElementsPerAccess), byte_offset_(0), + k_group_idx_(0) { + + int access_strided = lane_id / 8; + int access_contiguous = (lane_id % 8); + + byte_offset_ = (access_contiguous + access_strided * stride_) * sizeof(AccessType); + + pointer_= reinterpret_cast(ref.data()); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + pointer_ += offset / kElementsPerAccess; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + int offset = (tile_offset.contiguous() * InstructionShape::kContiguous) * + stride_ * kElementsPerAccess + + tile_offset.strided() * Shape::kStrided; + + add_pointer_offset(offset); + + int old_k_group_idx = k_group_idx_; + + k_group_idx_ += tile_offset.contiguous(); + + if ((k_group_idx_ & 2) ^ (old_k_group_idx & 2)) { + byte_offset_ ^= 0x40; + } + + return *this; + } + + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative(TensorCoord const &tile_offset) { + + add_tile_offset(tile_offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + pointer_ += stride_ * InstructionShape::kContiguous; + + if (k_group_idx_ & 0x1) { + // xor ptr + byte_offset_ ^= 0x40; + } + + ++k_group_idx_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_byte_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + AccessType *fetch_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::Iterations::kContiguous; ++c) { + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::Iterations::kStrided; ++s) { + + int access_idx = c + s * Policy::Iterations::kContiguous; + + AccessType const *source_ptr = pointer_ + + Policy::Delta::kContiguous * c * stride_ + + Policy::Delta::kStrided * s / kElementsPerAccess; + + char const *source_byte_ptr = reinterpret_cast(source_ptr) + byte_offset + byte_offset_; + + AccessType const *source = reinterpret_cast(source_byte_ptr); + + fetch_ptr[access_idx] = *source; + } + } + + Element *exchange_ptr = reinterpret_cast(&frag); + + if (k_group_idx_ & 1) { + // exchange on 64b granularity + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Fragment::kElements; i += 2) { + Element tmp = exchange_ptr[i]; + exchange_ptr[i] = exchange_ptr[i + 1]; + exchange_ptr[i + 1] = tmp; + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = tile_offset.contiguous() * + InstructionShape::kContiguous / + Layout::kElementsPerAccess + + tile_offset.strided() * Shape::kStrided * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + k_group_idx_ = k_group; + } +}; + +//////////////////////////////////////////////////////////////////////////////// +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::RowMajorTensorOpMultiplicand64bCrosswise, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajorTensorOpMultiplicand64bCrosswise; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicand64bCrosswise, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) { + int t_id = lane_id; + int row = t_id & 0x7; + int col = (t_id >> 4) & 0x3; + t_id = row * 4 + col; + iterator_ = Base({ref.data(), ref.stride()}, t_id); + } + // MCTLASS_HOST_DEVICE + // MmaTensorOpMultiplicandTileIterator( + // TensorRef const &ref, + // int lane_id + // ): iterator_({ref.data(), ref.stride()}, lane_id) { + // } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset_negative({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.strided(), tile_offset.contiguous()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::ColumnMajorTensorOpMultiplicand64bCrosswise, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::ColumnMajorTensorOpMultiplicand64bCrosswise; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIterator< + layout::PitchLinearShape, kOperand, Element, + layout::TensorOpMultiplicand64bCrosswise, + layout::PitchLinearShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) { + int t_id = lane_id; + int row = t_id & 0x7; + int col = (t_id >> 4) & 0x3; + t_id = row * 4 + col; + iterator_ = Base({ref.data(), ref.stride()}, t_id); + } + // MCTLASS_HOST_DEVICE + // MmaTensorOpMultiplicandTileIterator( + // TensorRef const &ref, + // int lane_id + // ): iterator_({ref.data(), ref.stride()}, lane_id) { + // } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset_negative(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset_negative({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.contiguous(), tile_offset.strided()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + + +/// Tile iterator specialized for canonical matrix layouts +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Operand identity + Operand Operand_, + /// Data type of A elements + typename Element_, + /// Layout of operand + typename Layout_, + /// Shape of one matrix production operation (concept: MatrixShape) + typename InstructionShape_, + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + int OpDelta_, + /// Number of threads participating in one matrix operation + int Threads = 32, + /// Number of partitions along K dimension + int PartitionsK_ = 1> +class MmaTensorOpMultiplicandTileIteratorCanonical { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + /// Basic check + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = Layout_; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = Threads; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Number of elements accessed per Shared Memory load + static int const kElementsPerAccess = + (sizeof_bits::value >= 32 ? 1 : 32 / sizeof_bits::value); + +private: + + static int const kWarpShapeOuter = + (kOperand == Operand::kA ? Shape::kRow : Shape::kColumn); + + static int const kWarpShapeInner = + (kOperand == Operand::kA ? Shape::kColumn : Shape::kRow); + + + /// Rounded up instruction counts + using InstructionCount = MatrixShape< + Shape::kRow / InstructionShape::kRow, + Shape::kColumn / InstructionShape::kColumn + >; + + /// Rounded up tile dimensions + using WarpShapeDivisible = MatrixShape< + InstructionCount::kRow * InstructionShape::kRow, + InstructionCount::kColumn * InstructionShape::kColumn + >; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array< + Element, + WarpShapeDivisible::kRow * WarpShapeDivisible::kColumn / kThreads + >; + + /// Memory access type + using AccessType = AlignedArray; + +private: + + /// Underlying tensor reference + TensorRef ref_; + + /// Extent of tensor + MatrixCoord extent_; + + /// Origin + MatrixCoord origin_; + + /// Used to conditionally enable extents checking + bool divisible_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical(): divisible_(true) { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical( + TensorRef const &ref, + int lane_id + ): ref_(ref), extent_(Shape::kRow, Shape::kColumn), divisible_(true) { + + if (mctlass::platform::is_same::value) { //Only for m16n16k16 + if (kOperand == Operand::kA) { + origin_ = MatrixCoord(lane_id & 0xf, ((lane_id >> 4) << 2) ^ 0xf); + } + else { + origin_ = MatrixCoord(((lane_id >> 4) << 2) ^ 0xf, lane_id & 0xf); + } + } + else { + if (kOperand == Operand::kA) { + origin_ = MatrixCoord(lane_id / 4, (lane_id % 4) * kElementsPerAccess); + } + else { + origin_ = MatrixCoord((lane_id % 4) * kElementsPerAccess, lane_id / 4); + } + } + + ref_.add_coord_offset(origin_); + } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical( + TensorRef const &ref, + TensorCoord extent, + int lane_id + ): ref_(ref), extent_(extent), divisible_(false) { + + if (kOperand == Operand::kA) { + origin_ = MatrixCoord(lane_id / 4, (lane_id % 4) * kElementsPerAccess); + } + else { + origin_ = MatrixCoord((lane_id % 4) * kElementsPerAccess, lane_id / 4); + } + + ref_.add_coord_offset(origin_); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical &add_pointer_offset(LongIndex offset) { + + ref_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical &add_tile_offset(TensorCoord const &tile_offset) { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + origin_ += coord_offset; + + ref_.add_coord_offset(coord_offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical & operator++() { + + if (kOperand == Operand::kA) { + add_tile_offset({0, 1}); + } + else { + add_tile_offset({1, 0}); + } + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical & operator--() { + + if (kOperand == Operand::kA) { + add_tile_offset({0, -1}); + } + else { + add_tile_offset({-1, 0}); + } + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + + int const kWarpShapeDivisibleInner = + (kOperand == Operand::kA ? WarpShapeDivisible::kColumn : WarpShapeDivisible::kRow); + + // Take advantage of Tensor Op's 8 x 4T access pattern + int const kAccessesInner = (kWarpShapeDivisibleInner / kElementsPerAccess) / 4; + + AccessType *access_ptr = reinterpret_cast(&frag); + + if (kOperand == Operand::kA) { + if (mctlass::platform::is_same::value) { + + int const kTilesPerInstruction = InstructionShape::kRow / 16; + + MCTLASS_PRAGMA_UNROLL + for (int inst_m_idx = 0; inst_m_idx < InstructionCount::kRow; ++inst_m_idx) { + + MCTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { + + int access_idx = inner_idx + inst_m_idx * kAccessesInner; + for (int idx = 0; idx < kElementsPerAccess; ++idx) { + MatrixCoord offset( + inst_m_idx * InstructionShape::kRow, + -1 * (inner_idx * kElementsPerAccess + idx)); + + MatrixCoord access_coord = origin_ + offset; + + if (divisible_ || + (access_coord.row() < extent_.row() && access_coord.column() < extent_.column())) { + + Element const *ref_ptr = reinterpret_cast(ref_.data() + ref_.offset(offset)); + access_ptr[access_idx][idx] = ref_ptr[0]; + } + else { + AccessType zero; + zero.clear(); + access_ptr[access_idx] = zero; + } + } + } + } + + } + else { + int const kTilesPerInstruction = InstructionShape::kRow / 8; + + MCTLASS_PRAGMA_UNROLL + for (int inst_m_idx = 0; inst_m_idx < InstructionCount::kRow; ++inst_m_idx) { + + MCTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { + + MCTLASS_PRAGMA_UNROLL + for (int access_m_idx = 0; access_m_idx < kTilesPerInstruction; ++access_m_idx) { + int access_idx = + access_m_idx + kTilesPerInstruction * (inner_idx + kAccessesInner * inst_m_idx); + + MatrixCoord offset( + access_m_idx * 8 + inst_m_idx * InstructionShape::kRow, + inner_idx * 4 * kElementsPerAccess); + + MatrixCoord access_coord = origin_ + offset; + + if (divisible_ || + (access_coord.row() < extent_.row() && access_coord.column() < extent_.column())) { + + access_ptr[access_idx] = *reinterpret_cast( + ref_.data() + ref_.offset(offset)); + } + else { + AccessType zero; + zero.clear(); + access_ptr[access_idx] = zero; + } + } + } + } + } + } + else { + if (mctlass::platform::is_same::value) { + MCTLASS_PRAGMA_UNROLL + for (int inst_n_idx = 0; inst_n_idx < InstructionCount::kColumn; ++inst_n_idx) { + MCTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { + int access_idx = inner_idx + kAccessesInner * inst_n_idx; + for (int idx = 0; idx < kElementsPerAccess; ++idx) { + MatrixCoord offset( + -1 * (inner_idx * kElementsPerAccess + idx), + inst_n_idx * InstructionShape::kColumn); + + MatrixCoord access_coord = origin_ + offset; + + if (divisible_ || + (access_coord.row() < extent_.row() && access_coord.column() < extent_.column())) { + + Element const *ref_ptr = reinterpret_cast(ref_.data() + ref_.offset(offset)); + access_ptr[access_idx][idx] = ref_ptr[0]; + } + else { + AccessType zero; + zero.clear(); + access_ptr[access_idx] = zero; + } + } + } + } + } + else { + MCTLASS_PRAGMA_UNROLL + for (int inst_n_idx = 0; inst_n_idx < InstructionCount::kColumn; ++inst_n_idx) { + + MCTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { + int access_idx = inner_idx + kAccessesInner * inst_n_idx; + + MatrixCoord offset( + inner_idx * 4 * kElementsPerAccess, + inst_n_idx * 8); + + MatrixCoord access_coord = origin_ + offset; + + if (divisible_ || + (access_coord.row() < extent_.row() && access_coord.column() < extent_.column())) { + + access_ptr[access_idx] = *reinterpret_cast( + ref_.data() + ref_.offset(offset)); + } + else { + AccessType zero; + zero.clear(); + access_ptr[access_idx] = zero; + } + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + + load_with_pointer_offset(frag, byte_offset * 8 / sizeof_bits::value); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + + load_with_pointer_offset(frag, ref_.offset(coord_offset)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + + load_with_pointer_offset(frag, ref_.offset(coord_offset) + pointer_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + + load_with_pointer_offset(frag, ref_.offset(coord_offset) + byte_offset * 8 / sizeof_bits::value); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation + } +}; + +/// Tile iterator specialized for canonical matrix layouts +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Operand identity + Operand Operand_, + /// Data type of A elements + typename Layout_, + /// Shape of one matrix production operation (concept: MatrixShape) + typename InstructionShape_, + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + int OpDelta_, + /// Number of threads participating in one matrix operation + int Threads, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIteratorCanonical { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + /// Basic check + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = mctlass::tfloat32_t; + + /// Layout of source tile + using Layout = Layout_; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Number of elements accessed per Shared Memory load + static int const kElementsPerAccess = + (sizeof_bits::value >= 32 ? 1 : 32 / sizeof_bits::value); + +private: + + static int const kWarpShapeOuter = + (kOperand == Operand::kA ? Shape::kRow : Shape::kColumn); + + static int const kWarpShapeInner = + (kOperand == Operand::kA ? Shape::kColumn : Shape::kRow); + + + /// Rounded up instruction counts + using InstructionCount = MatrixShape< + Shape::kRow / InstructionShape::kRow, + Shape::kColumn / InstructionShape::kColumn + >; + + /// Rounded up tile dimensions + using WarpShapeDivisible = MatrixShape< + InstructionCount::kRow * InstructionShape::kRow, + InstructionCount::kColumn * InstructionShape::kColumn + >; + +public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array< + Element, + WarpShapeDivisible::kRow * WarpShapeDivisible::kColumn / kThreads + >; + + /// Memory access type + using AccessType = AlignedArray; + +private: + + /// Underlying tensor reference + TensorRef ref_; + + /// Extent of tensor + MatrixCoord extent_; + + /// Origin + MatrixCoord origin_; + + /// Used to conditionally enable extents checking + bool divisible_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical(): divisible_(true) { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical( + TensorRef const &ref, + int lane_id + ): ref_(ref), extent_(Shape::kRow, Shape::kColumn), divisible_(true) { + int t_id = __lane_id(); + if (kOperand == Operand::kA) { + int row = t_id & 0xf; + int col = ((t_id >> 4) << 1) ^ 0x7; + int x = col; + int y = row; + origin_ = MatrixCoord(y, x); + } else { + int row = ((t_id >> 4) << 1) ^ 0x7; + int col = t_id & 0x7; + int y = col; + int x = row; + origin_ = MatrixCoord(x, y); + } + + ref_.add_coord_offset(origin_); + } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical( + TensorRef const &ref, + TensorCoord extent, + int lane_id + ): ref_(ref), extent_(extent), divisible_(false) { + + if (kOperand == Operand::kA) { + origin_ = MatrixCoord(lane_id / 4, (lane_id % 4) * kElementsPerAccess); + } + else { + origin_ = MatrixCoord((lane_id % 4) * kElementsPerAccess, lane_id / 4); + } + + ref_.add_coord_offset(origin_); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical &add_pointer_offset(LongIndex offset) { + + ref_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical &add_tile_offset(TensorCoord const &tile_offset) { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + origin_ += coord_offset; + + ref_.add_coord_offset(coord_offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical & operator++() { + + if (kOperand == Operand::kA) { + add_tile_offset({0, 1}); + } + else { + add_tile_offset({1, 0}); + } + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical & operator--() { + + if (kOperand == Operand::kA) { + add_tile_offset({0, -1}); + } + else { + add_tile_offset({-1, 0}); + } + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIteratorCanonical & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + load_with_pointer_offset(frag, 0); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + + int const kWarpShapeDivisibleInner = + (kOperand == Operand::kA ? WarpShapeDivisible::kColumn : WarpShapeDivisible::kRow); + + // Take advantage of Tensor Op's 8 x 4T access pattern + //int const kAccessesInner = (kWarpShapeDivisibleInner / kElementsPerAccess) / 4; + int const kAccessesInner = (kWarpShapeDivisibleInner / kElementsPerAccess) / 8; + + AccessType *access_ptr = reinterpret_cast(&frag); + + if (kOperand == Operand::kA) { + int const kTilesPerInstruction = InstructionShape::kRow / 8; + MCTLASS_PRAGMA_UNROLL + for (int inst_m_idx = 0; inst_m_idx < InstructionCount::kRow; ++inst_m_idx) { + MCTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner; ++inner_idx) { + MCTLASS_PRAGMA_UNROLL + for (int access_m_idx = 0; access_m_idx < kTilesPerInstruction; ++access_m_idx) { + int access_idx = access_m_idx + kTilesPerInstruction * + (inner_idx + 2 * kAccessesInner * inst_m_idx); + + MatrixCoord offset(inst_m_idx * InstructionShape::kRow, (-1 * access_m_idx)); + MatrixCoord access_coord = origin_ + offset; + if (divisible_ || + (access_coord.row() < extent_.row() && access_coord.column() < extent_.column())) { + + access_ptr[access_idx] = *reinterpret_cast( + ref_.data() + ref_.offset(offset)); + } + else { + AccessType zero; + zero.clear(); + access_ptr[access_idx] = zero; + } + } + } + } + } + else { + int kAccessesInner0 = 2; + MCTLASS_PRAGMA_UNROLL + for (int inst_n_idx = 0; inst_n_idx < InstructionCount::kColumn; ++inst_n_idx) { + + MCTLASS_PRAGMA_UNROLL + for (int inner_idx = 0; inner_idx < kAccessesInner0; ++inner_idx) { + int access_idx = inner_idx + kAccessesInner0 * inst_n_idx; + + MatrixCoord offset( + inner_idx * (-1), + inst_n_idx * 8); + + MatrixCoord access_coord = origin_ + offset; + + if (divisible_ || + (access_coord.row() < extent_.row() && access_coord.column() < extent_.column())) { + + access_ptr[access_idx] = *reinterpret_cast( + ref_.data() + ref_.offset(offset)); + } + else { + AccessType zero; + zero.clear(); + access_ptr[access_idx] = zero; + } + } + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + + load_with_pointer_offset(frag, byte_offset * 8 / sizeof_bits::value); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + + load_with_pointer_offset(frag, ref_.offset(coord_offset)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + + load_with_pointer_offset(frag, ref_.offset(coord_offset) + pointer_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + + TensorCoord coord_offset(tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn); + + load_with_pointer_offset(frag, ref_.offset(coord_offset) + byte_offset * 8 / sizeof_bits::value); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation + } +}; + +/// Wrapper for ColumnMajor +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::ColumnMajor, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::ColumnMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads, for half_t m16n16k16,we use kThreads=64 + static int const kThreads = mctlass::platform::is_same::value ? 64 : 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIteratorCanonical< + Shape, kOperand, Element, + layout::ColumnMajor, + InstructionShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) { + int t_id = lane_id; + int row = t_id & 0x7; + int col = (t_id >> 4) & 0x3; + t_id = row * 4 + col; + iterator_ = Base({ref.data(), ref.stride()}, t_id); + + if (mctlass::platform::is_same::value) { + iterator_ = Base({ref.data(), ref.stride()}, lane_id); + } + + } + // MCTLASS_HOST_DEVICE + // MmaTensorOpMultiplicandTileIterator( + // TensorRef const &ref, + // int lane_id + // ): iterator_({ref.data(), ref.stride()}, lane_id) { + // } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + TensorCoord const & extent, + int lane_id + ): iterator_({ref.data(), ref.stride()}, extent, lane_id) { + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.contiguous(), tile_offset.strided()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + + +/// Wrapper for RowMajor +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Identifies A or B multiplicand + Operand Operand_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Interval between adjacent *MMA instructions (in units of MMA + /// instructions) + int OpDelta_, + /// Number of partitions along K dimension + int PartitionsK_> +class MmaTensorOpMultiplicandTileIterator< + Shape_, Operand_, Element_, + mctlass::layout::RowMajor, + InstructionShape_, OpDelta_, 32, PartitionsK_> { + public: + + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand_; + + static_assert(kOperand == Operand::kA || kOperand== Operand::kB, + "MmaTensorOpMultiplicandIterator may only be instantiated for A or B operands to warp-level Mma."); + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads, for half_t m16n16k16,we use kThreads=64 + static int const kThreads = mctlass::platform::is_same::value ? 64 : 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Underlying tile iterator implementation + using Base = MmaTensorOpMultiplicandTileIteratorCanonical< + Shape, kOperand, Element, + layout::RowMajor, + InstructionShape, + kOpDelta, kThreads, PartitionsK_>; + + public: + + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + +private: + + /// Underlying tile iterator + Base iterator_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator(TensorRef const &ref, int lane_id) { + int t_id = lane_id; + int row = t_id & 0x7; + int col = (t_id >> 4) & 0x3; + t_id = row * 4 + col; + iterator_ = Base({ref.data(), ref.stride()}, t_id); + + if (mctlass::platform::is_same::value) { + iterator_ = Base({ref.data(), ref.stride()}, lane_id); + } + } + // MCTLASS_HOST_DEVICE + // MmaTensorOpMultiplicandTileIterator( + // TensorRef const &ref, + // int lane_id + // ): iterator_({ref.data(), ref.stride()}, lane_id) { + // } + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator( + TensorRef const &ref, + TensorCoord const &extent, + int lane_id + ): iterator_({ref.data(), ref.stride()}, extent, lane_id) { + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator++() { + + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpMultiplicandTileIterator & operator--() { + + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.row(), tile_offset.column())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + iterator_.load(frag); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, + {tile_offset.contiguous(), tile_offset.strided()}, + byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_sparse.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_sparse.h new file mode 100644 index 0000000..b347287 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_sparse.h @@ -0,0 +1,380 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines iterators to load sparse meta data used by warp-level matrix multiply operations + targeting Sparse Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/gemm/gemm.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" + +#include "mctlass/platform/platform.h" +#include "mctlass/fast_math.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of A elements + typename Element_, + /// Layout of operand + typename Layout_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + int OpDelta_, + /// Number of threads participating in one matrix operation + int Threads, + /// Number of partitions along K dimension + int PartitionsK_ = 1> +class SparseMmaTensorOpMetaTileIterator { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = Layout_; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: + /// MatrixShape) + static int const kOpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + static int const kSparse = 2; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + struct Policy { + static_assert( + !(Shape::kColumn % InstructionShape::kColumn), + "Shape of warp-level Mma must be divisible by operator shape."); + + static int const kElementsPerAccess = 128 / sizeof_bits::value; + + // Determine number of elements along outer dimension per individual LDSM op + static int const kLdsmOpOuter = InstructionShape::kColumn; + static int const kLdsmOpInner = 8 * kElementsPerAccess / kLdsmOpOuter; + + static_assert(!(Shape::kColumn % kLdsmOpOuter), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + static_assert(!(Shape::kRow % kLdsmOpInner), + "Shape of warp-level mma must be divisible by LDSM's " + "fundamental tile size."); + + /// Shape of one individual LDSM instruction + static int const LdsmShapeColumn = + InstructionShape::kColumn / kLdsmOpOuter; + static int const LdsmShapeRow = + ((4 / LdsmShapeColumn * kLdsmOpInner) > Shape::kRow) + ? (Shape::kRow / kLdsmOpInner) + : (4 / LdsmShapeColumn); + using LdsmShape = + layout::PitchLinearShape; + + /// Number and arrangement of LDSM instructions + using LdsmIterations = layout::PitchLinearShape< + Shape::kRow / kLdsmOpInner / LdsmShapeRow, + 1>; + + /// Number of groups for each tile + static int const kGroupsPerTile = + Shape::kColumn / InstructionShape::kColumn; + }; + + private: + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + /// Pointer type used for accesses + using AccessType = Array; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = + Array; + + private: + + /// Layout object storing stride values + Index stride_; + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + /// Internal counter used to determine when to increment byte offset and when + /// to XOR it + int k_group_idx_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + SparseMmaTensorOpMetaTileIterator() + : pointer_(nullptr), + stride_(0), + byte_offset_(0), + k_group_idx_(0) {} + + /// Constructor from TensorRef + MCTLASS_DEVICE + SparseMmaTensorOpMetaTileIterator(TensorRef const &ref, int lane_id) + : pointer_(reinterpret_cast(ref.data())), + stride_(ref.stride(0) / Policy::kElementsPerAccess), + byte_offset_(0), + k_group_idx_(0) { + + int access_contiguous = (lane_id % (Shape::kRow / Policy::kElementsPerAccess)); + int access_strided = (lane_id / (Shape::kRow / Policy::kElementsPerAccess)); + + byte_offset_ = (access_contiguous + access_strided * stride_) * + sizeof_bits::value * Policy::kElementsPerAccess / 8; + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + SparseMmaTensorOpMetaTileIterator &add_pointer_offset(LongIndex offset) { + byte_offset_ += offset * sizeof_bits::value / 8; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + SparseMmaTensorOpMetaTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + int offset = tile_offset.row() * Shape::kRow + + tile_offset.column() * InstructionShape::kColumn * stride_ * + Policy::kElementsPerAccess; + + add_pointer_offset(offset); + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + SparseMmaTensorOpMetaTileIterator &operator++() { + add_tile_offset({0, 1}); + + if (kPartitionsK > 1) { + ++k_group_idx_; + // Jump to next stage + if (k_group_idx_ == Policy::kGroupsPerTile) { + k_group_idx_ = 0; + add_tile_offset( + {0, ((kPartitionsK - 1) * Policy::kGroupsPerTile)}); + } + } + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + SparseMmaTensorOpMetaTileIterator &operator--(){ + byte_offset_ -= stride_ * InstructionShape::kColumn * + sizeof_bits::value * Policy::kElementsPerAccess / + 8; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE SparseMmaTensorOpMetaTileIterator & + operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + SparseMmaTensorOpMetaTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { load_with_byte_offset(frag, 0); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + Array *fetch_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < Policy::LdsmIterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsmIterations::kContiguous; ++c) { + + int access_idx = c + s * Policy::LdsmIterations::kContiguous; + + AccessType const *source_ptr = + pointer_ + + Policy::LdsmShape::kContiguous * Policy::kLdsmOpInner * c + + Policy::LdsmShape::kStrided * s * stride_; + + char const *source_byte_ptr = reinterpret_cast(source_ptr) + + byte_offset + byte_offset_; + + mctlass::arch::ldsm( + fetch_ptr[access_idx], source_byte_ptr); + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = + tile_offset.contiguous() * Shape::kRow / Layout::kElementsPerAccess + + tile_offset.strided() * InstructionShape::kColumn * stride_; + + byte_offset += sizeof(AccessType) * pointer_offset; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no op + } +}; + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_wmma.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_wmma.h new file mode 100644 index 0000000..a405582 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_tile_iterator_wmma.h @@ -0,0 +1,800 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines iterators used by warp-level matrix multiply operations targeting Tensor Cores. +*/ + +#pragma once + + +#include "mctlass/mctlass.h" +#include "mctlass/arch/wmma.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) + +#include "mctlass/wmma_array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/gemm/gemm.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" + +#include "mctlass/platform/platform.h" +#include "mctlass/fast_math.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +//////////////////////////////////////////////////////////////////////////////// +template < + ///< Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Operand identity (A or B) + Operand Operand, + /// Data type of operand + typename Element_, + /// Layout of operand + typename Layout_, + /// Delta between *MMA operations (in units of *WMMA operations, concept:MatrixShape) + int OpDelta_, + /// Number of threads participating in one matrix operation + int Threads, + /// Shape of the warp in units of thread (concept: MmaTensorOpPolicy) + typename Policy_> +class MmaTensorOpWmmaMultiplicandTileIterator; + +//////////////////////////////////////////////////////////////////////////////// +/// This tile iterator is specialized for 64-thread WMMA operation. +/// It uses wmma::load_matrix_sync to load from shared +/// memory and therefore must be initialized with a TensorRef to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +//////////////////////////////////////////////////////////////////////////////// +template < + ///< Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Layout of operand + typename Layout_, + /// Interval between adjacent *WMMA instructions (in units of WMMA instructions) + int OpDelta_, + /// Shape of the warp in units of thread (concept: MmaTensorOpPolicy) + typename Policy_> +class MmaTensorOpWmmaMultiplicandTileIterator< + Shape_, Operand::kA, Element_, Layout_, + OpDelta_, 64, Policy_> { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kA; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = Layout_; + + /// Delta between *WMMA operations + static int const kOpDelta = OpDelta_; + + /// Wmma Operator information and operation delta + using Policy = Policy_; + + + // + // Derived quantities + // + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Stride Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Native Wmma shape for operand A (concept MatrixShape) + using WmmaShape = MatrixShape< + Policy::Operator::Shape::kM, + Policy::Operator::Shape::kK + >; + + /// Map mctlass dataype to wmma datatype + using WmmaDataType = typename mctlass::arch::MctlassToWmmaDataType::Type; + + /// Shape of individual WMMA load / stores for operand A + using Iterations = MatrixShape< + Shape::kRow / WmmaShape::kRow, + 1 + >; + + /// Fragment object holding a warps part + using Fragment = WmmaFragmentArray; + + + ////////////////////////////////////////////////////////////////////////////////////////////////////// + /// statically assert this specialization + ///////////////////////////////////////////////////////////////////////////////////////////////////// + /// This iterator is specalized for Operand A + static_assert(kOperand == Operand::kA, + "MmaTensorOpWmmaMultiplicandTileIterator may only be instantiated for A operands to warp-level Mma."); + + /// Supported memory layouts + static_assert( + platform::is_same::value || + platform::is_same::value, + "Supported list of memory layouts for WMMA are: RowMajor, ColumnMajor"); + + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + ///////////////////////////////////////////////////////////////////////////////////////////////////// + +private: + + /// Shared memory base pointers - not advanced + char const *pointer_; + + /// Byte offset into shared memory - advanced + Index byte_offset_; + + /// Stride in units of number of elements + StrideIndex stride_; + + /// Layout of shared memory + Layout layout_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): pointer_(reinterpret_cast(ref.data())), byte_offset_(0), stride_(ref.stride(0)), layout_(ref.stride(0)) { + + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + byte_offset_ += (offset * sizeof_bits::value) / 8; + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + Index elements_offset = layout_({tile_offset.row() * Shape::kRow, tile_offset.column() * WmmaShape::kColumn}); + + byte_offset_ += (elements_offset * sizeof_bits::value) / 8; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator & operator++() { + + Index elements_offset = layout_({0, WmmaShape::kColumn}); + + byte_offset_ += (elements_offset * sizeof_bits::value) / 8; + + return *this; + } + + /// Advances the iterator along the opposite of the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator & operator--() { + + Index elements_offset = layout_({0, WmmaShape::kColumn}); + + byte_offset_ -= (elements_offset * sizeof_bits::value) / 8; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load_with_byte_offset(Fragment &frag, Index byte_offset) const { + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kColumn; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Iterations::kRow; ++m) { + + Index load_byte_offset = layout_({m * WmmaShape::kRow, k * WmmaShape::kColumn}) * sizeof_bits::value / 8; + + const WmmaDataType *ptr = reinterpret_cast(pointer_ + byte_offset_ + load_byte_offset + byte_offset); + // mxmaca::wmma::load_matrix_sync(frag[m], ptr, stride_); + mxmaca::wmma::load_matrix_sync(frag[m], ptr, stride_); + } + } + } + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_byte_offset(frag, 0); + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store_with_byte_offset(Fragment const &frag, Index byte_offset) const { + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kColumn; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Iterations::kRow; ++m) { + + Index store_byte_offset = layout_({m * WmmaShape::kRow, k * WmmaShape::kColumn}) * sizeof_bits::value / 8; + + WmmaDataType *ptr = reinterpret_cast(pointer_ + byte_offset_ + store_byte_offset + byte_offset); + // mxmaca::wmma::store_matrix_sync(ptr, frag[m], stride_); + mxmaca::wmma::store_matrix_sync(ptr, frag[m], stride_); + } + } + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_byte_offset(frag, 0); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation here + } +}; + +//////////////////////////////////////////////////////////////////////////////// +/// This tile iterator is specialized for 64-thread WMMA operation. +/// It uses wmma::load_matrix_sync to load from shared +/// memory and therefore must be initialized with a TensorRef to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +//////////////////////////////////////////////////////////////////////////////// + +template < + ///< Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Layout of operand + typename Layout_, + /// Interval between adjacent *WMMA instructions (in units of WMMA instructions) + int OpDelta_, + /// Shape of the warp in units of thread (concept: MmaTensorOpPolicy) + typename Policy_> +class MmaTensorOpWmmaMultiplicandTileIterator< + Shape_, Operand::kB, Element_, Layout_, + OpDelta_, 64, Policy_> { + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Operand tag + static Operand const kOperand = Operand::kB; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = Layout_; + + /// Delta between *WMMA operations + static int const kOpDelta = OpDelta_; + + /// Wmma Operator information and operation delta + using Policy = Policy_; + + + // + // Derived quantities + // + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Stride Index type + using StrideIndex = typename TensorRef::Layout::Stride::Index; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Native Wmma shape (concept MatrixShape) + using WmmaShape = MatrixShape< + Policy::Operator::Shape::kK, + Policy::Operator::Shape::kN + >; + + /// Map mctlass dataype to wmma datatype + using WmmaDataType = typename mctlass::arch::MctlassToWmmaDataType::Type; + + /// Shape of individual WMMA load / stores for operand B + using Iterations = MatrixShape< + 1, + Shape::kColumn / WmmaShape::kColumn + >; + + /// Fragment object holding a warps part + using Fragment = WmmaFragmentArray; + + + ////////////////////////////////////////////////////////////////////////////////////////////////////// + /// statically asserts this specialization + ///////////////////////////////////////////////////////////////////////////////////////////////////// + /// This iterator is specalized for Operand B + static_assert(kOperand == Operand::kB, + "MmaTensorOpWmmaMultiplicandTileIterator may only be instantiated for B operands to warp-level Mma."); + + /// Supported memory layouts + static_assert( + platform::is_same::value || + platform::is_same::value, + "Supported list of memory layouts for WMMA are: RowMajor, ColumnMajor"); + + /// Not working on this feature at the moment. + static_assert(kOpDelta == 1, + "Alternative arrangements not supported at present."); + + ///////////////////////////////////////////////////////////////////////////////////////////////////// + +private: + + /// Shared memory base pointers - not advanced + char const *pointer_; + + /// Byte offset into shared memory - advanced + Index byte_offset_; + + /// Stride in units of number of elements + StrideIndex stride_; + + /// Layout of shared memory + Layout layout_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator( + TensorRef const &ref, + int lane_id + ): pointer_(reinterpret_cast(ref.data())), byte_offset_(0), stride_(ref.stride(0)), layout_(ref.stride(0)) { + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator &add_pointer_offset(LongIndex offset) { + + byte_offset_ += (offset * sizeof_bits::value) / 8; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + + Index elements_offset = layout_({tile_offset.row() * WmmaShape::kRow, tile_offset.column() * Shape::kColumn}); + + byte_offset_ += (elements_offset * sizeof_bits::value) / 8; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator & operator++() { + + Index elements_offset = layout_({WmmaShape::kRow, 0}); + + byte_offset_ += (elements_offset * sizeof_bits::value) / 8; + + return *this; + } + + /// Advances the iterator along the opposite of the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator & operator--() { + + Index elements_offset = layout_({WmmaShape::kRow, 0}); + + byte_offset_ -= (elements_offset * sizeof_bits::value) / 8; + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpWmmaMultiplicandTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load_with_byte_offset(Fragment &frag, Index byte_offset) const { + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kRow; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Iterations::kColumn; ++n) { + + Index load_byte_offset = layout_({k * WmmaShape::kRow, n * WmmaShape::kColumn}) * sizeof_bits::value / 8; + + const WmmaDataType *ptr = reinterpret_cast(pointer_ + byte_offset_ + load_byte_offset + byte_offset); + // mxmaca::wmma::load_matrix_sync(frag[n], ptr, stride_); + mxmaca::wmma::load_matrix_sync(frag[n], ptr, stride_); + } + } + } + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_byte_offset(frag, 0); + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store_with_byte_offset(Fragment const &frag, Index byte_offset) const { + + MCTLASS_PRAGMA_UNROLL + for (int k = 0; k < Iterations::kRow; ++k) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Iterations::kColumn; ++n) { + + Index store_byte_offset = layout_({k * WmmaShape::kRow, n * WmmaShape::kColumn}) * sizeof_bits::value / 8; + + WmmaDataType *ptr = reinterpret_cast(pointer_ + byte_offset_ + store_byte_offset + byte_offset); + // mxmaca::wmma::store_matrix_sync(ptr, frag[n], stride_); + mxmaca::wmma::store_matrix_sync(ptr, frag[n], stride_); + } + } + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_byte_offset(frag, 0); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation here + } +}; + +//////////////////////////////////////////////////////////////////////////////// +template < + ///< Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Layout of operand in memory + typename Layout_, + /// Interval between adjacent *WMMA instructions (in units of WMMA instructions, concept: MatrixShape) + typename OpDelta_, + /// Shape of the warp in units of thread (concept: MmaTensorOpPolicy) + typename Policy_> +class MmaTensorOpWmmaAccumulatorTileIterator; + +//////////////////////////////////////////////////////////////////////////////// +/// This tile iterator is specialized for 64-thread WMMA operation. +/// It uses wmma::store_matrix_sync to load from shared +/// memory and therefore must be initialized with a TensorRef to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept | +/// WriteableRandomAccessContiguousTileIteratorConcept +/// +//////////////////////////////////////////////////////////////////////////////// + +template < + ///< Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Layout of operand in memory + typename Layout_, + /// Interval between adjacent *WMMA instructions (in units of WMMA instructions) + typename OpDelta_, + /// Shape of the warp in units of thread (concept: MmaTensorOpPolicy) + typename Policy_> +class MmaTensorOpWmmaAccumulatorTileIterator +{ + public: + + /// Shape of tile to load (concept: MatrixShape) + using Shape = Shape_; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = Layout_; + + /// Delta between *MMA operations (in units of *MMA operations, concept: MatrixShape) + using OpDelta = OpDelta_; + + /// Number of participating threads + static int const kThreads = 64; + + /// Wmma Operator information and operation delta + using Policy = Policy_; + + + // + // Derived quantities + // + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Native Wmma shape (concept MatrixShape) + using WmmaShape = MatrixShape< + Policy::Operator::Shape::kM, + Policy::Operator::Shape::kN + >; + + /// Map mctlass dataype to wmma datatype + using WmmaDataType = typename mctlass::arch::MctlassToWmmaDataType::Type; + + /// Map mctlass::layout to nvuda::wmma::layout_t enum + // static mxmaca::wmma::layout_t const WmmaLayout = mctlass::arch::MctlassToWmmaLayout::value; + static mxmaca::wmma::layout_t const WmmaLayout = mctlass::arch::MctlassToWmmaLayout::value; + + /// Shape of individual WMMA load / stores for accumulator + using Iterations = MatrixShape< + Shape::kRow / WmmaShape::kRow, + Shape::kColumn / WmmaShape::kColumn + >; + + /// Fragment object holding a thread's part of a tile + using Fragment = WmmaFragmentArray; + + ////////////////////////////////////////////////////////////////////////////////////////////////////// + /// statically asserts this specialization + ///////////////////////////////////////////////////////////////////////////////////////////////////// + /// Supported layouts + static_assert( + platform::is_same::value || + platform::is_same::value, + "Supported list of memory layouts for WMMA are: RowMajor, ColumnMajor"); + +private: + + /// Internal reference + mctlass::TensorRef ref_; + +public: + + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + MmaTensorOpWmmaAccumulatorTileIterator() { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + MmaTensorOpWmmaAccumulatorTileIterator( + TensorRef const &ref, + int lane_id + ): ref_(ref) { } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + MmaTensorOpWmmaAccumulatorTileIterator &add_pointer_offset(LongIndex offset) { + ref_.add_pointer_offset(offset); + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + MmaTensorOpWmmaAccumulatorTileIterator &add_tile_offset(TensorCoord const &tile_offset) { + ref_.add_coord_offset({tile_offset.row() * Shape::kRow, tile_offset.column() * Shape::kColumn}); + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + MmaTensorOpWmmaAccumulatorTileIterator & operator++() { + ref_.add_coord_offset({Shape::kRow, 0}); + return *this; + } + + /// Advances the iterator along the opposite of the advance dimension + MCTLASS_HOST_DEVICE + MmaTensorOpWmmaAccumulatorTileIterator & operator--() { + ref_.add_coord_offset({-Shape::kRow, 0}); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpWmmaAccumulatorTileIterator & operator+=(TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + MmaTensorOpWmmaAccumulatorTileIterator & operator-=(TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) const { + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Iterations::kRow; ++m) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Iterations::kColumn; ++n) { + + const WmmaDataType * ptr = reinterpret_cast (ref_.data() + ref_.offset({m * WmmaShape::kRow, n * WmmaShape::kColumn}) + pointer_offset); + // mxmaca::wmma::load_matrix_sync(frag[m * Iterations::kColumn + n], ptr, ref_.stride()[0], WmmaLayout); + mxmaca::wmma::load_matrix_sync(frag[m * Iterations::kColumn + n], ptr, ref_.stride()[0], WmmaLayout); + + } + } + } + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) const { + + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < Iterations::kRow; ++m) { + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < Iterations::kColumn; ++n) { + + WmmaDataType * ptr = reinterpret_cast (ref_.data() + ref_.offset({m * WmmaShape::kRow, n * WmmaShape::kColumn}) + pointer_offset); + // mxmaca::wmma::store_matrix_sync(ptr, frag[m * Iterations::kColumn + n], ref_.stride()[0], WmmaLayout); + mxmaca::wmma::store_matrix_sync(ptr, frag[m * Iterations::kColumn + n], ref_.stride()[0], WmmaLayout); + } + } + } + + /// Stores a fragment to memory at the location pointed to by the iterator + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) const { + store_with_pointer_offset(frag, 0); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + // no operation here + } +}; + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// + +#endif // if defined(MCTLASS_ARCH_WMMA_ENABLED) + + diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_wmma.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_wmma.h new file mode 100644 index 0000000..e8ada86 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_tensor_op_wmma.h @@ -0,0 +1,223 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations targeting + Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/arch/wmma.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) + +#include "mctlass/wmma_array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/arch/mma_sm75.h" +#include "mctlass/arch/mma_sm80.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/warp/mma_tensor_op_policy.h" + +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_wmma.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +///< Structure to compute the matrix product targeting CUDA cores via WMMA. +template < + ///< Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + ///< Data type of A elements + typename ElementA_, + ///< Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + ///< Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + ///< Element type of C matrix + typename ElementC_, + ///< Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + ///< Policy describing warp-level Wmma operation (concept: MmaTensorOpPolicy) + typename Policy_, + ///< Number of partitions along K dimension + int PartitionsK_ = 1, + ///< Used for partial specialization + typename Enable = bool +> +class MmaTensorOpWmma { +public: + ///< Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + ///< Data type of multiplicand A + using ElementA = ElementA_; + + ///< Layout of multiplicand A + using LayoutA = LayoutA_; + + ///< Data type of multiplicand B + using ElementB = ElementB_; + + ///< Layout of multiplicand B + using LayoutB = LayoutB_; + + ///< Data type of accumulator matrix C + using ElementC = ElementC_; + + ///< Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaTensorOpPolicy) + using Policy = Policy_; + + /// Underlying instruction shape + using InstructionShape = typename Policy::Operator::Shape; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Policy::Operator; + + /// Indicates math operator + using MathOperator = typename ArchMmaOperator::Operator; + + /// Underlying architecture tag + using ArchTag = typename Policy::Operator::ArchTag; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = ComplexTransform::kNone; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = ComplexTransform::kNone; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassWmmaTensorOp; + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 64; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaTensorOpWmmaMultiplicandTileIterator< + MatrixShape, Operand::kA, ElementA, LayoutA, + Policy::OpDelta::kRow, kThreadCount, Policy>; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Iterates over the B operand in memory + using IteratorB = MmaTensorOpWmmaMultiplicandTileIterator< + MatrixShape, Operand::kB, ElementB, LayoutB, + Policy::OpDelta::kRow, kThreadCount, Policy>; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Iterates over the C operand in memory + using IteratorC = MmaTensorOpWmmaAccumulatorTileIterator< + MatrixShape, ElementC, LayoutC, + typename Policy::OpDelta, Policy>; + + /// Storage for C tile + using FragmentC = typename IteratorC::Fragment; + +private: + + static_assert( + !(Shape::kM % Policy::Operator::Shape::kM) && + !(Shape::kN % Policy::Operator::Shape::kN), + "Shape of warp-level Wmma must be divisible by operator shape (wmma native size)"); + + /// Number of wmma operations performed + using WmmaIterations = MatrixShape< + Shape::kM / Policy::Operator::Shape::kM, + Shape::kN / Policy::Operator::Shape::kN + >; + +public: + + /// Underlying matrix multiply operator (concept: mctlass::arch::Wmma) + typename Policy::Operator wmma; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaTensorOpWmma() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + FragmentA const &A, + FragmentB const &B, + FragmentC const &C) const { + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < WmmaIterations::kColumn; ++n) { + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < WmmaIterations::kRow; ++m) { + + // accumulate wmma mma + wmma(D[m * WmmaIterations::kColumn + n], A[m], B[n], C[m * WmmaIterations::kColumn + n]); + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +#endif // if defined(MCTLASS_ARCH_WMMA_ENABLED) + diff --git a/csrc/mctlass/include/mctlass/gemm/warp/mma_with_reduction_tensor_op.h b/csrc/mctlass/include/mctlass/gemm/warp/mma_with_reduction_tensor_op.h new file mode 100644 index 0000000..3c1fe7f --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/mma_with_reduction_tensor_op.h @@ -0,0 +1,453 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations targeting + Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/platform/platform.h" + +#include "mctlass/numeric_conversion.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/arch/mma_sm75.h" +#include "mctlass/arch/mma_sm80.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/warp/mma_tensor_op_policy.h" +#include "mctlass/gemm/warp/mma_tensor_op.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure to compute the matrix product targeting CUDA cores and SIMT math instructions. +template < + /// Size of the Gemm problem - concept: gemm::GemmShape<> + typename Shape_, + /// Data type of A elements + typename ElementA_, + /// Layout of A matrix (concept: MatrixLayout) + typename LayoutA_, + /// Data type of B elements + typename ElementB_, + /// Layout of B matrix (concept: MatrixLayout) + typename LayoutB_, + /// Element type of C matrix + typename ElementC_, + /// Layout of C matrix (concept: MatrixLayout) + typename LayoutC_, + /// Policy describing warp-level MmaTensorOp (concept: MmaTensorOp policy) + typename Policy_, + /// Reduce operand A or B along K dimension + bool ReduceKForA_, + /// Number of partitions along K dimension + int PartitionsK_ = 1, + /// Store the accumulators in row major or column major. Row major is used + /// when output layout is interleaved. + bool AccumulatorsInRowMajor = false, + /// Used for partial specialization + typename Enable = bool +> +class MmaWithReductionTensorOp { +public: + /// Shape of warp-level matrix operation (concept: GemmShape) + using Shape = Shape_; + + /// Data type of multiplicand A + using ElementA = ElementA_; + + /// Layout of multiplicand A + using LayoutA = LayoutA_; + + /// Data type of multiplicand B + using ElementB = ElementB_; + + /// Layout of multiplicand B + using LayoutB = LayoutB_; + + /// Data type of accumulator matrix C + using ElementC = ElementC_; + + /// Layout of accumulator matrix C + using LayoutC = LayoutC_; + + /// Shape of the warp in units of thread (concept: MmaLanePolicySimt) + using Policy = Policy_; + + /// Underlying matrix multiply operator (concept: arch::Mma) + using ArchMmaOperator = typename Policy::Operator; + + /// Indicates math operator + using MathOperator = typename ArchMmaOperator::Operator; + + /// Architecture tag from underlying instruction + using ArchTag = typename ArchMmaOperator::ArchTag; + + /// Indicates class of matrix operator + using OperatorClass = arch::OpClassTensorOp; + + /// Shape of underlying instruction + using InstructionShape = typename ArchMmaOperator::Shape; + + /// Complex transform on A operand + static ComplexTransform const kTransformA = ComplexTransform::kNone; + + /// Complex transform on B operand + static ComplexTransform const kTransformB = ComplexTransform::kNone; + + /// Number of threads participating in warp-level matrix product + static int const kThreadCount = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + static bool const kReduceKForA = ReduceKForA_; + + static_assert(platform::is_same::value || + platform::is_same::value, + "ElementA needs to be fp16 or bf16."); + + static_assert(platform::is_same::value || + platform::is_same::value, + "ElementB needs to be fp16 or bf16."); + + static_assert(platform::is_same>::value, + "Only supports 16x8x16 tensor core instruction."); + + static_assert(!AccumulatorsInRowMajor, + "Only calls tensor core instructions in column major."); + +public: + + /// Iterates over the A operand in memory + using IteratorA = MmaTensorOpMultiplicandTileIterator< + MatrixShape, Operand::kA, ElementA, LayoutA, + MatrixShape, + Policy::OpDelta::kRow, kThreadCount, kPartitionsK>; + + /// Storage for A tile + using FragmentA = typename IteratorA::Fragment; + + /// Storage for transformed A tile + using TransformedFragmentA = + Array; + + /// Iterates over the B operand in memory + using IteratorB = MmaTensorOpMultiplicandTileIterator< + MatrixShape, Operand::kB, ElementB, LayoutB, + MatrixShape, + Policy::OpDelta::kRow, kThreadCount, kPartitionsK>; + + /// Storage for B tile + using FragmentB = typename IteratorB::Fragment; + + /// Storage for transformed B tile + using TransformedFragmentB = + Array; + + /// Iterates over the C operand in memory + using IteratorC = MmaTensorOpAccumulatorTileIterator< + MatrixShape, ElementC, LayoutC, + typename ArchMmaOperator::Shape, typename Policy::OpDelta>; + + /// Storage for C tile + using FragmentC = typename IteratorC::Fragment; + + /// Number of mma operations performed + using MmaIterations = MatrixShape< + (Shape::kM + ArchMmaOperator::Shape::kM - 1) / ArchMmaOperator::Shape::kM, + (Shape::kN + ArchMmaOperator::Shape::kN - 1) / ArchMmaOperator::Shape::kN + >; + + using FragmentReduction = Array; + +public: + + /// Underlying matrix multiply operator (concept: arch::Mma) + ArchMmaOperator mma; + +public: + + // + // Methods + // + + /// Ctor + MCTLASS_DEVICE + MmaWithReductionTensorOp() {} + + /// Performs a warp-level matrix multiply-accumulate operation + MCTLASS_DEVICE + void operator()( + FragmentC &D, + TransformedFragmentA const &A, + TransformedFragmentB const &B, + FragmentC const &C, + FragmentReduction &gemm_k_reduction + ) const { + + using MmaOperandA = typename ArchMmaOperator::FragmentA; + using MmaOperandB = typename ArchMmaOperator::FragmentB; + using MmaOperandC = typename ArchMmaOperator::FragmentC; + + D = C; + + MmaOperandA const *ptr_A = reinterpret_cast(&A); + MmaOperandB const *ptr_B = reinterpret_cast(&B); + MmaOperandC *ptr_D = reinterpret_cast(&D); + + //#if defined(__MACA_ARCH__) && (__TLASS_ARCH__ < 800) + #if 0 + assert(0); + #elif defined(__MACA_ARCH__) + // Serpentine visitation order maximizing reuse of Ra + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < MmaIterations::kRow; ++m) { + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < MmaIterations::kColumn; ++n) { + + int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n); + + mma(ptr_D[m + n_serpentine * MmaIterations::kRow], + ptr_A[m], + ptr_B[n_serpentine], + ptr_D[m + n_serpentine * MmaIterations::kRow]); + + if (!kReduceKForA && m == 0) { + #if 1 + gemm_k_reduction[n_serpentine] += float(B[n_serpentine * 4]); + gemm_k_reduction[n_serpentine] += float(B[n_serpentine * 4 + 1]); + gemm_k_reduction[n_serpentine] += float(B[n_serpentine * 4 + 2]); + gemm_k_reduction[n_serpentine] += float(B[n_serpentine * 4 + 3]); + #else + uint32_t const *tmp = reinterpret_cast(&B); + + if (platform::is_same::value) { + asm volatile( + "{\n\t" + " .reg .f16 low, high;\n\t" + " .reg .f32 tmp;\n\t" + " mov.b32 {low, high}, %1;\n\t" + " cvt.f32.f16 tmp, low;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " cvt.f32.f16 tmp, high;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " mov.b32 {low, high}, %2;\n\t" + " cvt.f32.f16 tmp, low;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " cvt.f32.f16 tmp, high;\n\t" + " add.f32 %0, tmp, %0;\n\t" + "}\n\t" + : "+f"(gemm_k_reduction[n_serpentine]) + : "r"(tmp[n_serpentine * 2]), "r"(tmp[n_serpentine * 2 + 1])); + } else if (platform::is_same::value) { + asm volatile( + "{\n\t" + " .reg .f32 tmp;\n\t" + " shl.b32 tmp, %1, 16;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " and.b32 tmp, %1, 0xffff0000;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " shl.b32 tmp, %2, 16;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " and.b32 tmp, %2, 0xffff0000;\n\t" + " add.f32 %0, tmp, %0;\n\t" + "}\n\t" + : "+f"(gemm_k_reduction[n_serpentine]) + : "r"(tmp[n_serpentine * 2]), "r"(tmp[n_serpentine * 2 + 1])); + } else { + assert(0); + } + #endif + #endif + } + + if (kReduceKForA && (n == 0)) { + #if 1 + gemm_k_reduction[m * 2] += float(A[m * 8]); + gemm_k_reduction[m * 2] += float(A[m * 8 + 1]); + gemm_k_reduction[m * 2] += float(A[m * 8 + 4]); + gemm_k_reduction[m * 2] += float(A[m * 8 + 5]); + + gemm_k_reduction[m * 2 + 1] += float(A[m * 8 + 2]); + gemm_k_reduction[m * 2 + 1] += float(A[m * 8 + 3]); + gemm_k_reduction[m * 2 + 1] += float(A[m * 8 + 6]); + gemm_k_reduction[m * 2 + 1] += float(A[m * 8 + 7]); + #else + uint32_t const *tmp = reinterpret_cast(&A); + + if (platform::is_same::value) { + asm volatile( + "{\n\t" + " .reg .f16 low, high;\n\t" + " .reg .f32 tmp;\n\t" + " mov.b32 {low, high}, %2;\n\t" + " cvt.f32.f16 tmp, low;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " cvt.f32.f16 tmp, high;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " mov.b32 {low, high}, %3;\n\t" + " cvt.f32.f16 tmp, low;\n\t" + " add.f32 %1, tmp, %1;\n\t" + " cvt.f32.f16 tmp, high;\n\t" + " add.f32 %1, tmp, %1;\n\t" + " mov.b32 {low, high}, %4;\n\t" + " cvt.f32.f16 tmp, low;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " cvt.f32.f16 tmp, high;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " mov.b32 {low, high}, %5;\n\t" + " cvt.f32.f16 tmp, low;\n\t" + " add.f32 %1, tmp, %1;\n\t" + " cvt.f32.f16 tmp, high;\n\t" + " add.f32 %1, tmp, %1;\n\t" + "}\n\t" + : "+f"(gemm_k_reduction[m * 2]), "+f"(gemm_k_reduction[m * 2 + 1]) + : "r"(tmp[m * 4]), "r"(tmp[m * 4 + 1]),"r"(tmp[m * 4 + 2]), "r"(tmp[m * 4 + 3])); + + } else if (platform::is_same::value) { + + asm volatile( + "{\n\t" + " .reg .f32 tmp;\n\t" + " shl.b32 tmp, %2, 16;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " and.b32 tmp, %2, 0xffff0000;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " shl.b32 tmp, %3, 16;\n\t" + " add.f32 %1, tmp, %1;\n\t" + " and.b32 tmp, %3, 0xffff0000;\n\t" + " add.f32 %1, tmp, %1;\n\t" + " shl.b32 tmp, %4, 16;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " and.b32 tmp, %4, 0xffff0000;\n\t" + " add.f32 %0, tmp, %0;\n\t" + " shl.b32 tmp, %5, 16;\n\t" + " add.f32 %1, tmp, %1;\n\t" + " and.b32 tmp, %5, 0xffff0000;\n\t" + " add.f32 %1, tmp, %1;\n\t" + "}\n\t" + : "+f"(gemm_k_reduction[m * 2]), "+f"(gemm_k_reduction[m * 2 + 1]) + : "r"(tmp[m * 4]), "r"(tmp[m * 4 + 1]),"r"(tmp[m * 4 + 2]), "r"(tmp[m * 4 + 3])); + + } else { + assert(0); + } + #endif + #endif + } + } + } + #else + assert(0); + #endif + } + + /// Transform the mma operands to the required types + MCTLASS_DEVICE + void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B, + FragmentA const &A, FragmentB const &B) const { + + // + // Define conversions from source type to instruction type + // + FloatRoundStyle const kRoundA = + PreferredRoundingMode::kRound; + FloatRoundStyle const kRoundB = + PreferredRoundingMode::kRound; + //#if defined(__MACA_ARCH__) && (__TLASS_ARCH__ < 800) + #if 0 + detail::ConvertAndPack + convert_A; + NumericArrayConverter + convert_B; + Array const *ptr_B = + reinterpret_cast const *>(&B); + Array * + ptr_dst_B = reinterpret_cast *>(&dst_B); + + dst_A = convert_A(A); + + ptr_dst_B[0] = convert_B(ptr_B[0]); + ptr_dst_B[1] = convert_B(ptr_B[1]); + + #elif defined(__MACA_ARCH__) + detail::ConvertAndPack + convert_A; + NumericArrayConverter + convert_B; + Array const *ptr_A = + reinterpret_cast const *>(&A); + Array * + ptr_dst_A = reinterpret_cast *>(&dst_A); + + dst_B = convert_B(B); + + ptr_dst_A[0] = convert_A(ptr_A[0]); + ptr_dst_A[1] = convert_A(ptr_A[1]); + #else + assert(0); + #endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/scale_bias_tile_iterator.h b/csrc/mctlass/include/mctlass/gemm/warp/scale_bias_tile_iterator.h new file mode 100644 index 0000000..a2fedac --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/scale_bias_tile_iterator.h @@ -0,0 +1,572 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Defines iterators used by warp-level loading scale and bias vectors. + Every scale/bias data only needs to be loaded once for every channel. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/gemm/gemm.h" + +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" + +#include "mctlass/platform/platform.h" +#include "mctlass/fast_math.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of A elements + typename Element_, + /// Layout of operand + typename Layout_, + /// Shape of one matrix production operation (concept: GemmShape) + typename InstructionShape_, + /// Policy of the details of LDSM shape and iterations + typename Policy_, + /// Number of threads participating in one matrix operation + int Threads, + /// Number of partitions along K dimension + int PartitionsK_ = 1> +class ScaleBiasTileIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: PitchLinearShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: PitchLinearShape) + typename InstructionShape_, + /// Policy of the details of LDSM shape and iterations + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK_> +class ScaleBiasTileIterator { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::PitchLinear; + + /// Shape of one matrix product operation (concept: GemmShape) + using InstructionShape = InstructionShape_; + + /// Number of participating threads + static int const kThreads = 32; + + /// Number of partitions along K dimension + static int const kPartitionsK = PartitionsK_; + + /// Number of partitions along K dimension + static int const kElementsPerAccess = 128 / sizeof_bits::value; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + using Policy = Policy_; + + private: + + /// Pointer type used for accesses + using AccessType = Array; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = Array; + + private: + + /// Shared memory base pointers - not advanced + AccessType const *pointer_; + + /// Byte offset incremented as iterator advances + Index byte_offset_; + + /// Internal counter used to determine when to increment byte offset and when + /// to XOR it + int k_group_idx_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + ScaleBiasTileIterator() + : pointer_(nullptr), + byte_offset_(0), + k_group_idx_(0) {} + + /// Constructor from TensorRef + MCTLASS_DEVICE + ScaleBiasTileIterator(TensorRef const &ref_scale_bias, + int lane_id) + : byte_offset_(0), k_group_idx_(0) { + /// 16816 only + pointer_ = reinterpret_cast(ref_scale_bias.data()) + + ((lane_id >> 3) & 1) * Shape::kContiguous / kElementsPerAccess + + (lane_id >> 4); + } + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + ScaleBiasTileIterator &add_pointer_offset(LongIndex offset) { + byte_offset_ += offset * sizeof_bits::value / 8; + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + ScaleBiasTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + int whole_tiles = tile_offset.contiguous() / Policy::kGroupsPerTile; + int k_groups_delta = tile_offset.contiguous() % Policy::kGroupsPerTile; + + byte_offset_ += k_groups_delta * sizeof_bits::value * + kElementsPerAccess * Policy::LdsmShape::kContiguous / 8; + + // Multiply by 2 because scale and bias belonging to the same stage are next + // to each other in the shared memory. + pointer_ += (2 * whole_tiles * Shape::kContiguous / kElementsPerAccess); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + ScaleBiasTileIterator &operator++() { + byte_offset_ += Policy::LdsmShape::kContiguous * + sizeof_bits::value * kElementsPerAccess / 8; + + k_group_idx_++; + + if (k_group_idx_ == (Policy::kGroupsPerTile / kPartitionsK)) { + k_group_idx_ = 0; + byte_offset_ -= (Policy::kGroupsPerTile / kPartitionsK) * + Policy::LdsmShape::kContiguous * + sizeof_bits::value * kElementsPerAccess / 8; + add_tile_offset({Policy::kGroupsPerTile, 0}); + } + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + ScaleBiasTileIterator &operator--() { assert(0); } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + ScaleBiasTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + ScaleBiasTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { load_with_byte_offset(frag, 0); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + Array *fetch_ptr = + reinterpret_cast *>(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < 1; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < Policy::LdsmIterations::kContiguous; ++c) { + int access_idx = c + s * Policy::LdsmIterations::kContiguous; + + AccessType const *source_ptr = + pointer_ + Policy::LdsmShape::kContiguous * c; + + char const *source_byte_ptr = + reinterpret_cast(source_ptr) + byte_offset + + byte_offset_; + + mctlass::arch::ldsm( + fetch_ptr[access_idx], source_byte_ptr); + } + } + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + load_with_byte_offset(frag, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + load_with_byte_offset(frag, tile_offset, 0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + load_with_byte_offset(frag, tile_offset, pointer_offset * sizeof(Element)); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + Index pointer_offset = tile_offset.contiguous() * + InstructionShape::kContiguous / + kElementsPerAccess; + + byte_offset += sizeof_bits::value * pointer_offset / 8; + + load_with_byte_offset(frag, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + k_group_idx_ = k_group % (Policy::kGroupsPerTile / kPartitionsK); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// This tile iterator is specialized for 32-thread TensorOps. It uses LDSM to +/// load from shared memory and therefore must be initialized with a TensorRef +/// to shared memory. +/// +/// Satisfies: +/// ReadableRandomAccessContiguousTileIteratorConcept +/// +template < + /// Size of the matrix to load (concept: MatrixShape) + typename Shape_, + /// Data type of elements + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + /// Policy of the details of LDSM shape and iterations + typename Policy_, + /// Number of partitions along K dimension + int PartitionsK_> +class ScaleBiasTileIterator { + public: + /// Shape of tile to load (concept: PitchLinearShape) + using Shape = Shape_; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Number of participating threads + static int const kThreads = 32; + + /// TensorRef type for loading element from a tensor + using TensorRef = TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Internal structure of iterator - made public to enable introspection + using Policy = Policy_; + + /// Underlying tile iterator implementation + using Base = ScaleBiasTileIterator< + layout::PitchLinearShape, Element, + layout::PitchLinear, + layout::PitchLinearShape, + Policy, kThreads, PartitionsK_>; + + public: + // + // Derived quantities + // + + /// Fragment object holding a thread's part of a tile + using Fragment = typename Base::Fragment; + + private: + /// Underlying tile iterator + Base iterator_; + + public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + ScaleBiasTileIterator() {} + + /// Constructor from TensorRef + MCTLASS_HOST_DEVICE + ScaleBiasTileIterator(TensorRef const &ref_scale_bias, int lane_id) + : iterator_({ref_scale_bias.data(), ref_scale_bias.stride()}, lane_id) {} + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_HOST_DEVICE + ScaleBiasTileIterator &add_pointer_offset(LongIndex offset) { + iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + ScaleBiasTileIterator &add_tile_offset( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_DEVICE + ScaleBiasTileIterator &add_tile_offset_negative( + TensorCoord const &tile_offset) { + iterator_.add_tile_offset_negative({tile_offset.column(), tile_offset.row()}); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + ScaleBiasTileIterator &operator++() { + ++iterator_; + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_HOST_DEVICE + ScaleBiasTileIterator &operator--() { + --iterator_; + + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + ScaleBiasTileIterator &operator+=( + TensorCoord const &tile_offset) { + add_tile_offset(PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of + ///< the tensor + MCTLASS_DEVICE + ScaleBiasTileIterator &operator-=( + TensorCoord const &tile_offset) { + add_tile_offset(-PitchLinearCoord(tile_offset.column(), tile_offset.row())); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { iterator_.load(frag); } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index byte_offset) const { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + assert(0); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + iterator_.load_with_byte_offset( + frag, {tile_offset.strided(), tile_offset.contiguous()}, byte_offset); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + iterator_.set_kgroup_index(k_group); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/softmax_scale_bias_transform.h b/csrc/mctlass/include/mctlass/gemm/warp/softmax_scale_bias_transform.h new file mode 100644 index 0000000..82cf87b --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/softmax_scale_bias_transform.h @@ -0,0 +1,117 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level per-channel softmax before + matrix multiply-accumulate operations targeting Tensor Cores. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/platform/platform.h" + +#include "mctlass/numeric_conversion.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/arch/memory_sm75.h" +#include "mctlass/arch/mma_sm75.h" +#include "mctlass/arch/mma_sm80.h" + +#include "mctlass/gemm/gemm.h" +#include "mctlass/gemm/warp/mma.h" + +#include "mctlass/gemm/warp/mma_tensor_op_policy.h" + +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator.h" +#include "mctlass/gemm/warp/mma_tensor_op_tile_iterator_sm80.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct SoftmaxScaleBiasTransform { + + using T = typename FragmentActivations::Element; + + static int const NumActivations = FragmentActivations::kElements; + static int const NumNormSum = FragmentNormSum::kElements; + static int const MmaElements = 2; + // One element has one scale and one bias + static int const MmaScaleBiasPair = 2; + // 16816 has 2 columns and 2 rows + static int const MmaCols = 2; + static int const MmaRows = 2; + + using MmaOperand = Array; + using NormSumOperand = Array<__half2, MmaScaleBiasPair>; + + MCTLASS_DEVICE + void transform(MmaOperand &activations, + NormSumOperand const &norm_sum) { + + __half2* packed_activations = reinterpret_cast<__half2*>(&activations); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < MmaElements / 2; ++i) { + __half2 out = ::h2exp(__hsub2(packed_activations[i], norm_sum[2*i])); + packed_activations[i] = __hmul2(out, norm_sum[2*i + 1]); + } + } + + MCTLASS_DEVICE + void operator()(FragmentActivations &activations, + FragmentNormSum const &norm_sum) { + MmaOperand *ptr_activations = reinterpret_cast(&activations); + NormSumOperand const *ptr_norm_sum = + reinterpret_cast(&norm_sum); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < (NumActivations / MmaElements); ++i) { + transform(ptr_activations[i], + ptr_norm_sum[i / (MmaCols * MmaRows) * MmaRows + i % MmaRows]); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/gemm/warp/tile_iterator_planar_complex.h b/csrc/mctlass/include/mctlass/gemm/warp/tile_iterator_planar_complex.h new file mode 100644 index 0000000..0f698c4 --- /dev/null +++ b/csrc/mctlass/include/mctlass/gemm/warp/tile_iterator_planar_complex.h @@ -0,0 +1,250 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing warp-level matrix multiply-accumulate operations. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/gemm/gemm.h" + +#include "mctlass/array_planar_complex.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace gemm { +namespace warp { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +class TileIteratorPlanarComplex { +public: + + /// Underlying iterator over real-valued tiles + using TileIterator = TileIterator_; + + /// Underlying element type + using Element = typename TileIterator::Element; + + /// Underlying layout type + using Layout = typename TileIterator::Layout; + + /// TensorRef type for loading element from a tensor + using TensorRef = typename TileIterator::TensorRef; + + /// Index type + using Index = typename TensorRef::Index; + + /// Long Index type + using LongIndex = typename TensorRef::LongIndex; + + /// Coordinate for an element in the tensor + using TensorCoord = typename TensorRef::TensorCoord; + + /// Planar complex fragment + using Fragment = ArrayPlanarComplex; + +public: + + /// Underlying tile iterator + TileIterator tile_iterator_; + + /// Offset (in units of bytes) to the imaginary part of the planar complex matrix + LongIndex imaginary_offset_; + +public: + /// Default ctor constructs null iterator + MCTLASS_HOST_DEVICE + TileIteratorPlanarComplex(): imaginary_offset_(0) { } + + /// Constructor from TensorRef + MCTLASS_DEVICE + TileIteratorPlanarComplex( + TensorRef const &ref, + int lane_id, + LongIndex imaginary_offset + ): + tile_iterator_(ref, lane_id), + imaginary_offset_((imaginary_offset * sizeof_bits::value) / 8) { } + + + /// Adds a pointer offset to internal pointer(s) to advance through memory + MCTLASS_DEVICE + TileIteratorPlanarComplex &add_pointer_offset(LongIndex offset) { + + tile_iterator_.add_pointer_offset(offset); + + return *this; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_HOST_DEVICE + TileIteratorPlanarComplex &add_tile_offset(TensorCoord const &tile_offset) { + + tile_iterator_.add_tile_offset(tile_offset); + + return *this; + } + + /// Advances the iterator along the advance dimension + MCTLASS_DEVICE + TileIteratorPlanarComplex & operator++() { + ++tile_iterator_; + return *this; + } + + // + // WIP + // + + /// Advances the iterator along the opposite of the advance dimension + MCTLASS_HOST_DEVICE + TileIteratorPlanarComplex & operator--() { + --tile_iterator_; + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + TileIteratorPlanarComplex & operator+=(TensorCoord const &tile_offset) { + tile_iterator_.add_tile_offset(tile_offset); + return *this; + } + + ///< advances in units of whole tiles along the logical coordinate space of the tensor + MCTLASS_DEVICE + TileIteratorPlanarComplex & operator-=(TensorCoord const &tile_offset) { + tile_iterator_.add_tile_offset(-tile_offset); + return *this; + } + + /// Loads a fragment from memory at the location pointed to by the iterator. + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + tile_iterator_.load_with_byte_offset(frag.real, 0); + tile_iterator_.load_with_byte_offset(frag.imag, imaginary_offset_); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset in units of bytes + Index byte_offset) const { + + tile_iterator_.load_with_byte_offset(frag.real, byte_offset); + tile_iterator_.load_with_byte_offset(frag.imag, byte_offset + imaginary_offset_); + } + + /// Loads a fragment from memory with additional logical offset + MCTLASS_DEVICE + void load_with_pointer_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a linear offset + Index pointer_offset) const { + + Index byte_offset = (pointer_offset * sizeof_bits::value)/8; + + tile_iterator_.load_with_byte_offset(frag.real, byte_offset); + tile_iterator_.load_with_byte_offset(frag.imag, byte_offset + imaginary_offset_); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset) const { + + tile_iterator_.load_with_byte_offset(frag.real, tile_offset, 0); + tile_iterator_.load_with_byte_offset(frag.imag, tile_offset, imaginary_offset_); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index pointer_offset) const { + + Index byte_offset = (pointer_offset * sizeof_bits::value)/8; + + tile_iterator_.load_with_byte_offset(frag.real, tile_offset, byte_offset); + tile_iterator_.load_with_byte_offset(frag.real, tile_offset, byte_offset + imaginary_offset_); + } + + /// Loads a fragment from memory with logical offset in units of whole tiles. + MCTLASS_DEVICE + void load_with_byte_offset( + /// fragment to load from the tensor + Fragment &frag, + /// loads a tile with a logical offset in units of whole tiles + TensorCoord const &tile_offset, + /// loads a tile with a logical offset AND a pointer offset + Index byte_offset) const { + + tile_iterator_.load_with_byte_offset(frag.real, tile_offset, byte_offset); + tile_iterator_.load_with_byte_offset(frag.imag, tile_offset, byte_offset + imaginary_offset_); + } + + /// Notify the iterator which k-group it is currently pointing to. + /// + /// This does not advance the iterator. Rather, it overrides its internal + /// tracking with constant-valued k-group index to enable the compiler to + /// fold constants and achieve more efficient code. + /// + /// This is used by some nontrivial permuted layouts. + MCTLASS_DEVICE + void set_kgroup_index(int k_group) { + tile_iterator_.set_kgroup_index(k_group); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace gemm +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/half.h b/csrc/mctlass/include/mctlass/half.h new file mode 100644 index 0000000..4c779b5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/half.h @@ -0,0 +1,918 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Defines a class for using IEEE half-precision floating-point types in host or + device code. +*/ +#pragma once + +#ifndef MCTLASS_ENABLE_F16C +#define MCTLASS_ENABLE_F16C 0 +#endif + +#if defined(__MACACC_RTC__) + +#include "mctlass/floating_point_nvrtc.h" + +// F16C extensions are not meaningful when compiling for NVRTC which only accommodates device code. +#undef MCTLASS_ENABLE_F16C +#define MCTLASS_ENABLE_F16C 0 + +#else +// +// Standard Library headers belong here to avoid conflicts with NVRTC. +// +#include +#include +#include +#include +#endif + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +#include + +#include "mctlass/mctlass.h" +#include "mctlass/float8.h" +#include "mctlass/platform/platform.h" + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// Optionally target F16C extentions to accelerate half-precision conversion. +#if !defined(__MACA_ARCH__) && (MCTLASS_ENABLE_F16C) +#if defined(_MSC_VER) + +#include + +#if defined(__i386__) || defined(__x86_64__) +#include +#endif + +#define F16C_ROUND_NEAREST 0 + +#if !defined(__MACA_ARCH__) +extern __inline float _cvtsh_ss (unsigned short __S) { + __m128i packed; + std::memcpy(&packed, &__S, sizeof(__S)); + + __m128 result = _mm_cvtph_ps(packed); + + float flt; + std::memcpy(&flt, &result, sizeof(flt)); + + return flt; +} + +__inline unsigned short _cvtss_sh (float __F, const int) { + __m128 packed; + std::memcpy(&packed, &__F, sizeof(__F)); + + __m128i result = _mm_cvtps_ph(packed, F16C_ROUND_NEAREST); + + unsigned short u; + std::memcpy(&u, &result, sizeof(u)); + + return u; +} +#endif + +#else + +// Linux +#include + +#if defined(__i386__) || defined(__x86_64__) +#include +#endif + +#define F16C_ROUND_NEAREST (_MM_FROUND_TO_NEAREST_INT |_MM_FROUND_NO_EXC) + +#endif // _MSC_VER + +class CpuId { + + bool f16c_enabled; + + CpuId() { + #if defined(__i386__) || defined(__x86_64__) + #if defined(_MSC_VER) + int exx[4]; + + __cpuid (exx, 1); + f16c_enabled = exx[2] & 0x20000000; + + #else + // GCC / Clang + int eax, ebx, ecx, edx; + + __cpuid (1 , eax, ebx, ecx, edx); + f16c_enabled = ecx & 0x20000000; + #endif + #else + // Arm / PowerPC etc. + f16c_enabled = false; + #endif + } + +public: + + bool is_f16c_supported() const { + return f16c_enabled; + } + + static const CpuId& instance() { + static CpuId cpu; + return cpu; + } +}; +#endif // !defined(__MACA_ARCH__) && MCTLASS_ENABLE_F16C + +/////////////////////////////////////////////////////////////////////////////////////////////////// + + +namespace mctlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// IEEE half-precision floating-point type +struct alignas(2) half_t { + + // + // Data members + // + + /// Storage type + uint16_t storage; + + // + // Static conversion operators + // + + /// Constructs from an unsigned short + MCTLASS_HOST_DEVICE + static half_t bitcast(uint16_t x) { + half_t h; + h.storage = x; + return h; + } + + /// FP32 -> FP16 conversion - rounds to nearest even + MCTLASS_HOST_DEVICE + static half_t convert(float const& flt) { + #if defined(__MACA_ARCH__) + //float2half_rn is equals with float2half + return half_t(__float2half(flt)); + #else + + // #if !defined(__MACA_ARCH__) && MCTLASS_ENABLE_F16C + // if( CpuId::instance().is_f16c_supported() ) { + // unsigned short u = _cvtss_sh(flt, F16C_ROUND_NEAREST); + // return bitcast(u); + // } + // #endif + + // software implementation rounds toward nearest even + unsigned s; + + #if defined(__MACA_ARCH__) + s = reinterpret_cast(flt); + #else + std::memcpy(&s, &flt, sizeof(s)); + #endif + + uint16_t sign = uint16_t((s >> 16) & 0x8000); + int16_t exp = uint16_t(((s >> 23) & 0xff) - 127); + int mantissa = s & 0x7fffff; + uint16_t u = 0; + + if ((s & 0x7fffffff) == 0) { + // sign-preserving zero + return bitcast(sign); + } + + if (exp > 15) { + if (exp == 128 && mantissa) { + // not a number + u = 0x7fff; + } else { + // overflow to infinity + u = sign | 0x7c00; + } + return bitcast(u); + } + + int sticky_bit = 0; + + if (exp >= -14) { + // normal fp32 to normal fp16 + exp = uint16_t(exp + uint16_t(15)); + u = uint16_t(((exp & 0x1f) << 10)); + u = uint16_t(u | (mantissa >> 13)); + } else { + // normal single-precision to subnormal half_t-precision representation + int rshift = (-14 - exp); + if (rshift < 32) { + mantissa |= (1 << 23); + + sticky_bit = ((mantissa & ((1 << rshift) - 1)) != 0); + + mantissa = (mantissa >> rshift); + u = (uint16_t(mantissa >> 13) & 0x3ff); + } else { + mantissa = 0; + u = 0; + } + } + + // round to nearest even + int round_bit = ((mantissa >> 12) & 1); + sticky_bit |= ((mantissa & ((1 << 12) - 1)) != 0); + + if ((round_bit && sticky_bit) || (round_bit && (u & 1))) { + u = uint16_t(u + 1); + } + + u |= sign; + + return bitcast(u); + #endif + } + + /// FP32 -> FP16 conversion - rounds to nearest even + MCTLASS_HOST_DEVICE + static half_t convert(int const& n) { + #if defined(__MACA_ARCH__) + return half_t(__int2half_rn(n)); + #else + return convert(float(n)); + #endif + } + + /// FP32 -> FP16 conversion - rounds to nearest even + MCTLASS_HOST_DEVICE + static half_t convert(unsigned const& n) { + #if defined(__MACA_ARCH__) + return half_t(__uint2half_rn(n)); + #else + return convert(float(n)); + #endif + } + + /// Converts a half-precision value stored as a uint16_t to a float + MCTLASS_HOST_DEVICE + static float convert(half_t const& x) { + #if defined(__MACA_ARCH__) + return __half2float(x.to_half()); + #else + + // #if !defined(__MACA_ARCH__) && MCTLASS_ENABLE_F16C + // if( CpuId::instance().is_f16c_supported() ) { + // unsigned short u = x.storage; + // return _cvtsh_ss(u); + // } + // #endif + + uint16_t const &h = x.storage; + int sign = ((h >> 15) & 1); + int exp = ((h >> 10) & 0x1f); + int mantissa = (h & 0x3ff); + unsigned f = 0; + + if (exp > 0 && exp < 31) { + // normal + exp += 112; + f = (sign << 31) | (exp << 23) | (mantissa << 13); + } else if (exp == 0) { + if (mantissa) { + // subnormal + exp += 113; + while ((mantissa & (1 << 10)) == 0) { + mantissa <<= 1; + exp--; + } + mantissa &= 0x3ff; + f = (sign << 31) | (exp << 23) | (mantissa << 13); + } else { + // sign-preserving zero + f = (sign << 31); + } + } else if (exp == 31) { + if (mantissa) { + f = 0x7fffffff; // not a number + } else { + f = (0xff << 23) | (sign << 31); // inf + } + } + #if defined(__MACA_ARCH__) + return reinterpret_cast(f); + #else + float flt; + std::memcpy(&flt, &f, sizeof(flt)); + return flt; + #endif + #endif + } + + // + // Methods + // + + /// Default constructor + MCTLASS_HOST_DEVICE + half_t() : storage(0) { } + + /// Reinterpret cast from MACA's half type + MCTLASS_HOST_DEVICE + explicit half_t(half const & x) { + #if defined(__MACA_ARCH__) + storage = reinterpret_cast(x); + #else + __half_raw raw(x); + std::memcpy(&storage, &raw.x, sizeof(storage)); + #endif + } + + /// Floating point conversion + MCTLASS_HOST_DEVICE + explicit half_t(float x) { + storage = convert(x).storage; + } + + /// Floating point conversion + MCTLASS_HOST_DEVICE + explicit half_t(double x): half_t(float(x)) { + + } + + /// float_e4m3_t conversion + MCTLASS_HOST_DEVICE + explicit half_t(float_e4m3_t x): half_t(float(x)) { + + } + + /// float_e5m2_t conversion + MCTLASS_HOST_DEVICE + explicit half_t(float_e5m2_t x): half_t(float(x)) { + + } + + /// Integer conversion - round to nearest even + MCTLASS_HOST_DEVICE + explicit half_t(int x) { + storage = convert(x).storage; + } + + /// Integer conversion - round toward zero + MCTLASS_HOST_DEVICE + explicit half_t(unsigned x) { + storage = convert(x).storage; + } + + /// Assignment + MCTLASS_HOST_DEVICE + half_t & operator=(half const &x) { + #if defined(__MACA_ARCH__) + storage = reinterpret_cast(x); + #else + __half_raw raw(x); + std::memcpy(&storage, &raw.x, sizeof(storage)); + #endif + return *this; + } + + /// Converts to float + MCTLASS_HOST_DEVICE + operator float() const { + return convert(*this); + } + + /// Converts to float + MCTLASS_HOST_DEVICE + explicit operator double() const { + return double(convert(*this)); + } + + /// Converts to float + MCTLASS_HOST_DEVICE + explicit operator int() const { + return int(convert(*this)); + } + + /// Casts to bool + MCTLASS_HOST_DEVICE + explicit operator bool() const { + return (convert(*this) != 0.0f); + } + + /// Bitcasts to MACA's half type + MCTLASS_HOST_DEVICE + half to_half() const { + #if defined(__MACA_ARCH__) + return reinterpret_cast(storage); + #else + __half_raw raw; + std::memcpy(&raw.x, &storage, sizeof(raw.x)); + return half(raw); + #endif + } + + /// Bitcasts to MACA's half type + #if defined(__MACA_ARCH__) + MCTLASS_HOST_DEVICE + const __fp16& to_macahalf() const { + return reinterpret_cast<__fp16 const &>(storage); + } + #endif + + /// Accesses raw internal state + MCTLASS_HOST_DEVICE + uint16_t& raw() { + return storage; + } + + /// Accesses raw internal state + MCTLASS_HOST_DEVICE + uint16_t raw() const { + return storage; + } + + /// Returns the sign bit + MCTLASS_HOST_DEVICE + bool signbit() const { + return ((storage & 0x8000) != 0); + } + + /// Returns the biased exponent + MCTLASS_HOST_DEVICE + int exponent_biased() const { + return int((storage >> 10) & 0x1f); + } + + /// Returns the unbiased exponent + MCTLASS_HOST_DEVICE + int exponent() const { + return exponent_biased() - 15; + } + + /// Returns the mantissa + MCTLASS_HOST_DEVICE + int mantissa() const { + return int(storage & 0x3ff); + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +MCTLASS_HOST_DEVICE +bool signbit(mctlass::half_t const& h) { + return ((h.raw() & 0x8000) != 0); +} + +MCTLASS_HOST_DEVICE +mctlass::half_t abs(mctlass::half_t const& h) { + return mctlass::half_t::bitcast(h.raw() & 0x7fff); +} + +MCTLASS_HOST_DEVICE +bool isnan(mctlass::half_t const& h) { + return (h.exponent_biased() == 0x1f) && h.mantissa(); +} + +MCTLASS_HOST_DEVICE +bool isfinite(mctlass::half_t const& h) { + return (h.exponent_biased() != 0x1f); +} + +MCTLASS_HOST_DEVICE +mctlass::half_t nanh(const char*) { + // NVIDIA canonical NaN + return mctlass::half_t::bitcast(0x7fff); +} + +MCTLASS_HOST_DEVICE +bool isinf(mctlass::half_t const& h) { + return (h.exponent_biased() == 0x1f) && !h.mantissa(); +} + +MCTLASS_HOST_DEVICE +bool isnormal(mctlass::half_t const& h) { + return h.exponent_biased() && h.exponent_biased() != 0x1f; +} + +MCTLASS_HOST_DEVICE +int fpclassify(mctlass::half_t const& h) { + int exp = h.exponent_biased(); + int mantissa = h.mantissa(); + if (exp == 0x1f) { + if (mantissa) { + return FP_NAN; + } + else { + return FP_INFINITE; + } + } + else if (!exp) { + if (mantissa) { + return FP_SUBNORMAL; + } + else { + return FP_ZERO; + } + } + return FP_NORMAL; +} + +MCTLASS_HOST_DEVICE +mctlass::half_t sqrt(mctlass::half_t const& h) { +#if defined(__MACACC_RTC__) + return mctlass::half_t(sqrtf(float(h))); +#else + return mctlass::half_t(std::sqrt(float(h))); +#endif +} + +MCTLASS_HOST_DEVICE +half_t copysign(half_t const& a, half_t const& b) { + + uint16_t a_mag = (a.raw() & 0x7fff); + uint16_t b_sign = (b.raw() & 0x8000); + uint16_t result = (a_mag | b_sign); + + return half_t::bitcast(result); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Standard Library operations and definitions +// +/////////////////////////////////////////////////////////////////////////////////////////////////// + +#if !defined(__MACACC_RTC__) +namespace std { + +/// Numeric limits +template <> +struct numeric_limits { + static bool const is_specialized = true; + static bool const is_signed = true; + static bool const is_integer = false; + static bool const is_exact = false; + static bool const has_infinity = true; + static bool const has_quiet_NaN = true; + static bool const has_signaling_NaN = false; + static std::float_denorm_style const has_denorm = std::denorm_present; + static bool const has_denorm_loss = true; + static std::float_round_style const round_style = std::round_to_nearest; + static bool const is_iec559 = true; + static bool const is_bounded = true; + static bool const is_modulo = false; + static int const digits = 10; + + /// Least positive value + static mctlass::half_t min() { return mctlass::half_t::bitcast(0x0001); } + + /// Minimum finite value + static mctlass::half_t lowest() { return mctlass::half_t::bitcast(0xfbff); } + + /// Maximum finite value + static mctlass::half_t max() { return mctlass::half_t::bitcast(0x7bff); } + + /// Returns smallest finite value + static mctlass::half_t epsilon() { return mctlass::half_t::bitcast(0x1800); } + + /// Returns maximum rounding error + static mctlass::half_t round_error() { return mctlass::half_t(0.5f); } + + /// Returns positive infinity value + static mctlass::half_t infinity() { return mctlass::half_t::bitcast(0x7c00); } + + /// Returns quiet NaN value + static mctlass::half_t quiet_NaN() { return mctlass::half_t::bitcast(0x7fff); } + + /// Returns signaling NaN value + static mctlass::half_t signaling_NaN() { return mctlass::half_t::bitcast(0x7fff); } + + /// Returns smallest positive subnormal value + static mctlass::half_t denorm_min() { return mctlass::half_t::bitcast(0x0001); } +}; +} // namespace std +#endif + +namespace platform { + +/// std::numeric_limits +template +struct numeric_limits; + +/// Numeric limits +template <> +struct numeric_limits { + static bool const is_specialized = true; + static bool const is_signed = true; + static bool const is_integer = false; + static bool const is_exact = false; + static bool const has_infinity = true; + static bool const has_quiet_NaN = true; + static bool const has_signaling_NaN = false; +#if !defined(__MACACC_RTC__) + static std::float_denorm_style const has_denorm = std::denorm_present; +#endif + static bool const has_denorm_loss = true; +#if !defined(__MACACC_RTC__) + static std::float_round_style const round_style = std::round_to_nearest; +#endif + static bool const is_iec559 = true; + static bool const is_bounded = true; + static bool const is_modulo = false; + static int const digits = 10; + + /// Least positive value + MCTLASS_HOST_DEVICE + static mctlass::half_t min() { return mctlass::half_t::bitcast(0x0001); } + + /// Minimum finite value + MCTLASS_HOST_DEVICE + static mctlass::half_t lowest() { return mctlass::half_t::bitcast(0xfbff); } + + /// Maximum finite value + MCTLASS_HOST_DEVICE + static mctlass::half_t max() { return mctlass::half_t::bitcast(0x7bff); } + + /// Returns smallest finite value + MCTLASS_HOST_DEVICE + static mctlass::half_t epsilon() { return mctlass::half_t::bitcast(0x1800); } + + /// Returns maximum rounding error + MCTLASS_HOST_DEVICE + static mctlass::half_t round_error() { return mctlass::half_t(0.5f); } + + /// Returns positive infinity value + MCTLASS_HOST_DEVICE + static mctlass::half_t infinity() { return mctlass::half_t::bitcast(0x7c00); } + + /// Returns quiet NaN value + MCTLASS_HOST_DEVICE + static mctlass::half_t quiet_NaN() { return mctlass::half_t::bitcast(0x7fff); } + + /// Returns signaling NaN value + MCTLASS_HOST_DEVICE + static mctlass::half_t signaling_NaN() { return mctlass::half_t::bitcast(0x7fff); } + + /// Returns smallest positive subnormal value + MCTLASS_HOST_DEVICE + static mctlass::half_t denorm_min() { return mctlass::half_t::bitcast(0x0001); } +}; +} // namespace platform + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Arithmetic operators +// +/////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// +MCTLASS_HOST_DEVICE +bool operator==(half_t const& lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + return __heq(lhs.to_half(), rhs.to_half()); +#else + return float(lhs) == float(rhs); +#endif +} + +MCTLASS_HOST_DEVICE +bool operator!=(half_t const& lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + return __hne(lhs.to_half(), rhs.to_half()); +#else + return float(lhs) != float(rhs); +#endif +} + +MCTLASS_HOST_DEVICE +bool operator<(half_t const& lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + return __hlt(lhs.to_half(), rhs.to_half()); +#else + return float(lhs) < float(rhs); +#endif +} + +MCTLASS_HOST_DEVICE +bool operator<=(half_t const& lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + return __hle(lhs.to_half(), rhs.to_half()); +#else + return float(lhs) <= float(rhs); +#endif +} + +MCTLASS_HOST_DEVICE +bool operator>(half_t const& lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + return __hgt(lhs.to_half(), rhs.to_half()); +#else + return float(lhs) > float(rhs); +#endif +} + +MCTLASS_HOST_DEVICE +bool operator>=(half_t const& lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + return __hge(lhs.to_half(), rhs.to_half()); +#else + return float(lhs) >= float(rhs); +#endif +} + +MCTLASS_HOST_DEVICE +half_t operator+(half_t const& lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + return half_t(__hadd(lhs.to_half(), rhs.to_half())); +#else + return half_t(float(lhs) + float(rhs)); +#endif +} + +MCTLASS_HOST_DEVICE +half_t operator-(half_t const& lhs) { +#if defined(__MACA_ARCH__) + return half_t(__hneg(lhs.to_half())); +#else + return half_t(-float(lhs)); +#endif +} + +MCTLASS_HOST_DEVICE +half_t operator-(half_t const& lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + return half_t(__hsub(lhs.to_half(), rhs.to_half())); +#else + return half_t(float(lhs) - float(rhs)); +#endif +} + +MCTLASS_HOST_DEVICE +half_t operator*(half_t const& lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + return half_t(__hmul(lhs.to_half(), rhs.to_half())); +#else + return half_t(float(lhs) * float(rhs)); +#endif +} + +MCTLASS_HOST_DEVICE +half_t operator/(half_t const& lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + return half_t(__hdiv(lhs.to_half(), rhs.to_half())); +#else + return half_t(float(lhs) / float(rhs)); +#endif +} + +MCTLASS_HOST_DEVICE +half_t& operator+=(half_t & lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + lhs = half_t(__hadd(lhs.to_half(), rhs.to_half())); +#else + lhs = half_t(float(lhs) + float(rhs)); +#endif + return lhs; +} + +MCTLASS_HOST_DEVICE +half_t& operator-=(half_t & lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + lhs = half_t(__hsub(lhs.to_half(), rhs.to_half())); +#else + lhs = half_t(float(lhs) - float(rhs)); +#endif + return lhs; +} + +MCTLASS_HOST_DEVICE +half_t& operator*=(half_t & lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + lhs = half_t(__hmul(lhs.to_half(), rhs.to_half())); +#else + lhs = half_t(float(lhs) * float(rhs)); +#endif + return lhs; +} + +MCTLASS_HOST_DEVICE +half_t& operator/=(half_t & lhs, half_t const& rhs) { +#if defined(__MACA_ARCH__) + lhs = half_t(__hdiv(lhs.to_half(), rhs.to_half())); +#else + lhs = half_t(float(lhs) / float(rhs)); +#endif + return lhs; +} + +MCTLASS_HOST_DEVICE +half_t& operator++(half_t & lhs) { +#if defined(__MACA_ARCH__) + lhs = half_t(__hadd(lhs.to_half(), half_t(1.0f).to_half())); +#else + float tmp(lhs); + ++tmp; + lhs = half_t(tmp); +#endif + return lhs; +} + +MCTLASS_HOST_DEVICE +half_t& operator--(half_t & lhs) { +#if defined(__MACA_ARCH__) + lhs = half_t(__hsub(lhs.to_half(), half_t(1.0f).to_half())); +#else + float tmp(lhs); + --tmp; + lhs = half_t(tmp); +#endif + return lhs; +} + +MCTLASS_HOST_DEVICE +half_t operator++(half_t & lhs, int) { + half_t ret(lhs); +#if defined(__MACA_ARCH__) + lhs = half_t(__hadd(lhs.to_half(), half_t(1.0f).to_half())); +#else + float tmp(lhs); + tmp++; + lhs = half_t(tmp); +#endif + return ret; +} + +MCTLASS_HOST_DEVICE +half_t operator--(half_t & lhs, int) { + half_t ret(lhs); +#if defined(__MACA_ARCH__) + lhs = half_t(__hsub(lhs.to_half(), half_t(1.0f).to_half())); +#else + float tmp(lhs); + tmp--; + lhs = half_t(tmp); +#endif + return ret; +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// +// User-defined literals +// + +MCTLASS_HOST_DEVICE +mctlass::half_t operator "" _hf(long double x) { + return mctlass::half_t(float(x)); +} + +MCTLASS_HOST_DEVICE +mctlass::half_t operator "" _hf(unsigned long long int x) { + return mctlass::half_t(int(x)); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/integer_subbyte.h b/csrc/mctlass/include/mctlass/integer_subbyte.h new file mode 100644 index 0000000..8061c0b --- /dev/null +++ b/csrc/mctlass/include/mctlass/integer_subbyte.h @@ -0,0 +1,241 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Defines a class for using integer types smaller than one byte in host or + device code. +*/ +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#include "mctlass/platform/platform.h" + +namespace mctlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// 4-bit signed integer type +template +struct integer_subbyte { + + /// Number of bits + static int const kBits = Bits; + + /// Whether type is signed + static bool const kSigned = Signed; + + /// External type + using T = typename platform::conditional::type; + + /// Storage type + using Storage = uint8_t; + + /// Bitmask used to truncate from larger integers + static Storage const kMask = Storage((1 << kBits) - 1); + + // + // Data members + // + + Storage storage; + + // + // Methods + // + + /// No operation + MCTLASS_HOST_DEVICE + integer_subbyte() = default; + + /// Conversion from integer type + MCTLASS_HOST_DEVICE + integer_subbyte(int value) + : storage(reinterpret_cast(value) & kMask) {} + + MCTLASS_HOST_DEVICE + integer_subbyte(unsigned value) + : storage(reinterpret_cast(value) & kMask) {} + + MCTLASS_HOST_DEVICE + integer_subbyte(double value) { + T tmp = static_cast(value); + storage = Storage(reinterpret_cast(tmp) & kMask); + } + + /// + MCTLASS_HOST_DEVICE + operator T() const { + if (kSigned) { + // Sign extend + if (storage & Storage(1 << (kBits - 1))) { + return T(storage) | ~T(kMask); + } + } + return T(storage); + } + + /// Equality + MCTLASS_HOST_DEVICE + bool operator==(integer_subbyte const &rhs) const { + return storage == rhs.storage; + } + + /// Inequality + MCTLASS_HOST_DEVICE + bool operator!=(integer_subbyte const &rhs) const { + return storage != rhs.storage; + } + + /// Less than or equal + MCTLASS_HOST_DEVICE + bool operator<=(integer_subbyte const &rhs) const { + if (kSigned) { + if (storage & (1 << (kBits - 1))) { + return !(rhs.storage < storage); + } + } + return storage < rhs.storage; + } + + /// Less than + MCTLASS_HOST_DEVICE + bool operator<(integer_subbyte const &rhs) const { + if (kSigned) { + if (storage & (1 << (kBits - 1))) { + return !(rhs.storage <= storage); + } + } + return storage < rhs.storage; + } + + /// Greater than or equal + MCTLASS_HOST_DEVICE + bool operator>=(integer_subbyte const &rhs) const { + return !(*this < rhs); + } + + /// Greater than + MCTLASS_HOST_DEVICE + bool operator>(integer_subbyte const &rhs) const { + return !(*this <= rhs); + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + + +/// 1-bit Unsigned integer type +using uint1b_t = integer_subbyte<1, false>; + +/// 2-bit Integer type +using int2b_t = integer_subbyte<2, true>; + +/// 2-bit Unsigned integer type +using uint2b_t = integer_subbyte<2, false>; + +/// 4-bit Integer type +using int4b_t = integer_subbyte<4, true>; + +/// 4-bit Unsigned integer type +using uint4b_t = integer_subbyte<4, false>; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines the size of an element in bits - specialized for uint1b_t +template <> +struct sizeof_bits { + static int const value = 1; +}; + +/// Defines the size of an element in bits - specialized for int2b_t +template <> +struct sizeof_bits { + static int const value = 2; +}; + +/// Defines the size of an element in bits - specialized for uint2b_t +template <> +struct sizeof_bits { + static int const value = 2; +}; + +/// Defines the size of an element in bits - specialized for int4b_t +template <> +struct sizeof_bits { + static int const value = 4; +}; + +/// Defines the size of an element in bits - specialized for uint4b_t +template <> +struct sizeof_bits { + static int const value = 4; +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace platform { + +template <> +struct numeric_limits { + MCTLASS_HOST_DEVICE + static mctlass::int4b_t const lowest() noexcept { return -8;} + MCTLASS_HOST_DEVICE + static mctlass::int4b_t const max() noexcept { return 7;} + static constexpr bool is_integer = true; +}; + +template <> +struct numeric_limits { + MCTLASS_HOST_DEVICE + static mctlass::uint4b_t const lowest() noexcept { return 0;} + MCTLASS_HOST_DEVICE + static mctlass::uint4b_t const max() noexcept { return 15;} + static constexpr bool is_integer = true; +}; + +template <> +struct numeric_limits { + MCTLASS_HOST_DEVICE + static mctlass::uint1b_t const lowest() noexcept { return 0;} + MCTLASS_HOST_DEVICE + static mctlass::uint1b_t const max() noexcept { return 1;} + static constexpr bool is_integer = true; +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace platform +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/kernel_hardware_info.hpp b/csrc/mctlass/include/mctlass/kernel_hardware_info.hpp new file mode 100644 index 0000000..eb3aa63 --- /dev/null +++ b/csrc/mctlass/include/mctlass/kernel_hardware_info.hpp @@ -0,0 +1,73 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mc_runtime.h" +#include "mc_runtime_types.h" + +#include "mctlass/trace.h" + +namespace mctlass { + +struct KernelHardwareInfo { + // + // Data members + // + int device_id = 0; + int sm_count = 0; + + // + // Methods + // + + static int + query_device_multiprocessor_count(int device_id = 0) { + mcError_t result = mcGetDevice(&device_id); + if (result != mcSuccess) { + MCTLASS_TRACE_HOST( + " mcGetDevice() returned error " + << mcGetErrorString(result)); + return 0; + } + int multiprocessor_count; + result = mcDeviceGetAttribute(&multiprocessor_count, + mcDeviceAttributeMultiProcessorCount, device_id); + if (result != mcSuccess) { + MCTLASS_TRACE_HOST( + " mcDeviceGetAttribute() returned error " + << mcGetErrorString(result)); + return 0; + } + return multiprocessor_count; + } +}; + +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/kernel_launch.h b/csrc/mctlass/include/mctlass/kernel_launch.h new file mode 100644 index 0000000..dad420d --- /dev/null +++ b/csrc/mctlass/include/mctlass/kernel_launch.h @@ -0,0 +1,73 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines structures and helpers to launch MACA kernels within MCTLASS. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +namespace mctlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Structure containing the basic launch configuration of a CUDA kernel. +struct KernelLaunchConfiguration { + + /// CUDA grid dimensions + dim3 grid; + + /// CUDA threablock dimensions + dim3 block; + + /// Bytes of dynamically allocated SMEM in addition to static SMEM + size_t dynamic_smem; + + // + // Methods + // + + /// Constructs a KernellaunchConfiguration object + MCTLASS_HOST_DEVICE + KernelLaunchConfiguration( + dim3 _grid = dim3(1,1,1), + dim3 _block = dim3(1,1,1), + size_t _dynamic_smem = 0 + ): + grid(_grid), + block(_block), + dynamic_smem(_dynamic_smem) { } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/layout/layout.h b/csrc/mctlass/include/mctlass/layout/layout.h new file mode 100644 index 0000000..d91f7f6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/layout/layout.h @@ -0,0 +1,64 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines layout functions used by TensorRef and derived classes. + + Layout functions map logical coordinates to linear memory. They often require additional + data to describe strides between elements. + + Layout functions must implement all members in the public interface of IdentityTensorLayout<> + defined in mctlass/tensor_ref.h. +*/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/layout/vector.h" + +#include "mctlass/layout/tensor_op_multiplicand_sm70.h" +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" +/////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace layout { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace layout +} // namespace mctlass + +/////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/layout/matrix.h b/csrc/mctlass/include/mctlass/layout/matrix.h new file mode 100644 index 0000000..27c868e --- /dev/null +++ b/csrc/mctlass/include/mctlass/layout/matrix.h @@ -0,0 +1,1365 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines layout functions used by TensorRef and derived classes. + + Layout functions map logical coordinates to linear memory. They often require additional + data to describe strides between elements. + + Layout functions must implement all members in the public interface of IdentityTensorLayout<> + defined in mctlass/tensor_ref.h. +*/ +#pragma once + +#include "cute/layout.hpp" + +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/pitch_linear_coord.h" + +namespace mctlass { +namespace layout { + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Defines data layouts of various matrix formats usable by TensorRef and other classes. +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Mapping function for row-major matrices. +class RowMajor { +public: + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + +private: + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + RowMajor(LongIndex ldm = 0): stride_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajor(Stride stride): stride_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static RowMajor packed(MatrixCoord const &extent) { + return RowMajor(extent.column()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord const &coord) const { + return LongIndex(coord.row()) * LongIndex(stride_[0]) + coord.column(); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + MatrixCoord inverse(LongIndex offset) const { + return MatrixCoord(Index(offset / stride_[0]), Index(offset % stride_[0])); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const { + return stride_[idx]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index & stride(int idx) { + return stride_[idx]; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(MatrixCoord const &extent) const { + return LongIndex(extent.row()) * LongIndex(stride_[0]); + } + + MCTLASS_HOST_DEVICE + cute::Layout, cute::Stride > > + to_cute_layout(MatrixCoord const &extent) const { + return cute::Layout, cute::Stride > >{ + {extent[0], extent[1]}, + {stride(0), cute::Int<1>{}} + }; + } +}; + +/// Mapping function for column-major matrices. +class ColumnMajor { +public: + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + +private: + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajor(LongIndex ldm = 0): stride_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajor(Stride stride): stride_(stride) { } + + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static ColumnMajor packed(MatrixCoord const &extent) { + return ColumnMajor(extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord const &coord) const { + return LongIndex(coord.column()) * LongIndex(stride_[0]) + coord.row(); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + MatrixCoord inverse(LongIndex offset) const { + return MatrixCoord(Index(offset % stride_[0]), Index(offset / stride_[0])); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const { + return stride_[idx]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index & stride(int idx) { + return stride_[idx]; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(MatrixCoord const &extent) const { + return LongIndex(extent.column()) * LongIndex(stride_[0]); + } + + MCTLASS_HOST_DEVICE + cute::Layout, cute::Stride< cute::Int<1>, int64_t> > + to_cute_layout(MatrixCoord const &extent) const { + return cute::Layout, cute::Stride, int64_t> >{ + {extent[0], extent[1]}, + {cute::Int<1>{}, stride(0)} + }; + } +}; + +/// Mapping function for interleaved matrices. Matrix is structured +/// as row-major arrangement of fixed-size columns. +template +struct RowMajorInterleaved { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + /// Size of interleaved columns + static int const kInterleave = Interleave; + +private: + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorInterleaved(LongIndex ldm = 0): stride_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorInterleaved(Stride stride): stride_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static RowMajorInterleaved packed(MatrixCoord const &extent) { + return RowMajorInterleaved(extent.column() * kInterleave); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord const &coord) const { + Index row_major = coord.row() / kInterleave; + Index row_minor = coord.row() % kInterleave; + return LongIndex(row_major) * LongIndex(stride_[0]) + LongIndex(coord.column()) * kInterleave + row_minor; + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + MatrixCoord inverse(LongIndex offset) const { + + Index row_major = Index(offset / stride_[0]); + Index residual = Index(offset % stride_[0]); + + Index column = residual / kInterleave; + Index row_minor = residual % kInterleave; + + return MatrixCoord(row_major * kInterleave + row_minor, column); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const { + return stride_[idx]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index & stride(int idx) { + return stride_[idx]; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(MatrixCoord const &extent) const { + return (extent.row() + kInterleave - 1) / kInterleave * stride_[0]; + } +}; + +/// Mapping function for interleaved matrices. Matrix is structured +/// as column-major arrangement of fixed-size rows. +template +struct ColumnMajorInterleaved { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + /// Size of interleaved columns + static int const kInterleave = Interleave; + +private: + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorInterleaved(LongIndex ldm = 0): stride_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorInterleaved(Stride stride): stride_(stride) { } + + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static ColumnMajorInterleaved packed(MatrixCoord const &extent) { + return ColumnMajorInterleaved(extent.row() * kInterleave); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord const &coord) const { + Index column_major = coord.column() / kInterleave; + Index column_minor = coord.column() % kInterleave; + return LongIndex(column_major) * LongIndex(stride_[0]) + LongIndex(coord.row()) * kInterleave + column_minor; + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + MatrixCoord inverse(LongIndex offset) const { + + Index column_major = Index(offset / stride_[0]); + Index residual = Index(offset % stride_[0]); + + Index row = residual / kInterleave; + Index column_minor = residual % kInterleave; + + return MatrixCoord(row, column_major * kInterleave + column_minor); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const { + return stride_[idx]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index & stride(int idx) { + return stride_[idx]; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(MatrixCoord const &extent) const { + return (extent.column() + kInterleave - 1) / kInterleave * stride_[0]; + } +}; + +/// Enumerated type for canonical pitch-linear matrix layouts +enum class Matrix { + kColumnMajor, ///< leading dimension refers to stride between columns; stride along rows is 1 + kRowMajor ///< leading dimension refers to stride between rows; stride along columns is 1 +}; + +/// Mapping function for scenario in which layout is row-major or column-major but this information +/// is only available at runtime. +struct ContiguousMatrix { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + +private: + // + // Data members + // + + /// Stride data member + Stride stride_; + + /// Enumerated type indicating canonical matrix layout + Matrix layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + ContiguousMatrix( + Index ldm = 0, + Matrix layout = Matrix::kColumnMajor + ): + stride_(ldm), layout_(layout) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static ContiguousMatrix packed( + MatrixCoord const &extent, + Matrix layout = Matrix::kColumnMajor) { + + Index ldm = 0; + if (layout == Matrix::kColumnMajor) { + ldm = extent.row(); + } + else if (layout == Matrix::kRowMajor) { + ldm = extent.column(); + } + return ContiguousMatrix(ldm, layout); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord const &coord) const { + if (layout_ == Matrix::kColumnMajor) { + return coord.row() + coord.column() * stride_[0]; + } + else if (layout_ == Matrix::kRowMajor) { + return coord.row() * stride_[0] + coord.column(); + } + else { + // degenerate case + return 0; + } + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + MatrixCoord inverse(LongIndex offset) const { + return MatrixCoord(0, 0); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const { + return stride_[idx]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index & stride(int idx) { + return stride_[idx]; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(MatrixCoord const &extent) const { + if (layout_ == Matrix::kColumnMajor) { + return stride_[0] * extent.column(); + } + else if (layout_ == Matrix::kRowMajor) { + return stride_[0] * extent.row(); + } + else { + // degenerate case + return 0; + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Mapping function for scenario in which both rows and columns are separated by a stride. +template +struct AffineRankN { + + /// Logical rank of tensor + static int const kRank = Rank; + + /// Rank of stride vector + static int const kStrideRank = kRank; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = Coord; + + /// Stride vector + using Stride = Coord; + +private: + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + AffineRankN( + Stride const &stride = Stride() + ): + stride_(stride) { } + + /// Ctor + MCTLASS_HOST_DEVICE + AffineRankN( + Coord const &stride_m, + Coord const &stride_n + ) { + + // Concatenate the strides + MCTLASS_PRAGMA_UNROLL + for (int m = 0; m < kRank/2; ++m) { + stride_[m] = stride_m[m]; + } + + MCTLASS_PRAGMA_UNROLL + for (int n = 0; n < kRank/2; ++n) { + stride_[n + kRank/2] = stride_n[n]; + } + } + + /// Ctor for N = 2 + MCTLASS_HOST_DEVICE + AffineRankN( + LongIndex const &stride_m, + LongIndex const &stride_n + ) { + stride_[0] = stride_m; + stride_[1] = stride_n; + } + + /// Ctor for N = 2 + MCTLASS_HOST_DEVICE + AffineRankN( + LongIndex const &stride + ) { + stride_[0] = stride; + stride_[1] = 1; + } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static AffineRankN packed(TensorCoord const &extent) { + + AffineRankN layout; + layout.stride_[kRank - 1] = 1; + + MCTLASS_PRAGMA_UNROLL + for (int i = kRank - 1; i > 0; --i) { + layout.stride_[i - 1] = layout.stride_[i] * extent[i]; + } + + return layout; + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return dot(coord, stride_); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + return TensorCoord(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const { + return stride_[idx]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index & stride(int idx) { + return stride_[idx]; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + int idx = stride_.max_dim_index(); + return extent[idx] * stride_[idx]; + } +}; + +/// Mapping function for scenario in which both rows and columns are separated by a stride. +/// Row stride is smaller than column stride in AffineRank2ColumnMajor. +struct AffineRank2ColumnMajor { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 2; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + +private: + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + AffineRank2ColumnMajor( + Stride const &stride = Stride() + ): + stride_(stride) { } + + /// Ctor + MCTLASS_HOST_DEVICE + AffineRank2ColumnMajor( + LongIndex row_stride, ///< stride between elements in consecutive rows + LongIndex column_stride ///< stride between elements in consecutive columns + ) + { stride_[0] = row_stride; stride_[1] = column_stride;} + + /// Ctor + MCTLASS_HOST_DEVICE + AffineRank2ColumnMajor( + LongIndex stride + ) + { stride_[0] = 1; stride_[1] = stride;} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static AffineRank2ColumnMajor packed(MatrixCoord const &extent) { + return AffineRank2ColumnMajor(1, extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord const &coord) const { + return dot(coord, stride_); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + MatrixCoord inverse(LongIndex offset) const { + return MatrixCoord(0, 0); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const { + return stride_[idx]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index & stride(int idx) { + return stride_[idx]; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(MatrixCoord const &extent) const { + return extent.column() * stride_[1]; + } +}; + +/// Mapping function for scenario in which both rows and columns are separated by a stride. +/// Column stride is smaller than row stride in AffineRank2RowMajor. +struct AffineRank2RowMajor { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 2; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + +private: + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + AffineRank2RowMajor( + Stride const &stride = Stride() + ): + stride_(stride) { } + + /// Ctor + MCTLASS_HOST_DEVICE + AffineRank2RowMajor( + LongIndex row_stride, ///< stride between elements in consecutive rows + LongIndex column_stride ///< stride between elements in consecutive columns + ) { stride_[0] = row_stride; stride_[1] = column_stride;} + + /// Ctor + MCTLASS_HOST_DEVICE + AffineRank2RowMajor( + LongIndex stride + ) { stride_[0] = stride; stride_[1] = 1;} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static AffineRank2RowMajor packed(MatrixCoord const &extent) { + return AffineRank2RowMajor(1, extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord const &coord) const { + return dot(coord, stride_); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + MatrixCoord inverse(LongIndex offset) const { + return MatrixCoord(0, 0); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const { + return stride_[idx]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index & stride(int idx) { + return stride_[idx]; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(MatrixCoord const &extent) const { + return extent.row() * stride_[0]; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Utility functions to convert stride_factor to the strides used by the Affine2 layout. +// +// stride_factor is the logical distance between two coorinates. +// +// All Coodinates used here are matrix coordinates. stride[0] and extent[0] are for the +// rows. stride[1] and extent[1] are for the columns. +template + struct Affine2Layout_Factory { + MCTLASS_HOST_DEVICE + static Affine2Layout layout_factory(mctlass::Coord<2> const &extent, typename Affine2Layout::Stride stride_factor) { + return Affine2Layout::packed(extent); + } +}; + +template <> +struct Affine2Layout_Factory { +MCTLASS_HOST_DEVICE +static mctlass::layout::AffineRank2ColumnMajor layout_factory( + mctlass::Coord<2> const &extent, + typename mctlass::layout::AffineRank2ColumnMajor::Stride stride_factor) { + return mctlass::layout::AffineRank2ColumnMajor({ stride_factor[0], stride_factor[0] * stride_factor[1] * extent[0] }); + } +}; + +template <> +struct Affine2Layout_Factory { +MCTLASS_HOST_DEVICE +static mctlass::layout::AffineRank2RowMajor layout_factory( + mctlass::Coord<2> const &extent, + typename mctlass::layout::AffineRank2RowMajor::Stride stride_factor) { + return mctlass::layout::AffineRank2RowMajor({ stride_factor[0] * stride_factor[1] * extent[1], stride_factor[1] }); + } +}; + +// The base layout mctlass::layout::AffineRankN<2> is similar to AffineRank2ColumnMajor +template <> +struct Affine2Layout_Factory> { +MCTLASS_HOST_DEVICE +static mctlass::layout::AffineRankN<2> layout_factory( + mctlass::Coord<2> const &extent, + typename mctlass::layout::AffineRankN<2>::Stride stride_factor) { + return mctlass::layout::AffineRankN<2>({ stride_factor[0], stride_factor[0] * stride_factor[1] * extent[0] }); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Mapping function for block-linear matrices. Matrix is structured +/// as column-major arrangement of 2D tiles (that are column-major). +template +struct ColumnMajorBlockLinear { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + /// Size of a block in rows + static int const kBlockRows = BlockRows; + + /// Size of a block in columns + static int const kBlockColumns = BlockColumns; + +private: + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorBlockLinear(Index ldm = 0): stride_(ldm) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static ColumnMajorBlockLinear packed(MatrixCoord const &extent) { + return ColumnMajorBlockLinear(extent.row() * kBlockRows * kBlockColumns); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord const &coord) const { + return + (coord.row() % kBlockRows) + + (coord.column() % kBlockColumns) * kBlockRows + + (coord.row() / kBlockRows) * kBlockRows * kBlockColumns + + (coord.column() / kBlockColumns) * stride_[0]; + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + MatrixCoord inverse(LongIndex offset) const { + + return MatrixCoord(0, 0); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const { + return stride_[idx]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index & stride(int idx) { + return stride_[idx]; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(MatrixCoord const &extent) const { + return (extent.column() + kBlockColumns - 1) / kBlockColumns * stride_[0]; + } +}; + +/// Mapping function for block-linear matrices. Matrix is structured +/// as row-major arrangement of 2D tiles (that are row-major) +template +struct RowMajorBlockLinear { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + /// Size of a block in rows + static int const kBlockRows = BlockRows; + + /// Size of a block in columns + static int const kBlockColumns = BlockColumns; + +private: + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorBlockLinear(Index ldm = 0): stride_(ldm) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static RowMajorBlockLinear packed(MatrixCoord const &extent) { + return RowMajorBlockLinear(extent.column() * kBlockRows * kBlockColumns); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord const &coord) const { + return + (coord.column() % kBlockColumns) + + (coord.row() % kBlockRows) * kBlockColumns + + (coord.column() / kBlockColumns) * kBlockRows * kBlockColumns + + (coord.row() / kBlockRows) * stride_[0]; + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + MatrixCoord inverse(LongIndex offset) const { + return MatrixCoord(0, 0); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const { + return stride_[idx]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index & stride(int idx) { + return stride_[idx]; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(MatrixCoord const &extent) const { + return (extent.row() + kBlockRows - 1) / kBlockRows * stride_[0]; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +struct GeneralMatrix { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 2; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + +private: + // + // Data members + // + + Matrix layout_id_; + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + GeneralMatrix(): layout_id_(Matrix::kColumnMajor), stride_(make_Coord(0, 1)) { } + + /// Ctor + MCTLASS_HOST_DEVICE + GeneralMatrix( + Matrix layout_id, + Index ldm, + Index interleave): layout_id_(layout_id), stride_(make_Coord(ldm, interleave)) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static GeneralMatrix packed( + MatrixCoord const &extent, + Matrix layout_id = Matrix::kColumnMajor, + Index interleave = 1) { + + Index c; + if (layout_id == Matrix::kRowMajor) { + c = extent.column(); + } + else { + c = extent.row(); + } + + Index ldm = c * interleave; + + return GeneralMatrix(layout_id, ldm, interleave); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord const &coord) const { + Index c, s; + if (layout_id_ == Matrix::kRowMajor) { + c = coord.column(); + s = coord.row(); + } + else { + s = coord.column(); + c = coord.row(); + } + + Index v = s / stride_[1]; + Index residual = (s % stride_[1]); + + return LongIndex(c) * LongIndex(stride_[1]) + LongIndex(v) * LongIndex(stride_[0]) + residual; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + MCTLASS_HOST_DEVICE + Matrix layout_id() const { + return layout_id_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + MCTLASS_HOST_DEVICE + Matrix & layout_id() { + return layout_id_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const { + return stride_[idx]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + typename Stride::Index & stride(int idx) { + return stride_[idx]; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(MatrixCoord const &extent) const { + Index s; + if (layout_id_ == Matrix::kRowMajor) { + s = extent.row(); + } + else { + s = extent.column(); + } + + Index v = Index((s + stride_[1] - 1) / stride_[1]); + return LongIndex(v) * LongIndex(stride_[0]); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines transposes of matrix layouts +template +struct LayoutTranspose; + +/// Transpose of row-major is column-major +template <> +struct LayoutTranspose { + using type = layout::ColumnMajor; +}; + +/// Transpose of column-major is row-major +template <> +struct LayoutTranspose { + using type = layout::RowMajor; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace layout +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/layout/permute.h b/csrc/mctlass/include/mctlass/layout/permute.h new file mode 100644 index 0000000..1f2ca8c --- /dev/null +++ b/csrc/mctlass/include/mctlass/layout/permute.h @@ -0,0 +1,827 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines layout functions used by GEMM+permute path for common tensor or matrix formats. + + Like Layout functions, permute layout functions map logical coordinates to linear memory. They often require additional + data to describe strides between elements. + + Permute layout functions must implement all members in the interface of NoPermute<> defined in this file. Address offset + computation lies in operator() with private member variables {col_permute_, row_permute_ and stride_} as new addresses after permute op. +*/ +#pragma once +#if defined(__MACACC_RTC__) +#include +#else +#include "assert.h" +#endif +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/coord.h" +#include "mctlass/tensor_coord.h" + +namespace mctlass { +namespace layout { + +// template +// struct PermuteSelect { +// // Try to give a reasonable error message to the user +// static_assert(!platform::is_same::value, // aka always_false +// "You've tried to use a layout permutation for which the implementation is not availble. " +// "In order to provide an implementation for a particular combination of matrix layout " +// "and direction (direct/inverse), please specialize PermuteSelect trait."); +// }; + +// Base template for defining specializations of permutation inverses +template +struct InversePermute +{ + // Try to give a reasonable error message to the user + static_assert(!platform::is_same::value, // aka always_false + "To apply permutation to a GEMM input operand (A or B), an inverse permutation for the desired " + "permute class must be defined and enabled by specializing mctlass::layout::InversePermute trait."); +}; + +class PermuteBase { +public: + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; +}; + +class NoPermute : public PermuteBase { +public: + // + // Methods + // + + /// Constructor from matrix extent + MCTLASS_HOST_DEVICE + NoPermute(MatrixCoord extent, Index stride) { }; + + /// Constructor from pitch-linear extent + MCTLASS_HOST_DEVICE + NoPermute(PitchLinearCoord extent, Index stride) { }; + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord coord) const { return 0; } // not correct but should never be called + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord coord) const { return 0; } // not correct but should never be called +}; + +template<> +struct InversePermute { + using type = NoPermute; +}; + +/// Helper trait to detect if permute operation is a noop +template +bool constexpr is_trivial_permute = platform::is_same::value; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Defines permute layouts of various tensor formats. +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Tensor4DPermute0213 +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Permute layout function for 4-D permuted tensors with matrix (dimensions [M, N]) reshaped +/// as [M/D1, D1, D2, N/D2]. Then perform permute([0, 2, 1, 3]) on the corresponding tensor. +template +class Tensor4DPermute0213RowMajor : public PermuteBase { +private: + // + // Data members + // + + Index D3_; + + Index stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor4DPermute0213RowMajor(MatrixCoord extent, Index stride) { + + assert(extent.row() % D1 == 0); + assert(extent.column() % D2 == 0); + + D3_ = extent.column() / D2; + + stride_ = stride * D1 / D2; + } + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor4DPermute0213RowMajor(PitchLinearCoord extent, Index stride) + : Tensor4DPermute0213RowMajor(MatrixCoord(extent.strided(), extent.contiguous()), stride) {} + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord coord) const { + + // [i,j,k,l] -> [i,k,j,l] + Index l = coord.column() % D3_; + Index k = coord.column() / D3_; + Index j = coord.row() % D1; + Index i = coord.row() / D1; + + MatrixCoord permuted{k + i * D2, l + j * D3_}; + + return LongIndex(permuted.row()) * LongIndex(stride_) + LongIndex(permuted.column()); + } + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord coord) const { + return operator()(MatrixCoord(coord.strided(), coord.contiguous())); + } +}; + +// Inverse for Tensor4DPermute0213 can be implemented by simply swapping D1 and D2 +template +class Tensor4DPermute0213RowMajorInverse : public Tensor4DPermute0213RowMajor { +public: + using Base = Tensor4DPermute0213RowMajor; + using Base::Base; +}; + +template +struct InversePermute> { + using type = Tensor4DPermute0213RowMajorInverse; +}; + +template +struct InversePermute> { + using type = Tensor4DPermute0213RowMajor; +}; + +/// Permute layout function for 4-D permuted tensors with matrix (dimensions [M, N]) reshaped +/// as [M/D1, D1, D2, N/D2]. Then perform permute([0, 2, 1, 3]) on the corresponding tensor. +template +class Tensor4DPermute0213ColumnMajor : public PermuteBase { +private: + // + // Data members + // + + Index D0_; + + Index stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor4DPermute0213ColumnMajor(MatrixCoord extent, Index stride) { + + assert(extent.row() % D1 == 0); + assert(extent.column() % D2 == 0); + + D0_ = extent.row() / D1; + + stride_ = stride * D2 / D1; + } + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor4DPermute0213ColumnMajor(PitchLinearCoord extent, Index stride) + : Tensor4DPermute0213ColumnMajor(MatrixCoord(extent.contiguous(), extent.strided()), stride) {} + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord coord) const { + + // [i,j,k,l] -> [i,k,j,l] + Index l = coord.column() / D2; + Index k = coord.column() % D2; + Index j = coord.row() / D0_; + Index i = coord.row() % D0_; + + MatrixCoord permuted{i + k * D0_, j + l * D1}; + + return LongIndex(permuted.row()) + LongIndex(permuted.column()) * LongIndex(stride_); + } + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord coord) const { + return operator()(MatrixCoord(coord.contiguous(), coord.strided())); + } +}; + +// Inverse for Tensor4DPermute0213 can be implemented by simply swapping D1 and D2 +template +class Tensor4DPermute0213ColumnMajorInverse : public Tensor4DPermute0213ColumnMajor { +public: + using Base = Tensor4DPermute0213ColumnMajor; + using Base::Base; +}; + +template +struct InversePermute> { + using type = Tensor4DPermute0213ColumnMajorInverse; +}; + +template +struct InversePermute> { + using type = Tensor4DPermute0213ColumnMajor; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Tensor4DPermuteBMM0213 +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Permute layout function for 4-D permuted tensors for BMM with BMM tensor (dimensions [B, M, N]) reshaped +/// as [B/D1, D1, M, N]. Then perform permute([0, 2, 1, 3]) on the corresponding whole BMM tensor. +template +class Tensor4DPermuteBMM0213RowMajor : public PermuteBase { +private: + // + // Data members + // + + Index D3_; + + Index stride_; + + Index batch_stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor4DPermuteBMM0213RowMajor(MatrixCoord extent, Index stride) { + + Index D2 = extent.row(); + D3_ = extent.column(); + + stride_ = stride * D1; + batch_stride_ = D2 * stride_; + } + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor4DPermuteBMM0213RowMajor(PitchLinearCoord extent, Index stride) + : Tensor4DPermuteBMM0213RowMajor(MatrixCoord(extent.strided(), extent.contiguous()), stride) {} + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord coord) const { + + // The batch index for BMM + Index BMM_batch_idx = blockIdx.z; + + // [i,j,k,l] -> [i,k,j,l] + Index l = coord.column(); + Index k = coord.row(); + Index j = BMM_batch_idx % D1; + Index i = BMM_batch_idx / D1; + + Index pbatch = i; + MatrixCoord pcoord{k, l + j * D3_}; + + return pbatch * LongIndex(batch_stride_) + pcoord.row() * LongIndex(stride_) + pcoord.column(); + } + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord coord) const { + return operator()(MatrixCoord(coord.strided(), coord.contiguous())); + } +}; + +template +class Tensor4DPermuteBMM0213RowMajorInverse : public PermuteBase { +private: + // + // Data members + // + + Index D3_; + + Index stride_; + + Index batch_stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor4DPermuteBMM0213RowMajorInverse(MatrixCoord extent, Index stride) { + + assert(extent.column() % D1 == 0); + + Index D2 = extent.row(); + D3_ = extent.column() / D1; + + stride_ = stride / D1; + + batch_stride_ = D2 * stride_; + } + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor4DPermuteBMM0213RowMajorInverse(PitchLinearCoord extent, Index stride) + : Tensor4DPermuteBMM0213RowMajorInverse(MatrixCoord(extent.strided(), extent.contiguous()), stride) {} + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord coord) const { + + // The batch index for BMM + Index BMM_batch_idx = blockIdx.z; + + Index l = coord.column() % D3_; + Index j = coord.column() / D3_; + Index k = coord.row(); + Index i = BMM_batch_idx; + + // compute original [batch, row, col] index + Index pbatch = j + i * D1; + MatrixCoord pcoord{k, l}; + + return pbatch * LongIndex(batch_stride_) + pcoord.row() * LongIndex(stride_) + pcoord.column(); + } + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord coord) const { + return operator()(MatrixCoord(coord.strided(), coord.contiguous())); + } +}; + +template +struct InversePermute> { + using type = Tensor4DPermuteBMM0213RowMajorInverse; +}; + +template +struct InversePermute> { + using type = Tensor4DPermuteBMM0213RowMajor; +}; + +/// Permute layout function for 4-D permuted tensors for BMM with BMM tensor (dimensions [B, M, N]) reshaped +/// as [B/D1, D1, M, N]. Then perform permute([0, 3, 2, 1]) on the corresponding whole BMM tensor. +template +class Tensor4DPermuteBMM0321ColumnMajor : public PermuteBase { +private: + // + // Data members + // + + Index D2_; + + Index stride_; + + Index batch_stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor4DPermuteBMM0321ColumnMajor(MatrixCoord extent, Index stride) { + + D2_ = extent.row(); + Index D3 = extent.column(); + + stride_ = stride * D1; + batch_stride_ = stride_ * D3; + } + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor4DPermuteBMM0321ColumnMajor(PitchLinearCoord extent, Index stride) + : Tensor4DPermuteBMM0321ColumnMajor(MatrixCoord(extent.contiguous(), extent.strided()), stride) {} + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord coord) const { + + Index BMM_batch_idx = blockIdx.z; + + // [i,j,k,l] -> [i,k,j,l] + Index l = coord.column(); + Index k = coord.row(); + Index j = BMM_batch_idx % D1; + Index i = BMM_batch_idx / D1; + + Index pbatch = i; + MatrixCoord pcoord{k + j * D2_, l}; + + return pbatch * LongIndex(batch_stride_) + pcoord.row() + pcoord.column() * LongIndex(stride_); + } + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord coord) const { + return operator()(MatrixCoord(coord.contiguous(), coord.strided())); + } +}; + +template +class Tensor4DPermuteBMM0321ColumnMajorInverse : public PermuteBase { +private: + // + // Data members + // + + Index D2_; + + Index stride_; + + Index batch_stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor4DPermuteBMM0321ColumnMajorInverse(MatrixCoord extent, Index stride) { + + assert(extent.row() % D1 == 0); + + D2_ = extent.row() / D1; + Index D3 = extent.column(); + + stride_ = stride / D1; + batch_stride_ = stride_ * D3; + } + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor4DPermuteBMM0321ColumnMajorInverse(PitchLinearCoord extent, Index stride) + : Tensor4DPermuteBMM0321ColumnMajorInverse(MatrixCoord(extent.contiguous(), extent.strided()), stride) {} + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord coord) const { + + Index BMM_batch_idx = blockIdx.z; + + // The following assumes grouping [(D0)->batch, (D1,D2)->row, (D3)->col] + Index l = coord.column(); + Index k = coord.row() % D2_; + Index j = coord.row() / D2_; + Index i = BMM_batch_idx; + + Index pbatch = i * D1 + j; + MatrixCoord pcoord{k, l}; + + return pbatch * LongIndex(batch_stride_) + pcoord.row() + pcoord.column() * LongIndex(stride_); + } + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord coord) const { + return operator()(MatrixCoord(coord.contiguous(), coord.strided())); + } +}; + +template +struct InversePermute> { + using type = Tensor4DPermuteBMM0321ColumnMajorInverse; +}; + +template +struct InversePermute> { + using type = Tensor4DPermuteBMM0321ColumnMajor; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Tensor5DPermute20314 +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Permute layout function for 5-D permuted tensors with output matrix (dimension as [M, N]) reshaped +/// as [M/T1, T1, T2, T3, N/T2/T3]. Then perform permute([2, 0, 3, 1, 4]) on the corresponding output tensor. +template +class Tensor5DPermute20314RowMajor : public PermuteBase { +private: + // + // Data members + // + + Index T0_; + + Index T4_; + + Index stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor5DPermute20314RowMajor(MatrixCoord extent, Index stride) { + + assert(extent.row() % T1 == 0); + assert(extent.column() % (T2 * T3) == 0); + + T0_ = extent.row() / T1; + T4_ = extent.column() / (T2 * T3); + + /// Update stride_permute with stride + stride_ = stride / T2 * T1; // stride in Elements + } + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor5DPermute20314RowMajor(PitchLinearCoord extent, Index stride) + : Tensor5DPermute20314RowMajor(MatrixCoord(extent.strided(), extent.contiguous()), stride) {} + + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord coord) const { + + // Permute as torch.permute(X1, [2, 0, 3, 1, 4]) -> 5D Tensor indices as [i,j,k,l,m], the dimension of X + // is [T0, T1, T2, T3, T4], after permutation the dim of X1 is [T2, T0, T3, T1, T4]. + + Index m = coord.column() % T4_; + Index l = (coord.column() / T4_) % T3; + Index k = (coord.column() / T4_) / T3; + Index j = coord.row() % T1; + Index i = coord.row() / T1; + + MatrixCoord permuted{i + k * T0_, m + j * T4_ + l * T1 * T4_}; + + return LongIndex(permuted.row()) * LongIndex(stride_) + LongIndex(permuted.column()); + } + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord coord) const { + return operator()(MatrixCoord(coord.strided(), coord.contiguous())); + } +}; + +/// Inverse for Tensor5DPermute20314 (could also be given a proper name, e.g. Tensor5DPermute13024). +template +class Tensor5DPermute20314RowMajorInverse : public PermuteBase { +private: + // + // Data members + // + + Index T0_; + + Index T4_; + + // Permuted stride in units of elements + Index stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor5DPermute20314RowMajorInverse(MatrixCoord extent, Index stride) { + + assert(extent.row() % T2 == 0); + assert(extent.column() % (T1 * T3) == 0); + + T0_ = extent.row() / T2; + T4_ = extent.column() / (T1 * T3); + + stride_ = stride / T1 * T2; + } + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor5DPermute20314RowMajorInverse(PitchLinearCoord extent, Index stride) + : Tensor5DPermute20314RowMajorInverse(MatrixCoord(extent.strided(), extent.contiguous()), stride) {} + + /// Computes the offset after the inverse of permute operation in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord coord) const { + + Index m = coord.column() % T4_; + Index j = (coord.column() / T4_) % T1; + Index l = (coord.column() / T4_) / T1; + Index i = coord.row() % T0_; + Index k = coord.row() / T0_; + + MatrixCoord permuted{j + i * T1, m + l * T4_ + k * T3 * T4_}; + + return LongIndex(permuted.row()) * LongIndex(stride_) + LongIndex(permuted.column()); + } + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord coord) const { + return operator()(MatrixCoord(coord.strided(), coord.contiguous())); + } +}; + +template +struct InversePermute> { + using type = Tensor5DPermute20314RowMajorInverse; +}; + +template +struct InversePermute> { + using type = Tensor5DPermute20314RowMajor; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Tensor5DPermute02413 +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Permute layout function for 5-D permuted tensors with matrix (dimensions [M, N]) reshaped +/// as [M/T1, T1, T2, T3, N/T2/T3]. Then perform permute([0, 2, 4, 1, 3]) on the corresponding tensor. +template +class Tensor5DPermute02413ColumnMajor : public PermuteBase { +private: + // + // Data members + // + + Index T0_; + + Index T4_; + + Index stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor5DPermute02413ColumnMajor(MatrixCoord extent, Index stride) { + + assert(extent.row() % T1 == 0); + assert(extent.column() % (T2 * T3) == 0); + + T0_ = extent.row() / T1; + T4_ = extent.column() / (T2 * T3); + + /// Update stride_permute with stride + stride_ = stride / T1 * T2; // stride in Elements + } + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor5DPermute02413ColumnMajor(PitchLinearCoord extent, Index stride) + : Tensor5DPermute02413ColumnMajor(MatrixCoord(extent.contiguous(), extent.strided()), stride) {} + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord coord) const { + + // Permute as torch.permute(X1, [2, 0, 3, 1, 4]) -> 5D Tensor indices as [i,j,k,l,m], the dimension of X + // is [T0, T1, T2, T3, T4], after permutation the dim of X1 is [T0, T2, T4, T1, T3]. + + Index m = (coord.column() / T2) / T3; + Index l = (coord.column() / T2) % T3; + Index k = coord.column() % T2; + Index j = coord.row() / T0_; + Index i = coord.row() % T0_; + + MatrixCoord permuted{i + k * T0_, m + j * T4_ + l * T4_ * T1}; + + return LongIndex(permuted.row()) + LongIndex(permuted.column()) * LongIndex(stride_); + } + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord coord) const { + return operator()(MatrixCoord(coord.contiguous(), coord.strided())); + } +}; + +/// Inverse for Tensor5DPermute02413ColumnMajor +template +class Tensor5DPermute02413ColumnMajorInverse : public PermuteBase { +private: + // + // Data members + // + + Index T0_; + + Index T4_; + + // Permuted stride in units of elements + Index stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor5DPermute02413ColumnMajorInverse(MatrixCoord extent, Index stride) { + + assert(extent.row() % T2 == 0); + assert(extent.column() % (T1 * T3) == 0); + + T0_ = extent.row() / T2; + T4_ = extent.column() / (T1 * T3); + + stride_ = stride / T2 * T1; + } + + /// Constructor + MCTLASS_HOST_DEVICE + Tensor5DPermute02413ColumnMajorInverse(PitchLinearCoord extent, Index stride) + : Tensor5DPermute02413ColumnMajorInverse(MatrixCoord(extent.contiguous(), extent.strided()), stride) {} + + /// Computes the offset after the inverse of permute operation in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(MatrixCoord coord) const { + + Index m = coord.column() % T4_; + Index j = (coord.column() / T4_) % T1; + Index l = (coord.column() / T4_) / T1; + Index i = coord.row() % T0_; + Index k = coord.row() / T0_; + + MatrixCoord permuted{i + j * T0_, k + l * T2 + m * T2 * T3}; + + return LongIndex(permuted.row()) + LongIndex(permuted.column()) * LongIndex(stride_); + } + + /// Computes the offset after Permute Op in logical elements + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord coord) const { + return operator()(MatrixCoord(coord.contiguous(), coord.strided())); + } +}; + +template +struct InversePermute> { + using type = Tensor5DPermute02413ColumnMajorInverse; +}; + +template +struct InversePermute> { + using type = Tensor5DPermute02413ColumnMajor; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace layout +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/layout/pitch_linear.h b/csrc/mctlass/include/mctlass/layout/pitch_linear.h new file mode 100644 index 0000000..8150cf2 --- /dev/null +++ b/csrc/mctlass/include/mctlass/layout/pitch_linear.h @@ -0,0 +1,148 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines layout functions used by TensorRef and derived classes for pitch-linear memory. +*/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/coord.h" +#include "mctlass/pitch_linear_coord.h" + +namespace mctlass { +namespace layout { + +template + using PitchLinearShape = mctlass::PitchLinearShape < Contiguous, Strided >; + using PitchLinearCoord = PitchLinearCoord; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Mapping function for pitch-linear memory +class PitchLinear { +public: + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + +private: + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + PitchLinear(LongIndex ldm = 0): stride_(ldm) { } + + /// Constructor + MCTLASS_HOST_DEVICE + PitchLinear(Stride _stride): stride_(_stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static PitchLinear packed(TensorCoord const &extent) { + return PitchLinear(extent.contiguous()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return LongIndex(coord.contiguous()) + LongIndex(coord.strided()) * LongIndex(stride_[0]); + } + + /// Returns the logical coordinate given an offset. + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex index) const { + return make_Coord( + TensorCoord::Index(index % stride_[0]), + TensorCoord::Index(index / stride_[0]) + ); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + LongIndex stride(int rank) const { + return stride_[rank]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + LongIndex & stride(int rank) { + return stride_[rank]; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return extent.strided() * stride_[0]; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace layout +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/layout/tensor.h b/csrc/mctlass/include/mctlass/layout/tensor.h new file mode 100644 index 0000000..51639b0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/layout/tensor.h @@ -0,0 +1,636 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines layout functions used by TensorRef and derived classes for common 4-D and 5-D + tensor formats. + + Layout functions map logical coordinates to linear memory. They often require additional + data to describe strides between elements. + + Layout functions must implement all members in the public interface of IdentityTensorLayout<> + defined in mctlass/tensor_ref.h. +*/ +#pragma once +#if defined(__MACACC_RTC__) +#include +#else +#include "assert.h" +#endif +#include "mctlass/mctlass.h" +#include "mctlass/fast_math.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/coord.h" +#include "mctlass/tensor_coord.h" + +namespace mctlass { +namespace layout { + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Defines data layouts of various tensor formats usable by TensorRef and other classes. +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Mapping function for 4-D NHWC tensors. +class TensorNHWC { +public: + /// Logical rank of tensor + static int const kRank = 4; + + /// Rank of stride vector + static int const kStrideRank = 3; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate (n, h, w, c) + using TensorCoord = Tensor4DCoord; + + /// Stride vector + using Stride = Coord; + +private: + // + // Data members + // + + /// Stride data member - [stride_w, stride_h, stride_n] + Stride stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + TensorNHWC(Stride const &stride = Stride(0)): stride_(stride) { } + + /// Constructor + MCTLASS_HOST_DEVICE + TensorNHWC( + typename Stride::Index stride_w, ///< number of elements between adjacent W coordinates + typename Stride::Index stride_h, ///< number of elements between adjacent H coordinates + typename Stride::Index stride_n ///< number of elements between adjacent N coordinates + ): + stride_(make_Coord(stride_w, stride_h, stride_n)) { } + + /// Constructor + // Once convolutions implement 64b stride this ctor can be deleted + MCTLASS_HOST_DEVICE + TensorNHWC(Coord const &stride): + stride_(make_Coord( + static_cast(stride[0]), + static_cast(stride[1]), + static_cast(stride[2])) + ) { } + + /// Helper returns a layout to a tightly packed NHWC tensor. + MCTLASS_HOST_DEVICE + static TensorNHWC packed(TensorCoord const &extent) { + return TensorNHWC( + make_Coord( + extent.c(), + extent.w() * extent.c(), + extent.h() * extent.w() * extent.c() + ) + ); + } + + /// Returns the offset of a coordinate (n, h, w, c) in linear memory. + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return coord.c() + + LongIndex(stride_[0] * coord.w()) + + LongIndex(stride_[1] * coord.h()) + + LongIndex(stride_[2] * coord.n()); + } + + /// Returns the offset of a pitchlinear coordinate in linear memory. + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord coord) const { + return coord.contiguous() + LongIndex(coord.strided() * stride_[2]); + } + + /// Returns the logical coordinate (n, h, w, c) from a given offset in linear memory. + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex index) const { + + int n = 0, h = 0, w = 0, c = 0; + + #if defined(__MACA_ARCH__) + int tmp = 0; + c = int(index % static_cast(stride_[0])); + + unsigned int hw_mul, hw_shr, w_mul, w_shr, c_mul, c_shr; + + find_divisor(hw_mul, hw_shr, stride_[2]); + find_divisor(w_mul, w_shr, stride_[1]); + find_divisor(c_mul, c_shr, stride_[0]); + + fast_divmod(n, tmp, index, int(stride_[2]), hw_mul, hw_shr); + fast_divmod(h, w, tmp, int(stride_[1]), w_mul, w_shr); + fast_divmod(w, tmp, w, int(stride_[0]), c_mul, c_shr); + #else + + n = int(index / stride_[2]); + LongIndex residual = index % stride_[2]; + + h = int(residual / stride_[1]); + residual = (residual % stride_[1]); + + w = int(residual / stride_[0]); + c = int(residual % stride_[0]); + + #endif + return TensorCoord(n, h, w, c); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + // it does not make sense if the extent is larger than stride + // and we could not rely on the capacity calculation in such cases + // we could move this checkers to debug code only + if ((extent.c() > stride_[0]) + || (extent.w() * stride_[0] > stride_[1]) + || (extent.h() * stride_[1] > stride_[2])) { + assert(0); + } + return extent.n() * stride_[2]; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Mapping function for 4-D NCHW tensors. +class TensorNCHW { +public: + /// Logical rank of tensor + static int const kRank = 4; + + /// Rank of stride vector + static int const kStrideRank = 3; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = Tensor4DCoord; + + /// Stride vector + using Stride = Coord; + +private: + // + // Data members + // + + /// Stride data member - [w, hw, chw] + Stride stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + TensorNCHW(Stride const &stride = Stride(0)): stride_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static TensorNCHW packed(TensorCoord const &extent) { + return TensorNCHW( + make_Coord( + extent.w(), + extent.w() * extent.h(), + extent.h() * extent.w() * extent.c() + ) + ); + } + + /// Returns the offset of a coordinate in linear memory. + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return coord.w() + + LongIndex(stride_[0] * coord.h()) + + LongIndex(stride_[1] * coord.c()) + + LongIndex(stride_[2] * coord.n()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return extent.n() * stride_[2]; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Mapping function for 4-D NC/xHWx tensors. +template +class TensorNCxHWx { +public: + + /// Interleaving quantity + static int const kInterleave = Interleave; + + /// Logical rank of tensor + static int const kRank = 4; + + /// Rank of stride vector + static int const kStrideRank = 3; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = Tensor4DCoord; + + /// Stride vector + using Stride = Coord; + +private: + // + // Data members + // + + /// Stride data member - [Interleave x w, Interleave x wh, hwc] + Stride stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + TensorNCxHWx(Stride const &stride = Stride(0)): stride_(stride) { } + + /// Constructor + MCTLASS_HOST_DEVICE + TensorNCxHWx( + typename Stride::Index stride_w, ///< number of elements between adjacent W coordinates + typename Stride::Index stride_h, ///< number of elements between adjacent H coordinates + typename Stride::Index stride_n ///< number of elements between adjacent N coordinates + ): + stride_(make_Coord(stride_w, stride_h, stride_n)) { } + + /// Constructor + // Once convolutions implement 64b stride this ctor can be deleted + MCTLASS_HOST_DEVICE + TensorNCxHWx(Coord const &stride): + stride_(make_Coord( + static_cast(stride[0]), + static_cast(stride[1]), + static_cast(stride[2])) + ) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static TensorNCxHWx packed(TensorCoord const &extent) { + return TensorNCxHWx( + make_Coord( + kInterleave * extent.w(), + kInterleave * extent.w() * extent.h(), + extent.h() * extent.w() * extent.c() + ) + ); + } + + /// Returns the offset of a coordinate in linear memory. + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + + Index c_minor = (coord.c() % kInterleave); + Index c_major = (coord.c() / kInterleave); + + return c_minor + + LongIndex(kInterleave * coord.w()) + + LongIndex(stride_[0] * coord.h()) + + LongIndex(stride_[1] * c_major) + + LongIndex(stride_[2] * coord.n()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return extent.n() * stride_[2]; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Mapping function for 4-D CxRSKx tensors. +template +class TensorCxRSKx { +public: + + /// Interleaving quantity + static int const kInterleave = Interleave; + + /// Logical rank of tensor + static int const kRank = 4; + + /// Rank of stride vector + static int const kStrideRank = 3; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = Tensor4DCoord; + + /// Stride vector + using Stride = Coord; + +private: + // + // Data members + // + + /// Stride data member - [Interleave x n, Interleave x nw, Interleave x nwh] + Stride stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + TensorCxRSKx(Stride const &stride = Stride(0)): stride_(stride) { } + + /// Constructor + MCTLASS_HOST_DEVICE + TensorCxRSKx( + typename Stride::Index stride_w, ///< number of elements between adjacent W coordinates + typename Stride::Index stride_h, ///< number of elements between adjacent H coordinates + typename Stride::Index stride_n ///< number of elements between adjacent N coordinates + ): + stride_(make_Coord(stride_w, stride_h, stride_n)) { } + + /// Constructor + // Once convolutions implement 64b stride this ctor can be deleted + MCTLASS_HOST_DEVICE + TensorCxRSKx(Coord const &stride): + stride_(make_Coord( + static_cast(stride[0]), + static_cast(stride[1]), + static_cast(stride[2])) + ) { } + + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static TensorCxRSKx packed(TensorCoord const &extent) { + return TensorCxRSKx( + make_Coord( + kInterleave * extent.n(), + kInterleave * extent.n() * extent.w(), + kInterleave * extent.n() * extent.w() * extent.h() + ) + ); + } + + /// Returns the offset of a coordinate in linear memory. + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + + Index c_minor = (coord.c() % kInterleave); + Index c_major = (coord.c() / kInterleave); + + return c_minor + + LongIndex(kInterleave * coord.n()) + + LongIndex(stride_[0] * coord.w()) + + LongIndex(stride_[1] * coord.h()) + + LongIndex(stride_[2] * c_major); + } + + /// Returns the offset of a pitchlinear coordinate in linear memory. + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord const &coord) const { + return (coord.contiguous() % kInterleave) + + LongIndex((coord.contiguous() / kInterleave) * stride_[2]) + + LongIndex(coord.strided() * kInterleave); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return (extent.c() / kInterleave * stride_[2]); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Mapping function for 5-D NDHWC tensors. +class TensorNDHWC { +public: + /// Logical rank of tensor + static int const kRank = 5; + + /// Rank of stride vector + static int const kStrideRank = 4; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate (n, d, h, w, c) + using TensorCoord = Tensor5DCoord; + + /// Stride vector + using Stride = Coord; + +private: + // + // Data members + // + + /// Stride data member - [c, wc, hwc, dhwc] + Stride stride_; + +public: + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + TensorNDHWC(Stride const &stride = Stride(0)): stride_(stride) { } + + /// Constructor + MCTLASS_HOST_DEVICE + TensorNDHWC( + typename Stride::Index c, + typename Stride::Index wc, + typename Stride::Index hwc, + typename Stride::Index dhwc): + stride_(make_Coord(c, wc, hwc, dhwc)) { } + + /// Constructor + // Once convolutions implement 64b stride this ctor can be deleted + MCTLASS_HOST_DEVICE + TensorNDHWC(Coord const &stride): + stride_(make_Coord( + static_cast(stride[0]), + static_cast(stride[1]), + static_cast(stride[2]), + static_cast(stride[3])) + ) { } + + /// Helper returns a layout to a tightly packed NHWC tensor. + MCTLASS_HOST_DEVICE + static TensorNDHWC packed(TensorCoord const &extent) { + return TensorNDHWC( + make_Coord( + extent.c(), + extent.w() * extent.c(), + extent.h() * extent.w() * extent.c(), + extent.d() * extent.h() * extent.w() * extent.c() + ) + ); + } + + /// Returns the offset of a coordinate (n, d, h, w, c) in linear memory. + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return coord.c() + + LongIndex(stride_[0] * coord.w()) + + LongIndex(stride_[1] * coord.h()) + + LongIndex(stride_[2] * coord.d()) + + LongIndex(stride_[3] * coord.n()); + } + + /// Returns the offset of a pitchlinear coordinate in linear memory. + MCTLASS_HOST_DEVICE + LongIndex operator()(PitchLinearCoord coord) const { + return coord.contiguous() + LongIndex(coord.strided() * stride_[3]); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + // it does not make sense if the extent is larger than stride + // and we could not rely on the capacity calculation in such cases + // we could move this checkers to debug code only + if ((extent.c() > stride_[0]) + || (extent.w() * stride_[0] > stride_[1]) + || (extent.h() * stride_[1] > stride_[2]) + || (extent.d() * stride_[2] > stride_[3])) { + assert(0); + } + return extent.n() * stride_[3]; + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace layout +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/layout/tensor_op_multiplicand_sm70.h b/csrc/mctlass/include/mctlass/layout/tensor_op_multiplicand_sm70.h new file mode 100644 index 0000000..2f44891 --- /dev/null +++ b/csrc/mctlass/include/mctlass/layout/tensor_op_multiplicand_sm70.h @@ -0,0 +1,1044 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/coord.h" +#include "mctlass/layout/pitch_linear.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace layout { + +// template < +// int ElementSize, +// gemm::Operand Operand +// > +// struct VoltaTensorOpMultiplicandCongruous; + +// template < +// int ElementSize, +// gemm::Operand Operand +// > +// struct ColumnMajorVoltaTensorOpMultiplicandCongruous; +// template < +// int ElementSize, +// gemm::Operand Operand +// > +// struct RowMajorVoltaTensorOpMultiplicandCongruous; +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template based on element size (in bits) - defined in terms of pitch-linear memory. +template +struct VoltaTensorOpMultiplicandCongruous { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + /// This layout is optimized for 128b accesses + static int const kAccessSize = 128; + + /// Fundamental tile shape in units of vectors + using TileShape = PitchLinearShape<8, 4>; + + /// Fundamental partition shape in units of vectors + using PartitionShape = PitchLinearShape<8, 2>; + + // + // Static constants + // + + static int const kElementSize = ElementSize; + static int const kElementsPerAccess = kAccessSize / kElementSize; + + using PartitionCount = PitchLinearShape< + TileShape::kContiguous / PartitionShape::kContiguous, + TileShape::kStrided / PartitionShape::kStrided + >; + + using AccessCount = PitchLinearShape< + PartitionShape::kContiguous, + PartitionShape::kStrided + >; + +private: + + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + VoltaTensorOpMultiplicandCongruous(Index ldm = 0): stride_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + VoltaTensorOpMultiplicandCongruous(Stride stride): stride_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static VoltaTensorOpMultiplicandCongruous packed(TensorCoord const &extent) { + return VoltaTensorOpMultiplicandCongruous(extent[0]); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + + // First, compute c and s of vector within source (in units of vector accesses) + int vec_contiguous_idx = coord.contiguous() / kElementsPerAccess; + int vec_strided_idx = coord.strided(); + + // Compute the fundamental tile being accessed + int tile_contiguous_idx = vec_contiguous_idx / TileShape::kContiguous; + int tile_strided_idx = vec_strided_idx / TileShape::kStrided; + + int tile_contiguous_residual = vec_contiguous_idx % TileShape::kContiguous; + int tile_strided_residual = vec_strided_idx % TileShape::kStrided; + + // Then swizzle in a tile + // Swizzle pattern is (tid[2:0] << 2)|(tid[4:3] ^ tid[2:1]) + int permuted_strided_within_tile = (tile_contiguous_residual >> 1); + int permuted_contiguous_within_tile = (tile_strided_residual ^ permuted_strided_within_tile) | + ((tile_contiguous_residual & 1) << 2); + // Compute final element location + int element_contiguous = (tile_contiguous_idx * TileShape::kContiguous + + permuted_contiguous_within_tile) * kElementsPerAccess + (coord.contiguous() % kElementsPerAccess); + + int element_strided = tile_strided_idx * TileShape::kStrided + permuted_strided_within_tile; + + return element_contiguous + element_strided * stride_[0]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return extent[1] * stride_[0]; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a column-major view of pitch-linear memory to VoltaTensorOpMultiplicandCongruous +template +struct ColumnMajorVoltaTensorOpMultiplicandCongruous { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = VoltaTensorOpMultiplicandCongruous; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorVoltaTensorOpMultiplicandCongruous(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorVoltaTensorOpMultiplicandCongruous(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static ColumnMajorVoltaTensorOpMultiplicandCongruous packed(TensorCoord const &extent) { + return ColumnMajorVoltaTensorOpMultiplicandCongruous(extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.row(), coord.column())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.contiguous(), coord.strided()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.row(), extent.column())); + } +}; + +/// Template mapping a row-major view of pitch-linear memory to VoltaTensorOpMultiplicandCongruous +template +struct RowMajorVoltaTensorOpMultiplicandCongruous { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = VoltaTensorOpMultiplicandCongruous; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorVoltaTensorOpMultiplicandCongruous(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorVoltaTensorOpMultiplicandCongruous(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static RowMajorVoltaTensorOpMultiplicandCongruous packed(TensorCoord const &extent) { + return RowMajorVoltaTensorOpMultiplicandCongruous(extent.column()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.column(), coord.row())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.strided(), coord.contiguous()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.column(), extent.row())); + } +}; + + +/// Template based on element size (in bits) - defined in terms of pitch-linear memory. +// template +template +struct VoltaTensorOpMultiplicandBCongruous { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + /// This layout is optimized for 128b accesses + static int const kAccessSize = 128; + + /// Fundamental tile shape in units of vectors + using TileShape = PitchLinearShape<8, 4>; + + /// Fundamental partition shape in units of vectors + using PartitionShape = PitchLinearShape<4, 4>; + + // + // Static constants + // + + static int const kElementSize = ElementSize; + static int const kElementsPerAccess = kAccessSize / kElementSize; + + using PartitionCount = PitchLinearShape< + TileShape::kContiguous / PartitionShape::kContiguous, + TileShape::kStrided / PartitionShape::kStrided + >; + + using AccessCount = PitchLinearShape< + PartitionShape::kContiguous, + PartitionShape::kStrided + >; + +private: + + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + VoltaTensorOpMultiplicandBCongruous(Index ldm = 0): stride_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + VoltaTensorOpMultiplicandBCongruous(Stride stride): stride_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static VoltaTensorOpMultiplicandBCongruous packed(TensorCoord const &extent) { + return VoltaTensorOpMultiplicandBCongruous(extent[0]); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + + // First, compute c and s of vector within source (in units of vector accesses) + int vec_contiguous_idx = coord.contiguous() / kElementsPerAccess; + int vec_strided_idx = coord.strided(); + + // Compute the fundamental tile being accessed + int tile_contiguous_idx = vec_contiguous_idx / TileShape::kContiguous; + int tile_strided_idx = vec_strided_idx / TileShape::kStrided; + + int tile_contiguous_residual = vec_contiguous_idx % TileShape::kContiguous; + int tile_strided_residual = vec_strided_idx % TileShape::kStrided; + + // Then swizzle in a tile + // Swizzle pattern is (tid[1:0] << 3)|(tid & 0x4)|(tid[1:0]) + int permuted_strided_within_tile = (tile_contiguous_residual & 0x3); + int permuted_contiguous_within_tile = (tile_strided_residual ^ permuted_strided_within_tile) | + (tile_contiguous_residual & 0x4); + + // Compute final element location + int element_contiguous = (tile_contiguous_idx * TileShape::kContiguous + + permuted_contiguous_within_tile) * kElementsPerAccess + (coord.contiguous() % kElementsPerAccess); + + int element_strided = tile_strided_idx * TileShape::kStrided + permuted_strided_within_tile; + + return element_contiguous + element_strided * stride_[0]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return extent[1] * stride_[0]; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a column-major view of pitch-linear memory to VoltaTensorOpMultiplicandCongruous +template +struct ColumnMajorVoltaTensorOpMultiplicandBCongruous { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = VoltaTensorOpMultiplicandBCongruous; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorVoltaTensorOpMultiplicandBCongruous(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorVoltaTensorOpMultiplicandBCongruous(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static ColumnMajorVoltaTensorOpMultiplicandBCongruous packed(TensorCoord const &extent) { + return ColumnMajorVoltaTensorOpMultiplicandBCongruous(extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.row(), coord.column())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.contiguous(), coord.strided()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.row(), extent.column())); + } +}; + +/// Template mapping a row-major view of pitch-linear memory to VoltaTensorOpMultiplicandCongruous +template +struct RowMajorVoltaTensorOpMultiplicandBCongruous { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = VoltaTensorOpMultiplicandBCongruous; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorVoltaTensorOpMultiplicandBCongruous(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorVoltaTensorOpMultiplicandBCongruous(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static RowMajorVoltaTensorOpMultiplicandBCongruous packed(TensorCoord const &extent) { + return RowMajorVoltaTensorOpMultiplicandBCongruous(extent.column()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.column(), coord.row())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.strided(), coord.contiguous()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.column(), extent.row())); + } +}; + +/// Template based on element size (in bits) - defined in terms of pitch-linear +/// memory and KBlock size (in elements). +template +struct VoltaTensorOpMultiplicandCrosswise { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + /// This layout is optimized for 64b accesses + static int const kAccessSize = 64; + + // + // Static constants + // + + static int const kElementSize = ElementSize; + static int const kElementsPerAccess = kAccessSize / kElementSize; + static int const kKBlock = KBlock; + + private: + // + // Data members + // + + /// Stride data member. For GEMM, it equals to KBlock x stage. + Stride stride_; + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + VoltaTensorOpMultiplicandCrosswise(Index ldm = 0) : stride_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + VoltaTensorOpMultiplicandCrosswise(Stride stride) : stride_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static VoltaTensorOpMultiplicandCrosswise packed(TensorCoord const &extent) { + return VoltaTensorOpMultiplicandCrosswise(extent[1]); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + + // + // First, compute c and s of vector within source (in units of vector + // accesses) + // + int vec_contiguous_idx = coord.contiguous() / kElementsPerAccess; + int vec_strided_idx = coord.strided(); + + // + // Then swizzle + // The mapping is like this: + // id[1:0]|(id[3]^id[4])|id[2] + + int vec_strided_within_tile = vec_contiguous_idx & 0x7; + int permuted_vec_contiguous = + (vec_strided_idx & (~0xF)) + (vec_strided_idx & 0x3) * 4 + + (((vec_strided_idx >> 2) ^ ((vec_strided_idx & 0x10) >> 3)) & 0x3); + + permuted_vec_contiguous ^= ((vec_strided_within_tile >> 1) & 0x3); + + int permuted_vec_strided = vec_contiguous_idx; + + // + // Compute final element location + // + + int element_contiguous = permuted_vec_contiguous * kElementsPerAccess + + (coord.contiguous() % kElementsPerAccess); + + return element_contiguous + permuted_vec_strided * (stride_[0] * kElementsPerAccess); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return stride_; } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return stride_; } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return extent[0] * stride_[0]; + } +}; + +/// Template mapping a column-major view of pitch-linear memory to +/// VoltaTensorOpMultiplicandCrosswise +template +struct ColumnMajorVoltaTensorOpMultiplicandCrosswise { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = VoltaTensorOpMultiplicandCrosswise; + + /// This layout is optimized for 64b accesses + static int const kAccessSize = Base::kAccessSize; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + + private: + // + // Data members + // + + Base layout_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorVoltaTensorOpMultiplicandCrosswise(Index ldm = 0) : layout_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorVoltaTensorOpMultiplicandCrosswise(Stride stride) : layout_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static ColumnMajorVoltaTensorOpMultiplicandCrosswise packed( + TensorCoord const &extent) { + return ColumnMajorVoltaTensorOpMultiplicandCrosswise(extent.column()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.row(), coord.column())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.contiguous(), coord.strided()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return layout_.stride(); } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return layout_.stride(); } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.row(), extent.column())); + } +}; + +/// Template mapping a row-major view of pitch-linear memory to +/// TensorOpMultiplicandCrosswise +template +struct RowMajorVoltaTensorOpMultiplicandCrosswise { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = VoltaTensorOpMultiplicandCrosswise; + + /// This layout is optimized for 64b accesses + static int const kAccessSize = Base::kAccessSize; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + + private: + // + // Data members + // + + Base layout_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorVoltaTensorOpMultiplicandCrosswise(Index ldm = 0) : layout_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorVoltaTensorOpMultiplicandCrosswise(Stride stride) : layout_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static RowMajorVoltaTensorOpMultiplicandCrosswise packed( + TensorCoord const &extent) { + return RowMajorVoltaTensorOpMultiplicandCrosswise(extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.column(), coord.row())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.strided(), coord.contiguous()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return layout_.stride(); } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return layout_.stride(); } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.column(), extent.row())); + } +}; + +} // namespace layout +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/layout/tensor_op_multiplicand_sm75.h b/csrc/mctlass/include/mctlass/layout/tensor_op_multiplicand_sm75.h new file mode 100644 index 0000000..374fad0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/layout/tensor_op_multiplicand_sm75.h @@ -0,0 +1,1573 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/coord.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/layout/pitch_linear.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace layout { + +//////////////////////////////////////////////////////////////////////////////// + +/// Template based on element size (in bits) - defined in terms of pitch-linear +/// memory and Crosswise size (in elements). +/// This one is the base class of all Ampere/Turing fp16/bf16/int8/int4/int1 +/// tensor core kernels. tf32 TN uses this too. +template +struct TensorOpMultiplicand { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + + // + // Static constants + // + + /// This layout is optimized for 128b accesses + static int const kAccessSize = 128; + + static int const kElementSize = ElementSize; + static int const kElementsPerAccess = kAccessSize / kElementSize; + static int const kCrosswise = Crosswise; + + /// Contiguous dimension of the tile shape matches one shared memory cache + /// line - 128B. For 128bit access size, it equals to 8 accesses. + static int const kTileShapeContiguous = 128 / (kAccessSize / 8); + + /// Number of kblocks to store PartitionShape::kContiguous Elements + static int const kFactor = + kTileShapeContiguous * kElementsPerAccess / kCrosswise; + + static_assert( + (kFactor > 0), + "kCrosswise should be no large than one shared memory cache line."); + + /// The strided dimension needs to be at least (WarpSize(32) / + /// kTileShapeContiguous) for a warp to access. To ensure conflict free + /// access, it also needs to be at least (kTileShapeContiguous / kFactor). + /// See comments below + static int const kTileShapeStride = + ((kTileShapeContiguous / kFactor) > (32 / kTileShapeContiguous)) + ? (kTileShapeContiguous / kFactor) + : (32 / kTileShapeContiguous); + + /// Fundamental tile shape in units of vectors to guarantee bank conflict free + /// shared memory load/store. + /// For kFactor = 1, TileShape = <8, 8> + /// For kFactor > 1, TileShape = <8, 4> + using TileShape = PitchLinearShape; + + /// Fundamental partition shape in units of vectors + using PartitionShape = PitchLinearShape<4, 4>; + + using PartitionCount = + PitchLinearShape; + + using AccessCount = + PitchLinearShape; + + private: + // + // Data members + // + + /// Stride data member. For GEMM, it equals to kCrosswise x stage. + Stride stride_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicand(Index ldm = 0) : stride_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicand(Stride stride) : stride_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static TensorOpMultiplicand packed(TensorCoord const &extent) { + return TensorOpMultiplicand(extent[0]); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + // + // First, compute c and s of vector within source (in units of vector + // accesses) + // + + int vec_contiguous_idx = coord.contiguous() / kElementsPerAccess; + int vec_strided_idx = coord.strided() / kFactor; + + // Compute the fundamental tile being accessed + int tile_contiguous_idx = + vec_contiguous_idx / (TileShape::kContiguous / kFactor); + + int tile_contiguous_residual = + vec_contiguous_idx % (TileShape::kContiguous / kFactor) + + ((coord.strided() % kFactor) * (TileShape::kContiguous / kFactor)); + int tile_strided_residual = vec_strided_idx % TileShape::kStrided; + + // Compute the 'partition' within the fundamental tile + int partition_contiguous_idx = + tile_contiguous_residual / PartitionShape::kContiguous; + int partition_strided_idx = + tile_strided_residual / PartitionShape::kStrided; + + int partition_contiguous_residual = + tile_contiguous_residual % PartitionShape::kContiguous; + int partition_strided_residual = + tile_strided_residual % PartitionShape::kStrided; + + // + // Then swizzle + // + + int permuted_vec_contiguous_within_partition = + partition_contiguous_residual ^ (partition_strided_residual % 4); + + int permuted_partition_contiguous_within_tile = + partition_contiguous_idx ^ (partition_strided_idx % 2); + + // + // Compute final element location + // + + int element_contiguous = (tile_contiguous_idx * TileShape::kContiguous + + permuted_partition_contiguous_within_tile * + PartitionShape::kContiguous + + permuted_vec_contiguous_within_partition) * + kElementsPerAccess + + (coord.contiguous() % kElementsPerAccess); + + int element_strided = vec_strided_idx; + + return element_contiguous + element_strided * stride_[0] * kFactor; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return stride_; } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return stride_; } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return extent[1] * stride_[0]; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template based on element size (in bits) - defined in terms of pitch-linear +/// memory and Crosswise size (in elements). +template +struct TensorOpMultiplicandCongruous { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicand; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + + private: + // + // Data members + // + + Base layout_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandCongruous(Index ldm = 0) : layout_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandCongruous(Stride stride) : layout_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static TensorOpMultiplicandCongruous packed(TensorCoord const &extent) { + return TensorOpMultiplicandCongruous(extent[0]); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(coord); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return coord; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return layout_.stride(); } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return layout_.stride(); } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(extent); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template based on element size (in bits) - defined in terms of pitch-linear +/// memory and Crosswise size (in elements). +/// This one is just for TF32 NT kernel. +template +struct TensorOpMultiplicandCongruous<32, Crosswise> { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + /// This layout is optimized for 128b accesses + static int const kAccessSize = 128; + + /// Fundamental tile shape in units of vectors + using TileShape = PitchLinearShape<8, 4>; + + /// Partitionshape is the same as TileShape for this layout + using PartitionShape = PitchLinearShape<8, 4>; + + using PartitionCount = + PitchLinearShape; + + using AccessCount = + PitchLinearShape; + + // + // Static constants + // + static int const kElementSize = 32; + static int const kElementsPerAccess = kAccessSize / kElementSize; + + private: + // + // Data members + // + + /// Stride data member. + Stride stride_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandCongruous(Index ldm = 0) : stride_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandCongruous(Stride stride) : stride_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static TensorOpMultiplicandCongruous packed(TensorCoord const &extent) { + return TensorOpMultiplicandCongruous(extent[0]); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + int tc = coord.contiguous() / 32; + int ts = coord.strided() / 4; + + int c = (coord.contiguous() % 32) / kElementsPerAccess; + int s = coord.strided() % 4; + + LongIndex offset = (c ^ (2 * s)) * kElementsPerAccess + s * stride_[0] + + tc * 32 + ts * stride_[0] * 4 + coord.contiguous() % 4; + + return offset; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return stride_; } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return stride_; } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return extent[1] * stride_[0]; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a column-major view of pitch-linear memory to +/// TensorOpMultiplicand +template +struct ColumnMajorTensorOpMultiplicandCongruous { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCongruous; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorTensorOpMultiplicandCongruous(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorTensorOpMultiplicandCongruous(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static ColumnMajorTensorOpMultiplicandCongruous packed(TensorCoord const &extent) { + return ColumnMajorTensorOpMultiplicandCongruous(extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.row(), coord.column())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.contiguous(), coord.strided()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.row(), extent.column())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a row-major view of pitch-linear memory to +/// TensorOpMultiplicand +template +struct RowMajorTensorOpMultiplicandCongruous { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCongruous; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorTensorOpMultiplicandCongruous(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorTensorOpMultiplicandCongruous(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static RowMajorTensorOpMultiplicandCongruous packed(TensorCoord const &extent) { + return RowMajorTensorOpMultiplicandCongruous(extent.column()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.column(), coord.row())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.strided(), coord.contiguous()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.column(), extent.row())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template based on element size (in bits) - defined in terms of pitch-linear +/// memory and Crosswise size (in elements). +template +struct TensorOpMultiplicandCrosswise { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicand; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + static int const kCrosswise = Base::kCrosswise; + static int const kFactor = Base::kFactor; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + + private: + // + // Data members + // + + Base layout_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandCrosswise(Index ldm = 0) : layout_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandCrosswise(Stride stride) : layout_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static TensorOpMultiplicandCrosswise packed(TensorCoord const &extent) { + return TensorOpMultiplicandCrosswise(extent[0]); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(coord); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return coord; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return layout_.stride(); } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return layout_.stride(); } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(extent); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a column-major view of pitch-linear memory to +/// TensorOpMultiplicandCrosswise +template +struct ColumnMajorTensorOpMultiplicandCrosswise { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCrosswise; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + + private: + // + // Data members + // + + Base layout_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorTensorOpMultiplicandCrosswise(Index ldm = 0) : layout_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorTensorOpMultiplicandCrosswise(Stride stride) : layout_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static ColumnMajorTensorOpMultiplicandCrosswise packed( + TensorCoord const &extent) { + return ColumnMajorTensorOpMultiplicandCrosswise(extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.row(), coord.column())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.contiguous(), coord.strided()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return layout_.stride(); } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return layout_.stride(); } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.row(), extent.column())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a row-major view of pitch-linear memory to +/// TensorOpMultiplicandCrosswise +template +struct RowMajorTensorOpMultiplicandCrosswise { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCrosswise; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + + private: + // + // Data members + // + + Base layout_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorTensorOpMultiplicandCrosswise(Index ldm = 0) : layout_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorTensorOpMultiplicandCrosswise(Stride stride) : layout_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static RowMajorTensorOpMultiplicandCrosswise packed( + TensorCoord const &extent) { + return RowMajorTensorOpMultiplicandCrosswise(extent.column()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.column(), coord.row())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.strided(), coord.contiguous()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return layout_.stride(); } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return layout_.stride(); } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.column(), extent.row())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template based on element size (in bits) - defined in terms of pitch-linear memory. +template +struct TensorOpMultiplicandColumnMajorInterleaved { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + /// This layout is optimized for 128b accesses + static int const kAccessSize = 128; + + // + // Static constants + // + + static int const kElementSize = ElementSize; + static int const kElementsPerAccess = kAccessSize / kElementSize; + + //static int const kThreadBlockStrided = ThreadBlockStrided; + static int const kInterleavedK = InterleavedK; + +private: + + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandColumnMajorInterleaved(Index ldm = 0): stride_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandColumnMajorInterleaved(Stride stride): stride_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static TensorOpMultiplicandColumnMajorInterleaved packed(TensorCoord const &extent) { + return TensorOpMultiplicandColumnMajorInterleaved(extent[0] * kInterleavedK); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + int const rows_per_smem_cache_line = 128 / kInterleavedK; + + int row_id = coord.strided() / rows_per_smem_cache_line; + int col_id = (coord.strided() % rows_per_smem_cache_line) * kInterleavedK + coord.contiguous(); + + int access_block_id = col_id >> 4; + int swizzle_access_block_id = access_block_id ^ (row_id & 1); + + int swizzle_col_id = swizzle_access_block_id << 4; + + return row_id * 128 + swizzle_col_id; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return (extent[1] / kInterleavedK) * stride_[0]; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template based on element size (in bits) - defined in terms of pitch-linear memory. +template +struct TensorOpMultiplicandRowMajorInterleaved { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + /// This layout is optimized for 128b accesses + static int const kAccessSize = 128; + + // + // Static constants + // + + static int const kElementSize = ElementSize; + static int const kElementsPerAccess = kAccessSize / kElementSize; + + //static int const kThreadBlockStrided = ThreadBlockStrided; + static int const kInterleavedK = InterleavedK; + +private: + + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandRowMajorInterleaved(Index ldm = 0): stride_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandRowMajorInterleaved(Stride stride): stride_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static TensorOpMultiplicandRowMajorInterleaved packed(TensorCoord const &extent) { + return TensorOpMultiplicandRowMajorInterleaved(extent[1] * kInterleavedK); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + int const rows_per_smem_cache_line = 128 / kInterleavedK; + + int row_id = coord.strided() / rows_per_smem_cache_line; + int col_id = (coord.strided() % rows_per_smem_cache_line) * kInterleavedK + coord.contiguous(); + + int access_block_id = col_id >> 4; + int swizzle_access_block_id = access_block_id ^ (row_id & 1); + + int swizzle_col_id = swizzle_access_block_id << 4; + + return row_id * 128 + swizzle_col_id; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return (extent[0] / kInterleavedK) * stride_[0]; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a column-major view of pitch-linear memory to +/// TensorOpMultiplicand +template +struct MacaColumnMajorTensorOpMultiplicandCongruous { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCongruous; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + MacaColumnMajorTensorOpMultiplicandCongruous(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + MacaColumnMajorTensorOpMultiplicandCongruous(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static MacaColumnMajorTensorOpMultiplicandCongruous packed(TensorCoord const &extent) { + return MacaColumnMajorTensorOpMultiplicandCongruous(extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.row(), coord.column())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.contiguous(), coord.strided()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.row(), extent.column())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a row-major view of pitch-linear memory to +/// TensorOpMultiplicand +template +struct MacaRowMajorTensorOpMultiplicandCongruous { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCongruous; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + MacaRowMajorTensorOpMultiplicandCongruous(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + MacaRowMajorTensorOpMultiplicandCongruous(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static MacaRowMajorTensorOpMultiplicandCongruous packed(TensorCoord const &extent) { + return MacaRowMajorTensorOpMultiplicandCongruous(extent.column()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.column(), coord.row())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.strided(), coord.contiguous()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.column(), extent.row())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a column-major view of pitch-linear memory to +/// TensorOpMultiplicandCrosswise +template +struct MacaColumnMajorTensorOpMultiplicandCrosswise { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCrosswise; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + + private: + // + // Data members + // + + Base layout_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + MacaColumnMajorTensorOpMultiplicandCrosswise(Index ldm = 0) : layout_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + MacaColumnMajorTensorOpMultiplicandCrosswise(Stride stride) : layout_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static MacaColumnMajorTensorOpMultiplicandCrosswise packed( + TensorCoord const &extent) { + return MacaColumnMajorTensorOpMultiplicandCrosswise(extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.row(), coord.column())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.contiguous(), coord.strided()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return layout_.stride(); } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return layout_.stride(); } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.row(), extent.column())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a row-major view of pitch-linear memory to +/// TensorOpMultiplicandCrosswise +template +struct MacaRowMajorTensorOpMultiplicandCrosswise { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCrosswise; + + /// This layout is optimized for 128b accesses + static int const kAccessSize = Base::kAccessSize; + using TileShape = typename Base::TileShape; + using PartitionShape = typename Base::PartitionShape; + + // + // Static constants + // + + static int const kElementSize = Base::kElementSize; + static int const kElementsPerAccess = Base::kElementsPerAccess; + using PartitionCount = typename Base::PartitionCount; + using AccessCount = typename Base::AccessCount; + + private: + // + // Data members + // + + Base layout_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + MacaRowMajorTensorOpMultiplicandCrosswise(Index ldm = 0) : layout_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + MacaRowMajorTensorOpMultiplicandCrosswise(Stride stride) : layout_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static MacaRowMajorTensorOpMultiplicandCrosswise packed( + TensorCoord const &extent) { + return MacaRowMajorTensorOpMultiplicandCrosswise(extent.column()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.column(), coord.row())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.strided(), coord.contiguous()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return layout_.stride(); } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return layout_.stride(); } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.column(), extent.row())); + } +}; + +} // namespace layout +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/layout/tensor_op_multiplicand_sm80.h b/csrc/mctlass/include/mctlass/layout/tensor_op_multiplicand_sm80.h new file mode 100644 index 0000000..817522b --- /dev/null +++ b/csrc/mctlass/include/mctlass/layout/tensor_op_multiplicand_sm80.h @@ -0,0 +1,1139 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief layouts needed by Ampere fp64 tensor core kernels. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace layout { + +//////////////////////////////////////////////////////////////////////////////// + +/// Template based on element size (in bits) - defined in terms of pitch-linear +/// memory and Crosswise size (in elements). +struct TensorOpMultiplicandCongruous64b { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + + // + // Static constants + // + + static int const kElementSize = 64; + static int const kElementsPerAccess = 1; + + private: + + // + // Data members + // + + /// Stride data member. + Stride stride_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandCongruous64b(Index ldm = 0) : stride_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandCongruous64b(Stride stride) : stride_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static TensorOpMultiplicandCongruous64b packed(TensorCoord const &extent) { + return TensorOpMultiplicandCongruous64b(extent[0]); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + + int tc = coord.contiguous() / 16; + int ts = coord.strided() / 4; + + int c = coord.contiguous() % 16; + int s = coord.strided() % 4; + + + int bank = ((((c & 1) * 4 + (c & 6) / 2)) ^ (s & 1)) * 2 + (c / 8); + int row = (c & 6) / 2; + + bank ^= ((s & 2) * 2); + + LongIndex offset = tc * 16 + bank + (ts * 4 + row) * stride_[0]; + + return offset; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return stride_; } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return stride_; } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return extent[1] * stride_[0]; + } + + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + return TensorCoord(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a column-major view of pitch-linear memory to +/// TensorOpMultiplicand +struct ColumnMajorTensorOpMultiplicandCongruous64b { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCongruous64b; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorTensorOpMultiplicandCongruous64b(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorTensorOpMultiplicandCongruous64b(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static ColumnMajorTensorOpMultiplicandCongruous64b packed(TensorCoord const &extent) { + return ColumnMajorTensorOpMultiplicandCongruous64b(extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.row(), coord.column())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.contiguous(), coord.strided()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.row(), extent.column())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a row-major view of pitch-linear memory to +/// TensorOpMultiplicand +struct RowMajorTensorOpMultiplicandCongruous64b { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCongruous64b; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorTensorOpMultiplicandCongruous64b(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorTensorOpMultiplicandCongruous64b(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static RowMajorTensorOpMultiplicandCongruous64b packed(TensorCoord const &extent) { + return RowMajorTensorOpMultiplicandCongruous64b(extent.column()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.column(), coord.row())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.strided(), coord.contiguous()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.column(), extent.row())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template based on element size (in bits) - defined in terms of pitch-linear +/// memory and Crosswise size (in elements). +struct TensorOpMultiplicand64bCrosswise { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + + // + // Static constants + // + + static int const kElementSize = 64; + static int const kElementsPerAccess = 1; + + private: + + // + // Data members + // + + /// Stride data member. + Stride stride_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicand64bCrosswise(Index ldm = 0) : stride_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicand64bCrosswise(Stride stride) : stride_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static TensorOpMultiplicand64bCrosswise packed(TensorCoord const &extent) { + return TensorOpMultiplicand64bCrosswise(extent[0]); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + + int tc = coord.contiguous() / 16; + int ts = coord.strided() / 16; + + int c = coord.contiguous() % 16; + int s = coord.strided() % 16; + + int k_group = c / 4; + int access_s = s / 2; + + int row = access_s % 4; + int bank = ((k_group & 2) << 2) ^ ((s % 2) << 3) + (c % 4) * 2 + (access_s / 4) ^ (k_group & 1); + + int smem_row = (k_group * 4 + row) + tc * 16; + int smem_col = ts * 16 + bank; + + LongIndex offset = smem_row * stride_[0] + smem_col; + + return offset; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return stride_; } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return stride_; } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return extent[1] * stride_[0]; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template based on element size (in bits) - defined in terms of pitch-linear +/// memory and Crosswise size (in elements). +struct ColumnMajorTensorOpMultiplicand64bCrosswise { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicand64bCrosswise; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorTensorOpMultiplicand64bCrosswise(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorTensorOpMultiplicand64bCrosswise(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static ColumnMajorTensorOpMultiplicand64bCrosswise packed(TensorCoord const &extent) { + return ColumnMajorTensorOpMultiplicand64bCrosswise(extent.column()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.row(), coord.column())); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.row(), extent.column())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template based on element size (in bits) - defined in terms of pitch-linear +/// memory and Crosswise size (in elements). +struct RowMajorTensorOpMultiplicand64bCrosswise { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicand64bCrosswise; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorTensorOpMultiplicand64bCrosswise(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorTensorOpMultiplicand64bCrosswise(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static RowMajorTensorOpMultiplicand64bCrosswise packed(TensorCoord const &extent) { + return RowMajorTensorOpMultiplicand64bCrosswise(extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.column(), coord.row())); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.column(), extent.row())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template based on element size (in bits) - defined in terms of pitch-linear +/// memory and Crosswise size (in elements). +struct TensorOpMultiplicandCongruous128b { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + + // + // Static constants + // + + static int const kElementSize = 128; + static int const kElementsPerAccess = 1; + + private: + + // + // Data members + // + + /// Stride data member. + Stride stride_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandCongruous128b(Index ldm = 0) : stride_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandCongruous128b(Stride stride) : stride_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static TensorOpMultiplicandCongruous128b packed(TensorCoord const &extent) { + return TensorOpMultiplicandCongruous128b(extent[0]); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + + Index tc = coord.contiguous() / 8; + Index ts = coord.strided() / 4; + + Index c = coord.contiguous() % 8; + Index s = coord.strided() % 4; + + Index k_index = (c / 2); + + Index bank = (((c & 1) * 4) | (s ^ k_index)); + + LongIndex offset = tc * 8 + bank + (ts * 4 + k_index) * stride_[0]; + + return offset; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return stride_; } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return stride_; } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return extent[1] * stride_[0]; + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + return TensorCoord(); + } +}; + + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a column-major view of pitch-linear memory to +/// TensorOpMultiplicand +struct ColumnMajorTensorOpMultiplicandCongruous128b { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCongruous128b; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorTensorOpMultiplicandCongruous128b(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorTensorOpMultiplicandCongruous128b(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static ColumnMajorTensorOpMultiplicandCongruous128b packed(TensorCoord const &extent) { + return ColumnMajorTensorOpMultiplicandCongruous128b(extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.row(), coord.column())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.contiguous(), coord.strided()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.row(), extent.column())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a row-major view of pitch-linear memory to +/// TensorOpMultiplicand +struct RowMajorTensorOpMultiplicandCongruous128b { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCongruous128b; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorTensorOpMultiplicandCongruous128b(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorTensorOpMultiplicandCongruous128b(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static RowMajorTensorOpMultiplicandCongruous128b packed(TensorCoord const &extent) { + return RowMajorTensorOpMultiplicandCongruous128b(extent.column()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.column(), coord.row())); + } + + /// Inverse of layout function, mapping linear offset to logical coordinate + MCTLASS_HOST_DEVICE + TensorCoord inverse(LongIndex offset) const { + PitchLinearCoord coord = layout_.inverse(offset); + return MatrixCoord(coord.strided(), coord.contiguous()); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.column(), extent.row())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template based on element size (in bits) - defined in terms of pitch-linear +/// memory and Crosswise size (in elements). +struct TensorOpMultiplicandCrosswise128x4 { + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = PitchLinearCoord; + + /// Stride vector + using Stride = Coord; + + // + // Static constants + // + + static int const kElementSize = 128; + static int const kElementsPerAccess = 1; + + private: + + // + // Data members + // + + /// Stride data member. + Stride stride_; + + public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandCrosswise128x4(Index ldm = 0) : stride_(ldm) {} + + /// Ctor + MCTLASS_HOST_DEVICE + TensorOpMultiplicandCrosswise128x4(Stride stride) : stride_(stride) {} + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static TensorOpMultiplicandCrosswise128x4 packed(TensorCoord const &extent) { + return TensorOpMultiplicandCrosswise128x4(extent[0]); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + + Index tc = coord.contiguous() / 8; + Index ts = coord.strided() / 8; + + Index c = coord.contiguous() % 8; + Index s = coord.strided() % 8; + + Index liq = c % 4; + + Index bank = liq + ((s & 1) * 4) ^ (c & 4); + + Index k_index = (c & 4) + (s / 4) * 2 + ((s & 2) / 2); + + LongIndex offset = (tc * 8 + k_index) * stride_[0] + ts * 8 + bank; + + return offset; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { return stride_; } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride &stride() { return stride_; } + + /// Compute the number of contiguous elements needed to store a tensor with + /// the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return extent[1] * stride_[0]; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a column-major view of pitch-linear memory to +/// TensorOpMultiplicand +struct ColumnMajorTensorOpMultiplicandCrosswise128x4 { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCrosswise128x4; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorTensorOpMultiplicandCrosswise128x4(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + ColumnMajorTensorOpMultiplicandCrosswise128x4(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static ColumnMajorTensorOpMultiplicandCrosswise128x4 packed(TensorCoord const &extent) { + return ColumnMajorTensorOpMultiplicandCrosswise128x4(extent.column()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.row(), coord.column())); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.row(), extent.column())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Template mapping a row-major view of pitch-linear memory to +/// TensorOpMultiplicand +struct RowMajorTensorOpMultiplicandCrosswise128x4 { + + /// Logical rank of tensor + static int const kRank = 2; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = MatrixCoord; + + /// Stride vector + using Stride = Coord; + + // + // Invariants + // + + using Base = TensorOpMultiplicandCrosswise128x4; + +private: + + // + // Data members + // + + Base layout_; + +public: + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorTensorOpMultiplicandCrosswise128x4(Index ldm = 0): layout_(ldm) { } + + /// Ctor + MCTLASS_HOST_DEVICE + RowMajorTensorOpMultiplicandCrosswise128x4(Stride stride): layout_(stride) { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static RowMajorTensorOpMultiplicandCrosswise128x4 packed(TensorCoord const &extent) { + return RowMajorTensorOpMultiplicandCrosswise128x4(extent.row()); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (contiguous, strided) + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return layout_(PitchLinearCoord(coord.column(), coord.row())); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &extent) const { + return layout_.capacity(PitchLinearCoord(extent.column(), extent.row())); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace layout +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/layout/vector.h b/csrc/mctlass/include/mctlass/layout/vector.h new file mode 100644 index 0000000..660a9c7 --- /dev/null +++ b/csrc/mctlass/include/mctlass/layout/vector.h @@ -0,0 +1,104 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines layout functions used for rank=1 vectors. +*/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/coord.h" + +namespace mctlass { +namespace layout { + +/// Tensor layout for densely packed vectors. +class PackedVectorLayout { +public: + /// Logical rank of tensor + static int const kRank = 1; + + /// Rank of stride vector + static int const kStrideRank = 1; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = Coord; + + /// Stride vector + using Stride = Coord; + +private: + + // + // No actual stride vector stored + // + +public: + + // + // Methods + // + + MCTLASS_HOST_DEVICE + PackedVectorLayout() { } + + /// Helper returns a layout to a tightly packed tensor + MCTLASS_HOST_DEVICE + static PackedVectorLayout packed(TensorCoord const &size) { + return PackedVectorLayout(); + } + + /// Returns the offset of a coordinate in linear memory + MCTLASS_HOST_DEVICE + LongIndex operator()(TensorCoord const &coord) const { + return coord[0]; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return make_Coord(1); + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &size) const { + return size[0]; + } +}; + +} // namespace layout +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/matrix.h b/csrc/mctlass/include/mctlass/matrix.h new file mode 100644 index 0000000..7dcad3a --- /dev/null +++ b/csrc/mctlass/include/mctlass/matrix.h @@ -0,0 +1,14129 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/* + \file + \brief Matrix classes with value semantics. +*/ + +#pragma once + +#if !defined(__MACACC_RTC__) +#include +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/fast_math.h" +#include "mctlass/layout/matrix.h" + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Primary template with partial specializations to follow +template struct Matrix; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 1-by-2 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 1; + + /// Number of columns in matrix + static int const kColumns = 2; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 2; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 1-by-2 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 1-by-2 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, Element _0_1 + ) { + + data[0] = _0_0; data[1] = _0_1; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[1] = data[1]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 1 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 1 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 1] = m.data[1]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix row(int i) const { + return slice_1x2(i, 0); + } + + MCTLASS_HOST_DEVICE + Matrix &set_row(Matrix const &v, int i = 0) { + return set_slice_1x2(v, i, 0); + } + + /// Forms a 1-by-2 matrix by horizontally concatenating an Element with an Element + MCTLASS_HOST_DEVICE + static Matrix hcat(Element lhs, Element rhs) { + return Matrix( + lhs, rhs); + } + + /// Concatenates this matrix with a an Element to form a 1-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Element rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Concatenates this matrix with a a 1-by-2 matrix to form a 1-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Concatenates this matrix with a a 1-by-2 matrix to form a 2-by-2 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Concatenates this matrix with a a 2-by-2 matrix to form a 3-by-2 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Concatenates this matrix with a a 3-by-2 matrix to form a 4-by-2 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Elementwise add operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + result.data[1] = data[1] + rhs.data[1]; + + return result; + } + + /// Elementwise add operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + data[1] += rhs.data[1]; + + return *this; + } + + /// Elementwise subtract operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + result.data[1] = data[1] - rhs.data[1]; + + return result; + } + + /// Elementwise subtract operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + data[1] -= rhs.data[1]; + + return *this; + } + + /// Elementwise multiply operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + result.data[1] = data[1] * rhs.data[1]; + + return result; + } + + /// Scalar multiply operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + result.data[1] = data[1] * s; + + return result; + } + + /// Scalar multiply operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + data[1] *= s; + + return *this; + } + + /// Elementwise divide operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + result.data[1] = data[1] / rhs.data[1]; + + return result; + } + + /// Scalar divide operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + result.data[1] = data[1] / s; + + return result; + } + + /// Scalar divide operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + data[1] /= s; + + return *this; + } + + /// Elementwise divide operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (1-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + data[1] /= rhs.data[1]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + + return m; + } + + /// Matrix product of size 1-by-1-by-2 + MCTLASS_HOST_DEVICE + Element product(Matrix const &rhs, Element accum = Element()) const { + + // k=0 + accum += data[0] * rhs.data[0]; + + // k=1 + accum += data[1] * rhs.data[1]; + + return accum; + } + + /// Matrix product of size 1-by-1-by-2 + MCTLASS_HOST_DEVICE + Element operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 1-by-2-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + + // k=1 + accum.data[0] += data[1] * rhs.data[2]; + accum.data[1] += data[1] * rhs.data[3]; + + return accum; + } + + /// Matrix product of size 1-by-2-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 1-by-2-by-2 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Matrix product of size 1-by-3-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + + // k=1 + accum.data[0] += data[1] * rhs.data[3]; + accum.data[1] += data[1] * rhs.data[4]; + accum.data[2] += data[1] * rhs.data[5]; + + return accum; + } + + /// Matrix product of size 1-by-3-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 1-by-4-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + + // k=1 + accum.data[0] += data[1] * rhs.data[4]; + accum.data[1] += data[1] * rhs.data[5]; + accum.data[2] += data[1] * rhs.data[6]; + accum.data[3] += data[1] * rhs.data[7]; + + return accum; + } + + /// Matrix product of size 1-by-4-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Dot product of vectors with extent 2 + MCTLASS_HOST_DEVICE + Element dot(Matrix const &rhs, Element accum = Element()) const { + + accum += data[0] * rhs.data[0]; + accum += data[1] * rhs.data[1]; + return accum; + } + + /// Dot product of vectors with extent 2 + MCTLASS_HOST_DEVICE + Element dot(Matrix const &rhs, Element accum = Element()) const { + + accum += data[0] * rhs.data[0]; + accum += data[1] * rhs.data[1]; + return accum; + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + + return accum; + } + +}; + +/// Template alias for 1-by-2 matrix +template +using Matrix1x2 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix1x2 make_Matrix1x2( + Element _0_0, Element _0_1 +) { + return Matrix1x2( + _0_0, _0_1 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 1-by-3 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 1; + + /// Number of columns in matrix + static int const kColumns = 3; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 3; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 1-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 1-by-3 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, Element _0_1, Element _0_2 + ) { + + data[0] = _0_0; data[1] = _0_1; data[2] = _0_2; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + m.data[2] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[1] = data[1]; + mt.data[2] = data[2]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 1 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 1 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 2]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 2] = m.data[2]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix row(int i) const { + return slice_1x3(i, 0); + } + + MCTLASS_HOST_DEVICE + Matrix &set_row(Matrix const &v, int i = 0) { + return set_slice_1x3(v, i, 0); + } + + /// Forms a 1-by-3 matrix by horizontally concatenating an Element with a 1-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Element lhs, Matrix const & rhs) { + return Matrix( + lhs, rhs.at(0, 0), rhs.at(0, 1)); + } + + /// Forms a 1-by-3 matrix by horizontally concatenating a 1-by-2 matrix with an Element + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Element rhs) { + return Matrix( + lhs.at(0, 0), lhs.at(0, 1), rhs); + } + + /// Concatenates this matrix with a an Element to form a 1-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Element rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Concatenates this matrix with a a 1-by-3 matrix to form a 2-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Concatenates this matrix with a a 2-by-3 matrix to form a 3-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Concatenates this matrix with a a 3-by-3 matrix to form a 4-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Elementwise add operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + result.data[1] = data[1] + rhs.data[1]; + result.data[2] = data[2] + rhs.data[2]; + + return result; + } + + /// Elementwise add operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + data[1] += rhs.data[1]; + data[2] += rhs.data[2]; + + return *this; + } + + /// Elementwise subtract operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + result.data[1] = data[1] - rhs.data[1]; + result.data[2] = data[2] - rhs.data[2]; + + return result; + } + + /// Elementwise subtract operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + data[1] -= rhs.data[1]; + data[2] -= rhs.data[2]; + + return *this; + } + + /// Elementwise multiply operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + result.data[1] = data[1] * rhs.data[1]; + result.data[2] = data[2] * rhs.data[2]; + + return result; + } + + /// Scalar multiply operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + result.data[1] = data[1] * s; + result.data[2] = data[2] * s; + + return result; + } + + /// Scalar multiply operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + data[1] *= s; + data[2] *= s; + + return *this; + } + + /// Elementwise divide operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + result.data[1] = data[1] / rhs.data[1]; + result.data[2] = data[2] / rhs.data[2]; + + return result; + } + + /// Scalar divide operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + result.data[1] = data[1] / s; + result.data[2] = data[2] / s; + + return result; + } + + /// Scalar divide operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + data[1] /= s; + data[2] /= s; + + return *this; + } + + /// Elementwise divide operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (1-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + data[1] /= rhs.data[1]; + data[2] /= rhs.data[2]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + m.data[2] = -m.data[2]; + + return m; + } + + /// Matrix product of size 1-by-1-by-3 + MCTLASS_HOST_DEVICE + Element product(Matrix const &rhs, Element accum = Element()) const { + + // k=0 + accum += data[0] * rhs.data[0]; + + // k=1 + accum += data[1] * rhs.data[1]; + + // k=2 + accum += data[2] * rhs.data[2]; + + return accum; + } + + /// Matrix product of size 1-by-1-by-3 + MCTLASS_HOST_DEVICE + Element operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 1-by-2-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + + // k=1 + accum.data[0] += data[1] * rhs.data[2]; + accum.data[1] += data[1] * rhs.data[3]; + + // k=2 + accum.data[0] += data[2] * rhs.data[4]; + accum.data[1] += data[2] * rhs.data[5]; + + return accum; + } + + /// Matrix product of size 1-by-2-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 1-by-3-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + + // k=1 + accum.data[0] += data[1] * rhs.data[3]; + accum.data[1] += data[1] * rhs.data[4]; + accum.data[2] += data[1] * rhs.data[5]; + + // k=2 + accum.data[0] += data[2] * rhs.data[6]; + accum.data[1] += data[2] * rhs.data[7]; + accum.data[2] += data[2] * rhs.data[8]; + + return accum; + } + + /// Matrix product of size 1-by-3-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 1-by-3-by-3 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Matrix product of size 1-by-4-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + + // k=1 + accum.data[0] += data[1] * rhs.data[4]; + accum.data[1] += data[1] * rhs.data[5]; + accum.data[2] += data[1] * rhs.data[6]; + accum.data[3] += data[1] * rhs.data[7]; + + // k=2 + accum.data[0] += data[2] * rhs.data[8]; + accum.data[1] += data[2] * rhs.data[9]; + accum.data[2] += data[2] * rhs.data[10]; + accum.data[3] += data[2] * rhs.data[11]; + + return accum; + } + + /// Matrix product of size 1-by-4-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Dot product of vectors with extent 3 + MCTLASS_HOST_DEVICE + Element dot(Matrix const &rhs, Element accum = Element()) const { + + accum += data[0] * rhs.data[0]; + accum += data[1] * rhs.data[1]; + accum += data[2] * rhs.data[2]; + return accum; + } + + /// Dot product of vectors with extent 3 + MCTLASS_HOST_DEVICE + Element dot(Matrix const &rhs, Element accum = Element()) const { + + accum += data[0] * rhs.data[0]; + accum += data[1] * rhs.data[1]; + accum += data[2] * rhs.data[2]; + return accum; + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + accum += data[2]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + accum += data[2] * data[2]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + + return accum; + } + + /// Cross product + MCTLASS_HOST_DEVICE + Matrix cross(Matrix const &rhs) const { + return Matrix( + data[1] * rhs.data[2] - data[2] * rhs.data[1], + data[0] * rhs.data[2] - data[2] * rhs.data[1], + data[0] * rhs.data[1] - data[1] * rhs.data[0] + ); + } + +}; + +/// Template alias for 1-by-3 matrix +template +using Matrix1x3 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix1x3 make_Matrix1x3( + Element _0_0, Element _0_1, Element _0_2 +) { + return Matrix1x3( + _0_0, _0_1, _0_2 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 1-by-4 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 1; + + /// Number of columns in matrix + static int const kColumns = 4; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 4; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 1-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 1-by-4 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, Element _0_1, Element _0_2, Element _0_3 + ) { + + data[0] = _0_0; data[1] = _0_1; data[2] = _0_2; data[3] = _0_3; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + m.data[2] = s; + m.data[3] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[1] = data[1]; + mt.data[2] = data[2]; + mt.data[3] = data[3]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 1 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 1 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x4(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 3]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x4(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 3] = m.data[3]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix row(int i) const { + return slice_1x4(i, 0); + } + + MCTLASS_HOST_DEVICE + Matrix &set_row(Matrix const &v, int i = 0) { + return set_slice_1x4(v, i, 0); + } + + /// Forms a 1-by-4 matrix by horizontally concatenating an Element with a 1-by-3 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Element lhs, Matrix const & rhs) { + return Matrix( + lhs, rhs.at(0, 0), rhs.at(0, 1), rhs.at(0, 2)); + } + + /// Forms a 1-by-4 matrix by horizontally concatenating a 1-by-2 matrix with a 1-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0), rhs.at(0, 1)); + } + + /// Forms a 1-by-4 matrix by horizontally concatenating a 1-by-3 matrix with an Element + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Element rhs) { + return Matrix( + lhs.at(0, 0), lhs.at(0, 1), lhs.at(0, 2), rhs); + } + + /// Concatenates this matrix with a a 1-by-4 matrix to form a 2-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Concatenates this matrix with a a 2-by-4 matrix to form a 3-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Concatenates this matrix with a a 3-by-4 matrix to form a 4-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Elementwise add operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + result.data[1] = data[1] + rhs.data[1]; + result.data[2] = data[2] + rhs.data[2]; + result.data[3] = data[3] + rhs.data[3]; + + return result; + } + + /// Elementwise add operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + data[1] += rhs.data[1]; + data[2] += rhs.data[2]; + data[3] += rhs.data[3]; + + return *this; + } + + /// Elementwise subtract operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + result.data[1] = data[1] - rhs.data[1]; + result.data[2] = data[2] - rhs.data[2]; + result.data[3] = data[3] - rhs.data[3]; + + return result; + } + + /// Elementwise subtract operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + data[1] -= rhs.data[1]; + data[2] -= rhs.data[2]; + data[3] -= rhs.data[3]; + + return *this; + } + + /// Elementwise multiply operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + result.data[1] = data[1] * rhs.data[1]; + result.data[2] = data[2] * rhs.data[2]; + result.data[3] = data[3] * rhs.data[3]; + + return result; + } + + /// Scalar multiply operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + result.data[1] = data[1] * s; + result.data[2] = data[2] * s; + result.data[3] = data[3] * s; + + return result; + } + + /// Scalar multiply operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + data[1] *= s; + data[2] *= s; + data[3] *= s; + + return *this; + } + + /// Elementwise divide operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + result.data[1] = data[1] / rhs.data[1]; + result.data[2] = data[2] / rhs.data[2]; + result.data[3] = data[3] / rhs.data[3]; + + return result; + } + + /// Scalar divide operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + result.data[1] = data[1] / s; + result.data[2] = data[2] / s; + result.data[3] = data[3] / s; + + return result; + } + + /// Scalar divide operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + data[1] /= s; + data[2] /= s; + data[3] /= s; + + return *this; + } + + /// Elementwise divide operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (1-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + data[1] /= rhs.data[1]; + data[2] /= rhs.data[2]; + data[3] /= rhs.data[3]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + m.data[2] = -m.data[2]; + m.data[3] = -m.data[3]; + + return m; + } + + /// Matrix product of size 1-by-1-by-4 + MCTLASS_HOST_DEVICE + Element product(Matrix const &rhs, Element accum = Element()) const { + + // k=0 + accum += data[0] * rhs.data[0]; + + // k=1 + accum += data[1] * rhs.data[1]; + + // k=2 + accum += data[2] * rhs.data[2]; + + // k=3 + accum += data[3] * rhs.data[3]; + + return accum; + } + + /// Matrix product of size 1-by-1-by-4 + MCTLASS_HOST_DEVICE + Element operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 1-by-2-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + + // k=1 + accum.data[0] += data[1] * rhs.data[2]; + accum.data[1] += data[1] * rhs.data[3]; + + // k=2 + accum.data[0] += data[2] * rhs.data[4]; + accum.data[1] += data[2] * rhs.data[5]; + + // k=3 + accum.data[0] += data[3] * rhs.data[6]; + accum.data[1] += data[3] * rhs.data[7]; + + return accum; + } + + /// Matrix product of size 1-by-2-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 1-by-3-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + + // k=1 + accum.data[0] += data[1] * rhs.data[3]; + accum.data[1] += data[1] * rhs.data[4]; + accum.data[2] += data[1] * rhs.data[5]; + + // k=2 + accum.data[0] += data[2] * rhs.data[6]; + accum.data[1] += data[2] * rhs.data[7]; + accum.data[2] += data[2] * rhs.data[8]; + + // k=3 + accum.data[0] += data[3] * rhs.data[9]; + accum.data[1] += data[3] * rhs.data[10]; + accum.data[2] += data[3] * rhs.data[11]; + + return accum; + } + + /// Matrix product of size 1-by-3-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 1-by-4-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + + // k=1 + accum.data[0] += data[1] * rhs.data[4]; + accum.data[1] += data[1] * rhs.data[5]; + accum.data[2] += data[1] * rhs.data[6]; + accum.data[3] += data[1] * rhs.data[7]; + + // k=2 + accum.data[0] += data[2] * rhs.data[8]; + accum.data[1] += data[2] * rhs.data[9]; + accum.data[2] += data[2] * rhs.data[10]; + accum.data[3] += data[2] * rhs.data[11]; + + // k=3 + accum.data[0] += data[3] * rhs.data[12]; + accum.data[1] += data[3] * rhs.data[13]; + accum.data[2] += data[3] * rhs.data[14]; + accum.data[3] += data[3] * rhs.data[15]; + + return accum; + } + + /// Matrix product of size 1-by-4-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 1-by-4-by-4 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Dot product of vectors with extent 4 + MCTLASS_HOST_DEVICE + Element dot(Matrix const &rhs, Element accum = Element()) const { + + accum += data[0] * rhs.data[0]; + accum += data[1] * rhs.data[1]; + accum += data[2] * rhs.data[2]; + accum += data[3] * rhs.data[3]; + return accum; + } + + /// Dot product of vectors with extent 4 + MCTLASS_HOST_DEVICE + Element dot(Matrix const &rhs, Element accum = Element()) const { + + accum += data[0] * rhs.data[0]; + accum += data[1] * rhs.data[1]; + accum += data[2] * rhs.data[2]; + accum += data[3] * rhs.data[3]; + return accum; + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + accum += data[2]; + accum += data[3]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + accum += data[2] * data[2]; + accum += data[3] * data[3]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + + return accum; + } + +}; + +/// Template alias for 1-by-4 matrix +template +using Matrix1x4 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix1x4 make_Matrix1x4( + Element _0_0, Element _0_1, Element _0_2, Element _0_3 +) { + return Matrix1x4( + _0_0, _0_1, _0_2, _0_3 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 2-by-1 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 2; + + /// Number of columns in matrix + static int const kColumns = 1; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 2; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 2-by-1 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 2-by-1 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, + Element _1_0 + ) { + + data[0] = _0_0; + data[1] = _1_0; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[1] = data[1]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 2 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 2 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 1 + j + 0]; + m.data[1] = data[i * 1 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 1 + j + 0] = m.data[0]; + data[i * 1 + j + 1] = m.data[1]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix column(int j) const { + return slice_2x1(0, j); + } + + MCTLASS_HOST_DEVICE + Matrix &set_column(Matrix const &v, int j =0) { + return set_slice_2x1(v, 0, j); + } + + /// Concatenates this matrix with a a 2-by-1 matrix to form a 2-by-2 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Concatenates this matrix with a a 2-by-2 matrix to form a 2-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Concatenates this matrix with a a 2-by-3 matrix to form a 2-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Forms a 2-by-1 matrix by vertically concatenating an Element with an Element + MCTLASS_HOST_DEVICE + static Matrix vcat(Element upper, Element lower) { + return Matrix( + upper + , lower); + } + + /// Concatenates this matrix with a an Element to form a 3-by-1 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Element rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Concatenates this matrix with a a 2-by-1 matrix to form a 4-by-1 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Elementwise add operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + + result.data[1] = data[1] + rhs.data[1]; + + return result; + } + + /// Elementwise add operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + + data[1] += rhs.data[1]; + + return *this; + } + + /// Elementwise subtract operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + + result.data[1] = data[1] - rhs.data[1]; + + return result; + } + + /// Elementwise subtract operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + + data[1] -= rhs.data[1]; + + return *this; + } + + /// Elementwise multiply operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + + result.data[1] = data[1] * rhs.data[1]; + + return result; + } + + /// Scalar multiply operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + + result.data[1] = data[1] * s; + + return result; + } + + /// Scalar multiply operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + + data[1] *= s; + + return *this; + } + + /// Elementwise divide operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + + result.data[1] = data[1] / rhs.data[1]; + + return result; + } + + /// Scalar divide operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + + result.data[1] = data[1] / s; + + return result; + } + + /// Scalar divide operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + + data[1] /= s; + + return *this; + } + + /// Elementwise divide operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (2-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + + data[1] /= rhs.data[1]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + + return m; + } + + /// Matrix product of size 2-by-1-by-1 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[1] * rhs.data[0]; + + return accum; + } + + /// Matrix product of size 2-by-1-by-1 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 2-by-1-by-1 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Matrix product of size 2-by-2-by-1 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[1] * rhs.data[0]; + accum.data[3] += data[1] * rhs.data[1]; + + return accum; + } + + /// Matrix product of size 2-by-2-by-1 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 2-by-3-by-1 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[1] * rhs.data[0]; + accum.data[4] += data[1] * rhs.data[1]; + accum.data[5] += data[1] * rhs.data[2]; + + return accum; + } + + /// Matrix product of size 2-by-3-by-1 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 2-by-4-by-1 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + accum.data[4] += data[1] * rhs.data[0]; + accum.data[5] += data[1] * rhs.data[1]; + accum.data[6] += data[1] * rhs.data[2]; + accum.data[7] += data[1] * rhs.data[3]; + + return accum; + } + + /// Matrix product of size 2-by-4-by-1 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Dot product of vectors with extent 2 + MCTLASS_HOST_DEVICE + Element dot(Matrix const &rhs, Element accum = Element()) const { + + accum += data[0] * rhs.data[0]; + accum += data[1] * rhs.data[1]; + return accum; + } + + /// Dot product of vectors with extent 2 + MCTLASS_HOST_DEVICE + Element dot(Matrix const &rhs, Element accum = Element()) const { + + accum += data[0] * rhs.data[0]; + accum += data[1] * rhs.data[1]; + return accum; + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + + return accum; + } + +}; + +/// Template alias for 2-by-1 matrix +template +using Matrix2x1 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix2x1 make_Matrix2x1( + Element _0_0, + Element _1_0 +) { + return Matrix2x1( + _0_0, + _1_0 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 2-by-2 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 2; + + /// Number of columns in matrix + static int const kColumns = 2; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 4; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 2-by-2 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 2-by-2 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, Element _0_1, + Element _1_0, Element _1_1 + ) { + + data[0] = _0_0; data[1] = _0_1; + data[2] = _1_0; data[3] = _1_1; + } + + /// Constucts a 2-by-2 matrix from row vectors + MCTLASS_HOST_DEVICE + Matrix( + Matrix const &row_0, + Matrix const &row_1 + ) { + data[0] = row_0.data[0]; + data[1] = row_0.data[1]; + data[2] = row_1.data[0]; + data[3] = row_1.data[1]; + } + + /// Static method to construct a 2-by-2 matrix from column vectors + MCTLASS_HOST_DEVICE + static Matrix from_columns( + Matrix const &column_0, + Matrix const &column_1 + ) { + Matrix result; + + result.data[0] = column_0.data[0]; + result.data[1] = column_1.data[0]; + result.data[2] = column_0.data[1]; + result.data[3] = column_1.data[1]; + return result; + } + + /// Constructs an identity matrix + MCTLASS_HOST_DEVICE + static Matrix identity() { + Matrix m; + + m.data[0] = Element(1); + m.data[3] = Element(1); + + return m; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + m.data[2] = s; + m.data[3] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[3] = diag.data[1]; + + return m; + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[3] = diag.data[1]; + + return m; + } + + /// Gets an array of diagonal elements + MCTLASS_HOST_DEVICE + Matrix diagonal() const { + Matrix diag; + + diag.data[0] = data[0]; + diag.data[1] = data[3]; + + return diag; + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[2] = data[1]; + mt.data[1] = data[2]; + mt.data[3] = data[3]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 2 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 2 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 1] = m.data[1]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix row(int i) const { + return slice_1x2(i, 0); + } + + MCTLASS_HOST_DEVICE + Matrix &set_row(Matrix const &v, int i = 0) { + return set_slice_1x2(v, i, 0); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 2]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 2] = m.data[1]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix column(int j) const { + return slice_2x1(0, j); + } + + MCTLASS_HOST_DEVICE + Matrix &set_column(Matrix const &v, int j =0) { + return set_slice_2x1(v, 0, j); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 1]; + m.data[2] = data[i * 2 + j + 2]; + m.data[3] = data[i * 2 + j + 3]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 1] = m.data[1]; + data[i * 2 + j + 2] = m.data[2]; + data[i * 2 + j + 3] = m.data[3]; + + return *this; + } + + /// Forms a 2-by-2 matrix by horizontally concatenating a 2-by-1 matrix with a 2-by-1 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), rhs.at(0, 0) + , lhs.at(1, 0), rhs.at(1, 0)); + } + + /// Concatenates this matrix with a a 2-by-1 matrix to form a 2-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Concatenates this matrix with a a 2-by-2 matrix to form a 2-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Forms a 2-by-2 matrix by vertically concatenating a 1-by-2 matrix with a 1-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1) + , lower.at(0, 0), lower.at(0, 1)); + } + + /// Concatenates this matrix with a a 1-by-2 matrix to form a 3-by-2 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Concatenates this matrix with a a 2-by-2 matrix to form a 4-by-2 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Forms a 2-by-2 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Element A, Element B, + Element C, Element D) { + return Matrix( + A, B + , C, D + ); + } + + /// Elementwise add operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + result.data[1] = data[1] + rhs.data[1]; + + result.data[2] = data[2] + rhs.data[2]; + result.data[3] = data[3] + rhs.data[3]; + + return result; + } + + /// Elementwise add operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + data[1] += rhs.data[1]; + + data[2] += rhs.data[2]; + data[3] += rhs.data[3]; + + return *this; + } + + /// Elementwise subtract operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + result.data[1] = data[1] - rhs.data[1]; + + result.data[2] = data[2] - rhs.data[2]; + result.data[3] = data[3] - rhs.data[3]; + + return result; + } + + /// Elementwise subtract operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + data[1] -= rhs.data[1]; + + data[2] -= rhs.data[2]; + data[3] -= rhs.data[3]; + + return *this; + } + + /// Elementwise multiply operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + result.data[1] = data[1] * rhs.data[1]; + + result.data[2] = data[2] * rhs.data[2]; + result.data[3] = data[3] * rhs.data[3]; + + return result; + } + + /// Scalar multiply operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + result.data[1] = data[1] * s; + + result.data[2] = data[2] * s; + result.data[3] = data[3] * s; + + return result; + } + + /// Scalar multiply operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + data[1] *= s; + + data[2] *= s; + data[3] *= s; + + return *this; + } + + /// Elementwise divide operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + result.data[1] = data[1] / rhs.data[1]; + + result.data[2] = data[2] / rhs.data[2]; + result.data[3] = data[3] / rhs.data[3]; + + return result; + } + + /// Scalar divide operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + result.data[1] = data[1] / s; + + result.data[2] = data[2] / s; + result.data[3] = data[3] / s; + + return result; + } + + /// Scalar divide operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + data[1] /= s; + + data[2] /= s; + data[3] /= s; + + return *this; + } + + /// Elementwise divide operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (2-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + data[1] /= rhs.data[1]; + + data[2] /= rhs.data[2]; + data[3] /= rhs.data[3]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + m.data[2] = -m.data[2]; + m.data[3] = -m.data[3]; + + return m; + } + + /// Matrix product of size 2-by-1-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[2] * rhs.data[0]; + + // k=1 + accum.data[0] += data[1] * rhs.data[1]; + accum.data[1] += data[3] * rhs.data[1]; + + return accum; + } + + /// Matrix product of size 2-by-1-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 2-by-2-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[2] * rhs.data[0]; + accum.data[3] += data[2] * rhs.data[1]; + + // k=1 + accum.data[0] += data[1] * rhs.data[2]; + accum.data[1] += data[1] * rhs.data[3]; + accum.data[2] += data[3] * rhs.data[2]; + accum.data[3] += data[3] * rhs.data[3]; + + return accum; + } + + /// Matrix product of size 2-by-2-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 2-by-2-by-2 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Matrix product of size 2-by-3-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[2] * rhs.data[0]; + accum.data[4] += data[2] * rhs.data[1]; + accum.data[5] += data[2] * rhs.data[2]; + + // k=1 + accum.data[0] += data[1] * rhs.data[3]; + accum.data[1] += data[1] * rhs.data[4]; + accum.data[2] += data[1] * rhs.data[5]; + accum.data[3] += data[3] * rhs.data[3]; + accum.data[4] += data[3] * rhs.data[4]; + accum.data[5] += data[3] * rhs.data[5]; + + return accum; + } + + /// Matrix product of size 2-by-3-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 2-by-4-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + accum.data[4] += data[2] * rhs.data[0]; + accum.data[5] += data[2] * rhs.data[1]; + accum.data[6] += data[2] * rhs.data[2]; + accum.data[7] += data[2] * rhs.data[3]; + + // k=1 + accum.data[0] += data[1] * rhs.data[4]; + accum.data[1] += data[1] * rhs.data[5]; + accum.data[2] += data[1] * rhs.data[6]; + accum.data[3] += data[1] * rhs.data[7]; + accum.data[4] += data[3] * rhs.data[4]; + accum.data[5] += data[3] * rhs.data[5]; + accum.data[6] += data[3] * rhs.data[6]; + accum.data[7] += data[3] * rhs.data[7]; + + return accum; + } + + /// Matrix product of size 2-by-4-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + accum += data[2]; + accum += data[3]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + accum += data[2] * data[2]; + accum += data[3] * data[3]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + accum += data[3]; + + return accum; + } + + /// Returns 2-by-2 rotation matrix + MCTLASS_HOST_DEVICE + static Matrix rotation(Element theta) { + Element c = fast_cos(theta); + Element s = fast_sin(theta); + + return Matrix( + c, -s, + s, c + ); + } + + /// Computes the determinant of a 2-by-2 matrix + MCTLASS_HOST_DEVICE + Element determinant(Element accum = Element()) const { + accum += data[0] * data[3] - data[1] * data[2]; + + return accum; + } + + /// Computes the inverse of a 2-by-2 matrix given + /// the matrix's determinant + MCTLASS_HOST_DEVICE + Matrix inverse(Element det) const { + return Matrix( + data[3], -data[1], + -data[2], data[0] + ) * (Element(1) / det); + } + + /// Computes the inverse of a 2-by-2 matrix. + MCTLASS_HOST_DEVICE + Matrix inverse() const { + return inverse(determinant()); + } + +}; + +/// Template alias for 2-by-2 matrix +template +using Matrix2x2 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix2x2 make_Matrix2x2( + Element _0_0, Element _0_1, + Element _1_0, Element _1_1 +) { + return Matrix2x2( + _0_0, _0_1, + _1_0, _1_1 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 2-by-3 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 2; + + /// Number of columns in matrix + static int const kColumns = 3; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 6; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 2-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 2-by-3 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, Element _0_1, Element _0_2, + Element _1_0, Element _1_1, Element _1_2 + ) { + + data[0] = _0_0; data[1] = _0_1; data[2] = _0_2; + data[3] = _1_0; data[4] = _1_1; data[5] = _1_2; + } + + /// Constucts a 2-by-3 matrix from row vectors + MCTLASS_HOST_DEVICE + Matrix( + Matrix const &row_0, + Matrix const &row_1 + ) { + data[0] = row_0.data[0]; + data[1] = row_0.data[1]; + data[2] = row_0.data[2]; + data[3] = row_1.data[0]; + data[4] = row_1.data[1]; + data[5] = row_1.data[2]; + } + + /// Static method to construct a 2-by-3 matrix from column vectors + MCTLASS_HOST_DEVICE + static Matrix from_columns( + Matrix const &column_0, + Matrix const &column_1, + Matrix const &column_2 + ) { + Matrix result; + + result.data[0] = column_0.data[0]; + result.data[1] = column_1.data[0]; + result.data[2] = column_2.data[0]; + result.data[3] = column_0.data[1]; + result.data[4] = column_1.data[1]; + result.data[5] = column_2.data[1]; + return result; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + m.data[2] = s; + m.data[3] = s; + m.data[4] = s; + m.data[5] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[3] = diag.data[1]; + + return m; + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[3] = diag.data[1]; + + return m; + } + + /// Gets an array of diagonal elements + MCTLASS_HOST_DEVICE + Matrix diagonal() const { + Matrix diag; + + diag.data[0] = data[0]; + diag.data[1] = data[3]; + + return diag; + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[2] = data[1]; + mt.data[4] = data[2]; + mt.data[1] = data[3]; + mt.data[3] = data[4]; + mt.data[5] = data[5]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 2 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 2 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 2]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 2] = m.data[2]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix row(int i) const { + return slice_1x3(i, 0); + } + + MCTLASS_HOST_DEVICE + Matrix &set_row(Matrix const &v, int i = 0) { + return set_slice_1x3(v, i, 0); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 3]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 3] = m.data[1]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix column(int j) const { + return slice_2x1(0, j); + } + + MCTLASS_HOST_DEVICE + Matrix &set_column(Matrix const &v, int j =0) { + return set_slice_2x1(v, 0, j); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 3]; + m.data[3] = data[i * 3 + j + 4]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 3] = m.data[2]; + data[i * 3 + j + 4] = m.data[3]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 2]; + m.data[3] = data[i * 3 + j + 3]; + m.data[4] = data[i * 3 + j + 4]; + m.data[5] = data[i * 3 + j + 5]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 2] = m.data[2]; + data[i * 3 + j + 3] = m.data[3]; + data[i * 3 + j + 4] = m.data[4]; + data[i * 3 + j + 5] = m.data[5]; + + return *this; + } + + /// Forms a 2-by-3 matrix by horizontally concatenating a 2-by-1 matrix with a 2-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), rhs.at(0, 0), rhs.at(0, 1) + , lhs.at(1, 0), rhs.at(1, 0), rhs.at(1, 1)); + } + + /// Forms a 2-by-3 matrix by horizontally concatenating a 2-by-2 matrix with a 2-by-1 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0) + , lhs.at(1, 0), lhs.at(1, 1), rhs.at(1, 0)); + } + + /// Concatenates this matrix with a a 2-by-1 matrix to form a 2-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Forms a 2-by-3 matrix by vertically concatenating a 1-by-3 matrix with a 1-by-3 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1), upper.at(0, 2) + , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2)); + } + + /// Concatenates this matrix with a a 1-by-3 matrix to form a 3-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Concatenates this matrix with a a 2-by-3 matrix to form a 4-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Forms a 2-by-3 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Element A, Matrix const & B, + Element C, Matrix const & D) { + return Matrix( + A, B.at(0, 0), B.at(0, 1) + , C, D.at(0, 0), D.at(0, 1) + ); + } + + /// Forms a 2-by-3 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Element B, + Matrix const & C, Element D) { + return Matrix( + A.at(0, 0), A.at(0, 1), B + , C.at(0, 0), C.at(0, 1), D + ); + } + + /// Elementwise add operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + result.data[1] = data[1] + rhs.data[1]; + result.data[2] = data[2] + rhs.data[2]; + + result.data[3] = data[3] + rhs.data[3]; + result.data[4] = data[4] + rhs.data[4]; + result.data[5] = data[5] + rhs.data[5]; + + return result; + } + + /// Elementwise add operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + data[1] += rhs.data[1]; + data[2] += rhs.data[2]; + + data[3] += rhs.data[3]; + data[4] += rhs.data[4]; + data[5] += rhs.data[5]; + + return *this; + } + + /// Elementwise subtract operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + result.data[1] = data[1] - rhs.data[1]; + result.data[2] = data[2] - rhs.data[2]; + + result.data[3] = data[3] - rhs.data[3]; + result.data[4] = data[4] - rhs.data[4]; + result.data[5] = data[5] - rhs.data[5]; + + return result; + } + + /// Elementwise subtract operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + data[1] -= rhs.data[1]; + data[2] -= rhs.data[2]; + + data[3] -= rhs.data[3]; + data[4] -= rhs.data[4]; + data[5] -= rhs.data[5]; + + return *this; + } + + /// Elementwise multiply operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + result.data[1] = data[1] * rhs.data[1]; + result.data[2] = data[2] * rhs.data[2]; + + result.data[3] = data[3] * rhs.data[3]; + result.data[4] = data[4] * rhs.data[4]; + result.data[5] = data[5] * rhs.data[5]; + + return result; + } + + /// Scalar multiply operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + result.data[1] = data[1] * s; + result.data[2] = data[2] * s; + + result.data[3] = data[3] * s; + result.data[4] = data[4] * s; + result.data[5] = data[5] * s; + + return result; + } + + /// Scalar multiply operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + data[1] *= s; + data[2] *= s; + + data[3] *= s; + data[4] *= s; + data[5] *= s; + + return *this; + } + + /// Elementwise divide operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + result.data[1] = data[1] / rhs.data[1]; + result.data[2] = data[2] / rhs.data[2]; + + result.data[3] = data[3] / rhs.data[3]; + result.data[4] = data[4] / rhs.data[4]; + result.data[5] = data[5] / rhs.data[5]; + + return result; + } + + /// Scalar divide operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + result.data[1] = data[1] / s; + result.data[2] = data[2] / s; + + result.data[3] = data[3] / s; + result.data[4] = data[4] / s; + result.data[5] = data[5] / s; + + return result; + } + + /// Scalar divide operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + data[1] /= s; + data[2] /= s; + + data[3] /= s; + data[4] /= s; + data[5] /= s; + + return *this; + } + + /// Elementwise divide operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (2-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + data[1] /= rhs.data[1]; + data[2] /= rhs.data[2]; + + data[3] /= rhs.data[3]; + data[4] /= rhs.data[4]; + data[5] /= rhs.data[5]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + m.data[2] = -m.data[2]; + m.data[3] = -m.data[3]; + m.data[4] = -m.data[4]; + m.data[5] = -m.data[5]; + + return m; + } + + /// Matrix product of size 2-by-1-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[3] * rhs.data[0]; + + // k=1 + accum.data[0] += data[1] * rhs.data[1]; + accum.data[1] += data[4] * rhs.data[1]; + + // k=2 + accum.data[0] += data[2] * rhs.data[2]; + accum.data[1] += data[5] * rhs.data[2]; + + return accum; + } + + /// Matrix product of size 2-by-1-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 2-by-2-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[3] * rhs.data[0]; + accum.data[3] += data[3] * rhs.data[1]; + + // k=1 + accum.data[0] += data[1] * rhs.data[2]; + accum.data[1] += data[1] * rhs.data[3]; + accum.data[2] += data[4] * rhs.data[2]; + accum.data[3] += data[4] * rhs.data[3]; + + // k=2 + accum.data[0] += data[2] * rhs.data[4]; + accum.data[1] += data[2] * rhs.data[5]; + accum.data[2] += data[5] * rhs.data[4]; + accum.data[3] += data[5] * rhs.data[5]; + + return accum; + } + + /// Matrix product of size 2-by-2-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 2-by-3-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[3] * rhs.data[0]; + accum.data[4] += data[3] * rhs.data[1]; + accum.data[5] += data[3] * rhs.data[2]; + + // k=1 + accum.data[0] += data[1] * rhs.data[3]; + accum.data[1] += data[1] * rhs.data[4]; + accum.data[2] += data[1] * rhs.data[5]; + accum.data[3] += data[4] * rhs.data[3]; + accum.data[4] += data[4] * rhs.data[4]; + accum.data[5] += data[4] * rhs.data[5]; + + // k=2 + accum.data[0] += data[2] * rhs.data[6]; + accum.data[1] += data[2] * rhs.data[7]; + accum.data[2] += data[2] * rhs.data[8]; + accum.data[3] += data[5] * rhs.data[6]; + accum.data[4] += data[5] * rhs.data[7]; + accum.data[5] += data[5] * rhs.data[8]; + + return accum; + } + + /// Matrix product of size 2-by-3-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 2-by-3-by-3 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Matrix product of size 2-by-4-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + accum.data[4] += data[3] * rhs.data[0]; + accum.data[5] += data[3] * rhs.data[1]; + accum.data[6] += data[3] * rhs.data[2]; + accum.data[7] += data[3] * rhs.data[3]; + + // k=1 + accum.data[0] += data[1] * rhs.data[4]; + accum.data[1] += data[1] * rhs.data[5]; + accum.data[2] += data[1] * rhs.data[6]; + accum.data[3] += data[1] * rhs.data[7]; + accum.data[4] += data[4] * rhs.data[4]; + accum.data[5] += data[4] * rhs.data[5]; + accum.data[6] += data[4] * rhs.data[6]; + accum.data[7] += data[4] * rhs.data[7]; + + // k=2 + accum.data[0] += data[2] * rhs.data[8]; + accum.data[1] += data[2] * rhs.data[9]; + accum.data[2] += data[2] * rhs.data[10]; + accum.data[3] += data[2] * rhs.data[11]; + accum.data[4] += data[5] * rhs.data[8]; + accum.data[5] += data[5] * rhs.data[9]; + accum.data[6] += data[5] * rhs.data[10]; + accum.data[7] += data[5] * rhs.data[11]; + + return accum; + } + + /// Matrix product of size 2-by-4-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + accum += data[2]; + accum += data[3]; + accum += data[4]; + accum += data[5]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + accum += data[2] * data[2]; + accum += data[3] * data[3]; + accum += data[4] * data[4]; + accum += data[5] * data[5]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + accum += data[4]; + + return accum; + } + +}; + +/// Template alias for 2-by-3 matrix +template +using Matrix2x3 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix2x3 make_Matrix2x3( + Element _0_0, Element _0_1, Element _0_2, + Element _1_0, Element _1_1, Element _1_2 +) { + return Matrix2x3( + _0_0, _0_1, _0_2, + _1_0, _1_1, _1_2 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 2-by-4 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 2; + + /// Number of columns in matrix + static int const kColumns = 4; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 8; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 2-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 2-by-4 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, Element _0_1, Element _0_2, Element _0_3, + Element _1_0, Element _1_1, Element _1_2, Element _1_3 + ) { + + data[0] = _0_0; data[1] = _0_1; data[2] = _0_2; data[3] = _0_3; + data[4] = _1_0; data[5] = _1_1; data[6] = _1_2; data[7] = _1_3; + } + + /// Constucts a 2-by-4 matrix from row vectors + MCTLASS_HOST_DEVICE + Matrix( + Matrix const &row_0, + Matrix const &row_1 + ) { + data[0] = row_0.data[0]; + data[1] = row_0.data[1]; + data[2] = row_0.data[2]; + data[3] = row_0.data[3]; + data[4] = row_1.data[0]; + data[5] = row_1.data[1]; + data[6] = row_1.data[2]; + data[7] = row_1.data[3]; + } + + /// Static method to construct a 2-by-4 matrix from column vectors + MCTLASS_HOST_DEVICE + static Matrix from_columns( + Matrix const &column_0, + Matrix const &column_1, + Matrix const &column_2, + Matrix const &column_3 + ) { + Matrix result; + + result.data[0] = column_0.data[0]; + result.data[1] = column_1.data[0]; + result.data[2] = column_2.data[0]; + result.data[3] = column_3.data[0]; + result.data[4] = column_0.data[1]; + result.data[5] = column_1.data[1]; + result.data[6] = column_2.data[1]; + result.data[7] = column_3.data[1]; + return result; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + m.data[2] = s; + m.data[3] = s; + m.data[4] = s; + m.data[5] = s; + m.data[6] = s; + m.data[7] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[3] = diag.data[1]; + + return m; + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[3] = diag.data[1]; + + return m; + } + + /// Gets an array of diagonal elements + MCTLASS_HOST_DEVICE + Matrix diagonal() const { + Matrix diag; + + diag.data[0] = data[0]; + diag.data[1] = data[3]; + + return diag; + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[2] = data[1]; + mt.data[4] = data[2]; + mt.data[6] = data[3]; + mt.data[1] = data[4]; + mt.data[3] = data[5]; + mt.data[5] = data[6]; + mt.data[7] = data[7]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 2 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 2 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x4(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 3]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x4(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 3] = m.data[3]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix row(int i) const { + return slice_1x4(i, 0); + } + + MCTLASS_HOST_DEVICE + Matrix &set_row(Matrix const &v, int i = 0) { + return set_slice_1x4(v, i, 0); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 4]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 4] = m.data[1]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix column(int j) const { + return slice_2x1(0, j); + } + + MCTLASS_HOST_DEVICE + Matrix &set_column(Matrix const &v, int j =0) { + return set_slice_2x1(v, 0, j); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 4]; + m.data[3] = data[i * 4 + j + 5]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 4] = m.data[2]; + data[i * 4 + j + 5] = m.data[3]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 4]; + m.data[4] = data[i * 4 + j + 5]; + m.data[5] = data[i * 4 + j + 6]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 4] = m.data[3]; + data[i * 4 + j + 5] = m.data[4]; + data[i * 4 + j + 6] = m.data[5]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x4(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 3]; + m.data[4] = data[i * 4 + j + 4]; + m.data[5] = data[i * 4 + j + 5]; + m.data[6] = data[i * 4 + j + 6]; + m.data[7] = data[i * 4 + j + 7]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x4(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 3] = m.data[3]; + data[i * 4 + j + 4] = m.data[4]; + data[i * 4 + j + 5] = m.data[5]; + data[i * 4 + j + 6] = m.data[6]; + data[i * 4 + j + 7] = m.data[7]; + + return *this; + } + + /// Forms a 2-by-4 matrix by horizontally concatenating a 2-by-1 matrix with a 2-by-3 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), rhs.at(0, 0), rhs.at(0, 1), rhs.at(0, 2) + , lhs.at(1, 0), rhs.at(1, 0), rhs.at(1, 1), rhs.at(1, 2)); + } + + /// Forms a 2-by-4 matrix by horizontally concatenating a 2-by-2 matrix with a 2-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0), rhs.at(0, 1) + , lhs.at(1, 0), lhs.at(1, 1), rhs.at(1, 0), rhs.at(1, 1)); + } + + /// Forms a 2-by-4 matrix by horizontally concatenating a 2-by-3 matrix with a 2-by-1 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), lhs.at(0, 1), lhs.at(0, 2), rhs.at(0, 0) + , lhs.at(1, 0), lhs.at(1, 1), lhs.at(1, 2), rhs.at(1, 0)); + } + + /// Forms a 2-by-4 matrix by vertically concatenating a 1-by-4 matrix with a 1-by-4 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1), upper.at(0, 2), upper.at(0, 3) + , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2), lower.at(0, 3)); + } + + /// Concatenates this matrix with a a 1-by-4 matrix to form a 3-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Concatenates this matrix with a a 2-by-4 matrix to form a 4-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Forms a 2-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Element A, Matrix const & B, + Element C, Matrix const & D) { + return Matrix( + A, B.at(0, 0), B.at(0, 1), B.at(0, 2) + , C, D.at(0, 0), D.at(0, 1), D.at(0, 2) + ); + } + + /// Forms a 2-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), A.at(0, 1), B.at(0, 0), B.at(0, 1) + , C.at(0, 0), C.at(0, 1), D.at(0, 0), D.at(0, 1) + ); + } + + /// Forms a 2-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Element B, + Matrix const & C, Element D) { + return Matrix( + A.at(0, 0), A.at(0, 1), A.at(0, 2), B + , C.at(0, 0), C.at(0, 1), C.at(0, 2), D + ); + } + + /// Elementwise add operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + result.data[1] = data[1] + rhs.data[1]; + result.data[2] = data[2] + rhs.data[2]; + result.data[3] = data[3] + rhs.data[3]; + + result.data[4] = data[4] + rhs.data[4]; + result.data[5] = data[5] + rhs.data[5]; + result.data[6] = data[6] + rhs.data[6]; + result.data[7] = data[7] + rhs.data[7]; + + return result; + } + + /// Elementwise add operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + data[1] += rhs.data[1]; + data[2] += rhs.data[2]; + data[3] += rhs.data[3]; + + data[4] += rhs.data[4]; + data[5] += rhs.data[5]; + data[6] += rhs.data[6]; + data[7] += rhs.data[7]; + + return *this; + } + + /// Elementwise subtract operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + result.data[1] = data[1] - rhs.data[1]; + result.data[2] = data[2] - rhs.data[2]; + result.data[3] = data[3] - rhs.data[3]; + + result.data[4] = data[4] - rhs.data[4]; + result.data[5] = data[5] - rhs.data[5]; + result.data[6] = data[6] - rhs.data[6]; + result.data[7] = data[7] - rhs.data[7]; + + return result; + } + + /// Elementwise subtract operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + data[1] -= rhs.data[1]; + data[2] -= rhs.data[2]; + data[3] -= rhs.data[3]; + + data[4] -= rhs.data[4]; + data[5] -= rhs.data[5]; + data[6] -= rhs.data[6]; + data[7] -= rhs.data[7]; + + return *this; + } + + /// Elementwise multiply operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + result.data[1] = data[1] * rhs.data[1]; + result.data[2] = data[2] * rhs.data[2]; + result.data[3] = data[3] * rhs.data[3]; + + result.data[4] = data[4] * rhs.data[4]; + result.data[5] = data[5] * rhs.data[5]; + result.data[6] = data[6] * rhs.data[6]; + result.data[7] = data[7] * rhs.data[7]; + + return result; + } + + /// Scalar multiply operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + result.data[1] = data[1] * s; + result.data[2] = data[2] * s; + result.data[3] = data[3] * s; + + result.data[4] = data[4] * s; + result.data[5] = data[5] * s; + result.data[6] = data[6] * s; + result.data[7] = data[7] * s; + + return result; + } + + /// Scalar multiply operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + data[1] *= s; + data[2] *= s; + data[3] *= s; + + data[4] *= s; + data[5] *= s; + data[6] *= s; + data[7] *= s; + + return *this; + } + + /// Elementwise divide operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + result.data[1] = data[1] / rhs.data[1]; + result.data[2] = data[2] / rhs.data[2]; + result.data[3] = data[3] / rhs.data[3]; + + result.data[4] = data[4] / rhs.data[4]; + result.data[5] = data[5] / rhs.data[5]; + result.data[6] = data[6] / rhs.data[6]; + result.data[7] = data[7] / rhs.data[7]; + + return result; + } + + /// Scalar divide operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + result.data[1] = data[1] / s; + result.data[2] = data[2] / s; + result.data[3] = data[3] / s; + + result.data[4] = data[4] / s; + result.data[5] = data[5] / s; + result.data[6] = data[6] / s; + result.data[7] = data[7] / s; + + return result; + } + + /// Scalar divide operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + data[1] /= s; + data[2] /= s; + data[3] /= s; + + data[4] /= s; + data[5] /= s; + data[6] /= s; + data[7] /= s; + + return *this; + } + + /// Elementwise divide operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (2-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + data[1] /= rhs.data[1]; + data[2] /= rhs.data[2]; + data[3] /= rhs.data[3]; + + data[4] /= rhs.data[4]; + data[5] /= rhs.data[5]; + data[6] /= rhs.data[6]; + data[7] /= rhs.data[7]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + m.data[2] = -m.data[2]; + m.data[3] = -m.data[3]; + m.data[4] = -m.data[4]; + m.data[5] = -m.data[5]; + m.data[6] = -m.data[6]; + m.data[7] = -m.data[7]; + + return m; + } + + /// Matrix product of size 2-by-1-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[4] * rhs.data[0]; + + // k=1 + accum.data[0] += data[1] * rhs.data[1]; + accum.data[1] += data[5] * rhs.data[1]; + + // k=2 + accum.data[0] += data[2] * rhs.data[2]; + accum.data[1] += data[6] * rhs.data[2]; + + // k=3 + accum.data[0] += data[3] * rhs.data[3]; + accum.data[1] += data[7] * rhs.data[3]; + + return accum; + } + + /// Matrix product of size 2-by-1-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 2-by-2-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[4] * rhs.data[0]; + accum.data[3] += data[4] * rhs.data[1]; + + // k=1 + accum.data[0] += data[1] * rhs.data[2]; + accum.data[1] += data[1] * rhs.data[3]; + accum.data[2] += data[5] * rhs.data[2]; + accum.data[3] += data[5] * rhs.data[3]; + + // k=2 + accum.data[0] += data[2] * rhs.data[4]; + accum.data[1] += data[2] * rhs.data[5]; + accum.data[2] += data[6] * rhs.data[4]; + accum.data[3] += data[6] * rhs.data[5]; + + // k=3 + accum.data[0] += data[3] * rhs.data[6]; + accum.data[1] += data[3] * rhs.data[7]; + accum.data[2] += data[7] * rhs.data[6]; + accum.data[3] += data[7] * rhs.data[7]; + + return accum; + } + + /// Matrix product of size 2-by-2-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 2-by-3-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[4] * rhs.data[0]; + accum.data[4] += data[4] * rhs.data[1]; + accum.data[5] += data[4] * rhs.data[2]; + + // k=1 + accum.data[0] += data[1] * rhs.data[3]; + accum.data[1] += data[1] * rhs.data[4]; + accum.data[2] += data[1] * rhs.data[5]; + accum.data[3] += data[5] * rhs.data[3]; + accum.data[4] += data[5] * rhs.data[4]; + accum.data[5] += data[5] * rhs.data[5]; + + // k=2 + accum.data[0] += data[2] * rhs.data[6]; + accum.data[1] += data[2] * rhs.data[7]; + accum.data[2] += data[2] * rhs.data[8]; + accum.data[3] += data[6] * rhs.data[6]; + accum.data[4] += data[6] * rhs.data[7]; + accum.data[5] += data[6] * rhs.data[8]; + + // k=3 + accum.data[0] += data[3] * rhs.data[9]; + accum.data[1] += data[3] * rhs.data[10]; + accum.data[2] += data[3] * rhs.data[11]; + accum.data[3] += data[7] * rhs.data[9]; + accum.data[4] += data[7] * rhs.data[10]; + accum.data[5] += data[7] * rhs.data[11]; + + return accum; + } + + /// Matrix product of size 2-by-3-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 2-by-4-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + accum.data[4] += data[4] * rhs.data[0]; + accum.data[5] += data[4] * rhs.data[1]; + accum.data[6] += data[4] * rhs.data[2]; + accum.data[7] += data[4] * rhs.data[3]; + + // k=1 + accum.data[0] += data[1] * rhs.data[4]; + accum.data[1] += data[1] * rhs.data[5]; + accum.data[2] += data[1] * rhs.data[6]; + accum.data[3] += data[1] * rhs.data[7]; + accum.data[4] += data[5] * rhs.data[4]; + accum.data[5] += data[5] * rhs.data[5]; + accum.data[6] += data[5] * rhs.data[6]; + accum.data[7] += data[5] * rhs.data[7]; + + // k=2 + accum.data[0] += data[2] * rhs.data[8]; + accum.data[1] += data[2] * rhs.data[9]; + accum.data[2] += data[2] * rhs.data[10]; + accum.data[3] += data[2] * rhs.data[11]; + accum.data[4] += data[6] * rhs.data[8]; + accum.data[5] += data[6] * rhs.data[9]; + accum.data[6] += data[6] * rhs.data[10]; + accum.data[7] += data[6] * rhs.data[11]; + + // k=3 + accum.data[0] += data[3] * rhs.data[12]; + accum.data[1] += data[3] * rhs.data[13]; + accum.data[2] += data[3] * rhs.data[14]; + accum.data[3] += data[3] * rhs.data[15]; + accum.data[4] += data[7] * rhs.data[12]; + accum.data[5] += data[7] * rhs.data[13]; + accum.data[6] += data[7] * rhs.data[14]; + accum.data[7] += data[7] * rhs.data[15]; + + return accum; + } + + /// Matrix product of size 2-by-4-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 2-by-4-by-4 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + accum += data[2]; + accum += data[3]; + accum += data[4]; + accum += data[5]; + accum += data[6]; + accum += data[7]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + accum += data[2] * data[2]; + accum += data[3] * data[3]; + accum += data[4] * data[4]; + accum += data[5] * data[5]; + accum += data[6] * data[6]; + accum += data[7] * data[7]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + accum += data[5]; + + return accum; + } + +}; + +/// Template alias for 2-by-4 matrix +template +using Matrix2x4 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix2x4 make_Matrix2x4( + Element _0_0, Element _0_1, Element _0_2, Element _0_3, + Element _1_0, Element _1_1, Element _1_2, Element _1_3 +) { + return Matrix2x4( + _0_0, _0_1, _0_2, _0_3, + _1_0, _1_1, _1_2, _1_3 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 3-by-1 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 3; + + /// Number of columns in matrix + static int const kColumns = 1; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 3; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 3-by-1 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 3-by-1 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, + Element _1_0, + Element _2_0 + ) { + + data[0] = _0_0; + data[1] = _1_0; + data[2] = _2_0; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + m.data[2] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[1] = data[1]; + mt.data[2] = data[2]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 3 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 3 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 1 + j + 0]; + m.data[1] = data[i * 1 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 1 + j + 0] = m.data[0]; + data[i * 1 + j + 1] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 1 + j + 0]; + m.data[1] = data[i * 1 + j + 1]; + m.data[2] = data[i * 1 + j + 2]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 1 + j + 0] = m.data[0]; + data[i * 1 + j + 1] = m.data[1]; + data[i * 1 + j + 2] = m.data[2]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix column(int j) const { + return slice_3x1(0, j); + } + + MCTLASS_HOST_DEVICE + Matrix &set_column(Matrix const &v, int j =0) { + return set_slice_3x1(v, 0, j); + } + + /// Concatenates this matrix with a a 3-by-1 matrix to form a 3-by-2 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Concatenates this matrix with a a 3-by-2 matrix to form a 3-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Concatenates this matrix with a a 3-by-3 matrix to form a 3-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Forms a 3-by-1 matrix by vertically concatenating an Element with a 2-by-1 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Element upper, Matrix const & lower) { + return Matrix( + upper + , lower.at(0, 0) + , lower.at(1, 0)); + } + + /// Forms a 3-by-1 matrix by vertically concatenating a 2-by-1 matrix with an Element + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Element lower) { + return Matrix( + upper.at(0, 0) + , upper.at(1, 0) + , lower); + } + + /// Concatenates this matrix with a an Element to form a 4-by-1 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Element rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Elementwise add operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + + result.data[1] = data[1] + rhs.data[1]; + + result.data[2] = data[2] + rhs.data[2]; + + return result; + } + + /// Elementwise add operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + + data[1] += rhs.data[1]; + + data[2] += rhs.data[2]; + + return *this; + } + + /// Elementwise subtract operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + + result.data[1] = data[1] - rhs.data[1]; + + result.data[2] = data[2] - rhs.data[2]; + + return result; + } + + /// Elementwise subtract operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + + data[1] -= rhs.data[1]; + + data[2] -= rhs.data[2]; + + return *this; + } + + /// Elementwise multiply operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + + result.data[1] = data[1] * rhs.data[1]; + + result.data[2] = data[2] * rhs.data[2]; + + return result; + } + + /// Scalar multiply operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + + result.data[1] = data[1] * s; + + result.data[2] = data[2] * s; + + return result; + } + + /// Scalar multiply operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + + data[1] *= s; + + data[2] *= s; + + return *this; + } + + /// Elementwise divide operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + + result.data[1] = data[1] / rhs.data[1]; + + result.data[2] = data[2] / rhs.data[2]; + + return result; + } + + /// Scalar divide operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + + result.data[1] = data[1] / s; + + result.data[2] = data[2] / s; + + return result; + } + + /// Scalar divide operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + + data[1] /= s; + + data[2] /= s; + + return *this; + } + + /// Elementwise divide operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (3-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + + data[1] /= rhs.data[1]; + + data[2] /= rhs.data[2]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + m.data[2] = -m.data[2]; + + return m; + } + + /// Matrix product of size 3-by-1-by-1 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[1] * rhs.data[0]; + accum.data[2] += data[2] * rhs.data[0]; + + return accum; + } + + /// Matrix product of size 3-by-1-by-1 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 3-by-1-by-1 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Matrix product of size 3-by-2-by-1 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[1] * rhs.data[0]; + accum.data[3] += data[1] * rhs.data[1]; + accum.data[4] += data[2] * rhs.data[0]; + accum.data[5] += data[2] * rhs.data[1]; + + return accum; + } + + /// Matrix product of size 3-by-2-by-1 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 3-by-3-by-1 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[1] * rhs.data[0]; + accum.data[4] += data[1] * rhs.data[1]; + accum.data[5] += data[1] * rhs.data[2]; + accum.data[6] += data[2] * rhs.data[0]; + accum.data[7] += data[2] * rhs.data[1]; + accum.data[8] += data[2] * rhs.data[2]; + + return accum; + } + + /// Matrix product of size 3-by-3-by-1 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 3-by-4-by-1 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + accum.data[4] += data[1] * rhs.data[0]; + accum.data[5] += data[1] * rhs.data[1]; + accum.data[6] += data[1] * rhs.data[2]; + accum.data[7] += data[1] * rhs.data[3]; + accum.data[8] += data[2] * rhs.data[0]; + accum.data[9] += data[2] * rhs.data[1]; + accum.data[10] += data[2] * rhs.data[2]; + accum.data[11] += data[2] * rhs.data[3]; + + return accum; + } + + /// Matrix product of size 3-by-4-by-1 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Dot product of vectors with extent 3 + MCTLASS_HOST_DEVICE + Element dot(Matrix const &rhs, Element accum = Element()) const { + + accum += data[0] * rhs.data[0]; + accum += data[1] * rhs.data[1]; + accum += data[2] * rhs.data[2]; + return accum; + } + + /// Dot product of vectors with extent 3 + MCTLASS_HOST_DEVICE + Element dot(Matrix const &rhs, Element accum = Element()) const { + + accum += data[0] * rhs.data[0]; + accum += data[1] * rhs.data[1]; + accum += data[2] * rhs.data[2]; + return accum; + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + accum += data[2]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + accum += data[2] * data[2]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + + return accum; + } + + /// Cross product + MCTLASS_HOST_DEVICE + Matrix cross(Matrix const &rhs) const { + return Matrix( + data[1] * rhs.data[2] - data[2] * rhs.data[1], + data[0] * rhs.data[2] - data[2] * rhs.data[1], + data[0] * rhs.data[1] - data[1] * rhs.data[0] + ); + } + +}; + +/// Template alias for 3-by-1 matrix +template +using Matrix3x1 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix3x1 make_Matrix3x1( + Element _0_0, + Element _1_0, + Element _2_0 +) { + return Matrix3x1( + _0_0, + _1_0, + _2_0 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 3-by-2 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 3; + + /// Number of columns in matrix + static int const kColumns = 2; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 6; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 3-by-2 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 3-by-2 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, Element _0_1, + Element _1_0, Element _1_1, + Element _2_0, Element _2_1 + ) { + + data[0] = _0_0; data[1] = _0_1; + data[2] = _1_0; data[3] = _1_1; + data[4] = _2_0; data[5] = _2_1; + } + + /// Constucts a 3-by-2 matrix from row vectors + MCTLASS_HOST_DEVICE + Matrix( + Matrix const &row_0, + Matrix const &row_1, + Matrix const &row_2 + ) { + data[0] = row_0.data[0]; + data[1] = row_0.data[1]; + data[2] = row_1.data[0]; + data[3] = row_1.data[1]; + data[4] = row_2.data[0]; + data[5] = row_2.data[1]; + } + + /// Static method to construct a 3-by-2 matrix from column vectors + MCTLASS_HOST_DEVICE + static Matrix from_columns( + Matrix const &column_0, + Matrix const &column_1 + ) { + Matrix result; + + result.data[0] = column_0.data[0]; + result.data[1] = column_1.data[0]; + result.data[2] = column_0.data[1]; + result.data[3] = column_1.data[1]; + result.data[4] = column_0.data[2]; + result.data[5] = column_1.data[2]; + return result; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + m.data[2] = s; + m.data[3] = s; + m.data[4] = s; + m.data[5] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[4] = diag.data[1]; + m.data[8] = diag.data[2]; + + return m; + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[4] = diag.data[1]; + m.data[8] = diag.data[2]; + + return m; + } + + /// Gets an array of diagonal elements + MCTLASS_HOST_DEVICE + Matrix diagonal() const { + Matrix diag; + + diag.data[0] = data[0]; + diag.data[1] = data[4]; + diag.data[2] = data[8]; + + return diag; + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[3] = data[1]; + mt.data[1] = data[2]; + mt.data[4] = data[3]; + mt.data[2] = data[4]; + mt.data[5] = data[5]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 3 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 3 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 1] = m.data[1]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix row(int i) const { + return slice_1x2(i, 0); + } + + MCTLASS_HOST_DEVICE + Matrix &set_row(Matrix const &v, int i = 0) { + return set_slice_1x2(v, i, 0); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 2]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 2] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 1]; + m.data[2] = data[i * 2 + j + 2]; + m.data[3] = data[i * 2 + j + 3]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 1] = m.data[1]; + data[i * 2 + j + 2] = m.data[2]; + data[i * 2 + j + 3] = m.data[3]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 2]; + m.data[2] = data[i * 2 + j + 4]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 2] = m.data[1]; + data[i * 2 + j + 4] = m.data[2]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix column(int j) const { + return slice_3x1(0, j); + } + + MCTLASS_HOST_DEVICE + Matrix &set_column(Matrix const &v, int j =0) { + return set_slice_3x1(v, 0, j); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 1]; + m.data[2] = data[i * 2 + j + 2]; + m.data[3] = data[i * 2 + j + 3]; + m.data[4] = data[i * 2 + j + 4]; + m.data[5] = data[i * 2 + j + 5]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 1] = m.data[1]; + data[i * 2 + j + 2] = m.data[2]; + data[i * 2 + j + 3] = m.data[3]; + data[i * 2 + j + 4] = m.data[4]; + data[i * 2 + j + 5] = m.data[5]; + + return *this; + } + + /// Forms a 3-by-2 matrix by horizontally concatenating a 3-by-1 matrix with a 3-by-1 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), rhs.at(0, 0) + , lhs.at(1, 0), rhs.at(1, 0) + , lhs.at(2, 0), rhs.at(2, 0)); + } + + /// Concatenates this matrix with a a 3-by-1 matrix to form a 3-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Concatenates this matrix with a a 3-by-2 matrix to form a 3-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Forms a 3-by-2 matrix by vertically concatenating a 1-by-2 matrix with a 2-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1) + , lower.at(0, 0), lower.at(0, 1) + , lower.at(1, 0), lower.at(1, 1)); + } + + /// Forms a 3-by-2 matrix by vertically concatenating a 2-by-2 matrix with a 1-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1) + , upper.at(1, 0), upper.at(1, 1) + , lower.at(0, 0), lower.at(0, 1)); + } + + /// Concatenates this matrix with a a 1-by-2 matrix to form a 4-by-2 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Forms a 3-by-2 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Element A, Element B, + Matrix const & C, Matrix const & D) { + return Matrix( + A, B + , C.at(0, 0), D.at(0, 0) + , C.at(1, 0), D.at(1, 0) + ); + } + + /// Forms a 3-by-2 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Element C, Element D) { + return Matrix( + A.at(0, 0), B.at(0, 0) + , A.at(1, 0), B.at(1, 0) + , C, D + ); + } + + /// Elementwise add operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + result.data[1] = data[1] + rhs.data[1]; + + result.data[2] = data[2] + rhs.data[2]; + result.data[3] = data[3] + rhs.data[3]; + + result.data[4] = data[4] + rhs.data[4]; + result.data[5] = data[5] + rhs.data[5]; + + return result; + } + + /// Elementwise add operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + data[1] += rhs.data[1]; + + data[2] += rhs.data[2]; + data[3] += rhs.data[3]; + + data[4] += rhs.data[4]; + data[5] += rhs.data[5]; + + return *this; + } + + /// Elementwise subtract operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + result.data[1] = data[1] - rhs.data[1]; + + result.data[2] = data[2] - rhs.data[2]; + result.data[3] = data[3] - rhs.data[3]; + + result.data[4] = data[4] - rhs.data[4]; + result.data[5] = data[5] - rhs.data[5]; + + return result; + } + + /// Elementwise subtract operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + data[1] -= rhs.data[1]; + + data[2] -= rhs.data[2]; + data[3] -= rhs.data[3]; + + data[4] -= rhs.data[4]; + data[5] -= rhs.data[5]; + + return *this; + } + + /// Elementwise multiply operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + result.data[1] = data[1] * rhs.data[1]; + + result.data[2] = data[2] * rhs.data[2]; + result.data[3] = data[3] * rhs.data[3]; + + result.data[4] = data[4] * rhs.data[4]; + result.data[5] = data[5] * rhs.data[5]; + + return result; + } + + /// Scalar multiply operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + result.data[1] = data[1] * s; + + result.data[2] = data[2] * s; + result.data[3] = data[3] * s; + + result.data[4] = data[4] * s; + result.data[5] = data[5] * s; + + return result; + } + + /// Scalar multiply operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + data[1] *= s; + + data[2] *= s; + data[3] *= s; + + data[4] *= s; + data[5] *= s; + + return *this; + } + + /// Elementwise divide operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + result.data[1] = data[1] / rhs.data[1]; + + result.data[2] = data[2] / rhs.data[2]; + result.data[3] = data[3] / rhs.data[3]; + + result.data[4] = data[4] / rhs.data[4]; + result.data[5] = data[5] / rhs.data[5]; + + return result; + } + + /// Scalar divide operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + result.data[1] = data[1] / s; + + result.data[2] = data[2] / s; + result.data[3] = data[3] / s; + + result.data[4] = data[4] / s; + result.data[5] = data[5] / s; + + return result; + } + + /// Scalar divide operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + data[1] /= s; + + data[2] /= s; + data[3] /= s; + + data[4] /= s; + data[5] /= s; + + return *this; + } + + /// Elementwise divide operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (3-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + data[1] /= rhs.data[1]; + + data[2] /= rhs.data[2]; + data[3] /= rhs.data[3]; + + data[4] /= rhs.data[4]; + data[5] /= rhs.data[5]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + m.data[2] = -m.data[2]; + m.data[3] = -m.data[3]; + m.data[4] = -m.data[4]; + m.data[5] = -m.data[5]; + + return m; + } + + /// Matrix product of size 3-by-1-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[2] * rhs.data[0]; + accum.data[2] += data[4] * rhs.data[0]; + + // k=1 + accum.data[0] += data[1] * rhs.data[1]; + accum.data[1] += data[3] * rhs.data[1]; + accum.data[2] += data[5] * rhs.data[1]; + + return accum; + } + + /// Matrix product of size 3-by-1-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 3-by-2-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[2] * rhs.data[0]; + accum.data[3] += data[2] * rhs.data[1]; + accum.data[4] += data[4] * rhs.data[0]; + accum.data[5] += data[4] * rhs.data[1]; + + // k=1 + accum.data[0] += data[1] * rhs.data[2]; + accum.data[1] += data[1] * rhs.data[3]; + accum.data[2] += data[3] * rhs.data[2]; + accum.data[3] += data[3] * rhs.data[3]; + accum.data[4] += data[5] * rhs.data[2]; + accum.data[5] += data[5] * rhs.data[3]; + + return accum; + } + + /// Matrix product of size 3-by-2-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 3-by-2-by-2 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Matrix product of size 3-by-3-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[2] * rhs.data[0]; + accum.data[4] += data[2] * rhs.data[1]; + accum.data[5] += data[2] * rhs.data[2]; + accum.data[6] += data[4] * rhs.data[0]; + accum.data[7] += data[4] * rhs.data[1]; + accum.data[8] += data[4] * rhs.data[2]; + + // k=1 + accum.data[0] += data[1] * rhs.data[3]; + accum.data[1] += data[1] * rhs.data[4]; + accum.data[2] += data[1] * rhs.data[5]; + accum.data[3] += data[3] * rhs.data[3]; + accum.data[4] += data[3] * rhs.data[4]; + accum.data[5] += data[3] * rhs.data[5]; + accum.data[6] += data[5] * rhs.data[3]; + accum.data[7] += data[5] * rhs.data[4]; + accum.data[8] += data[5] * rhs.data[5]; + + return accum; + } + + /// Matrix product of size 3-by-3-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 3-by-4-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + accum.data[4] += data[2] * rhs.data[0]; + accum.data[5] += data[2] * rhs.data[1]; + accum.data[6] += data[2] * rhs.data[2]; + accum.data[7] += data[2] * rhs.data[3]; + accum.data[8] += data[4] * rhs.data[0]; + accum.data[9] += data[4] * rhs.data[1]; + accum.data[10] += data[4] * rhs.data[2]; + accum.data[11] += data[4] * rhs.data[3]; + + // k=1 + accum.data[0] += data[1] * rhs.data[4]; + accum.data[1] += data[1] * rhs.data[5]; + accum.data[2] += data[1] * rhs.data[6]; + accum.data[3] += data[1] * rhs.data[7]; + accum.data[4] += data[3] * rhs.data[4]; + accum.data[5] += data[3] * rhs.data[5]; + accum.data[6] += data[3] * rhs.data[6]; + accum.data[7] += data[3] * rhs.data[7]; + accum.data[8] += data[5] * rhs.data[4]; + accum.data[9] += data[5] * rhs.data[5]; + accum.data[10] += data[5] * rhs.data[6]; + accum.data[11] += data[5] * rhs.data[7]; + + return accum; + } + + /// Matrix product of size 3-by-4-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + accum += data[2]; + accum += data[3]; + accum += data[4]; + accum += data[5]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + accum += data[2] * data[2]; + accum += data[3] * data[3]; + accum += data[4] * data[4]; + accum += data[5] * data[5]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + accum += data[3]; + + return accum; + } + +}; + +/// Template alias for 3-by-2 matrix +template +using Matrix3x2 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix3x2 make_Matrix3x2( + Element _0_0, Element _0_1, + Element _1_0, Element _1_1, + Element _2_0, Element _2_1 +) { + return Matrix3x2( + _0_0, _0_1, + _1_0, _1_1, + _2_0, _2_1 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 3-by-3 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 3; + + /// Number of columns in matrix + static int const kColumns = 3; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 9; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 3-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 3-by-3 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, Element _0_1, Element _0_2, + Element _1_0, Element _1_1, Element _1_2, + Element _2_0, Element _2_1, Element _2_2 + ) { + + data[0] = _0_0; data[1] = _0_1; data[2] = _0_2; + data[3] = _1_0; data[4] = _1_1; data[5] = _1_2; + data[6] = _2_0; data[7] = _2_1; data[8] = _2_2; + } + + /// Constucts a 3-by-3 matrix from row vectors + MCTLASS_HOST_DEVICE + Matrix( + Matrix const &row_0, + Matrix const &row_1, + Matrix const &row_2 + ) { + data[0] = row_0.data[0]; + data[1] = row_0.data[1]; + data[2] = row_0.data[2]; + data[3] = row_1.data[0]; + data[4] = row_1.data[1]; + data[5] = row_1.data[2]; + data[6] = row_2.data[0]; + data[7] = row_2.data[1]; + data[8] = row_2.data[2]; + } + + /// Static method to construct a 3-by-3 matrix from column vectors + MCTLASS_HOST_DEVICE + static Matrix from_columns( + Matrix const &column_0, + Matrix const &column_1, + Matrix const &column_2 + ) { + Matrix result; + + result.data[0] = column_0.data[0]; + result.data[1] = column_1.data[0]; + result.data[2] = column_2.data[0]; + result.data[3] = column_0.data[1]; + result.data[4] = column_1.data[1]; + result.data[5] = column_2.data[1]; + result.data[6] = column_0.data[2]; + result.data[7] = column_1.data[2]; + result.data[8] = column_2.data[2]; + return result; + } + + /// Constructs an identity matrix + MCTLASS_HOST_DEVICE + static Matrix identity() { + Matrix m; + + m.data[0] = Element(1); + m.data[4] = Element(1); + m.data[8] = Element(1); + + return m; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + m.data[2] = s; + m.data[3] = s; + m.data[4] = s; + m.data[5] = s; + m.data[6] = s; + m.data[7] = s; + m.data[8] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[4] = diag.data[1]; + m.data[8] = diag.data[2]; + + return m; + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[4] = diag.data[1]; + m.data[8] = diag.data[2]; + + return m; + } + + /// Gets an array of diagonal elements + MCTLASS_HOST_DEVICE + Matrix diagonal() const { + Matrix diag; + + diag.data[0] = data[0]; + diag.data[1] = data[4]; + diag.data[2] = data[8]; + + return diag; + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[3] = data[1]; + mt.data[6] = data[2]; + mt.data[1] = data[3]; + mt.data[4] = data[4]; + mt.data[7] = data[5]; + mt.data[2] = data[6]; + mt.data[5] = data[7]; + mt.data[8] = data[8]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 3 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 3 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 2]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 2] = m.data[2]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix row(int i) const { + return slice_1x3(i, 0); + } + + MCTLASS_HOST_DEVICE + Matrix &set_row(Matrix const &v, int i = 0) { + return set_slice_1x3(v, i, 0); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 3]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 3] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 3]; + m.data[3] = data[i * 3 + j + 4]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 3] = m.data[2]; + data[i * 3 + j + 4] = m.data[3]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 2]; + m.data[3] = data[i * 3 + j + 3]; + m.data[4] = data[i * 3 + j + 4]; + m.data[5] = data[i * 3 + j + 5]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 2] = m.data[2]; + data[i * 3 + j + 3] = m.data[3]; + data[i * 3 + j + 4] = m.data[4]; + data[i * 3 + j + 5] = m.data[5]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 3]; + m.data[2] = data[i * 3 + j + 6]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 3] = m.data[1]; + data[i * 3 + j + 6] = m.data[2]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix column(int j) const { + return slice_3x1(0, j); + } + + MCTLASS_HOST_DEVICE + Matrix &set_column(Matrix const &v, int j =0) { + return set_slice_3x1(v, 0, j); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 3]; + m.data[3] = data[i * 3 + j + 4]; + m.data[4] = data[i * 3 + j + 6]; + m.data[5] = data[i * 3 + j + 7]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 3] = m.data[2]; + data[i * 3 + j + 4] = m.data[3]; + data[i * 3 + j + 6] = m.data[4]; + data[i * 3 + j + 7] = m.data[5]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 2]; + m.data[3] = data[i * 3 + j + 3]; + m.data[4] = data[i * 3 + j + 4]; + m.data[5] = data[i * 3 + j + 5]; + m.data[6] = data[i * 3 + j + 6]; + m.data[7] = data[i * 3 + j + 7]; + m.data[8] = data[i * 3 + j + 8]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 2] = m.data[2]; + data[i * 3 + j + 3] = m.data[3]; + data[i * 3 + j + 4] = m.data[4]; + data[i * 3 + j + 5] = m.data[5]; + data[i * 3 + j + 6] = m.data[6]; + data[i * 3 + j + 7] = m.data[7]; + data[i * 3 + j + 8] = m.data[8]; + + return *this; + } + + /// Forms a 3-by-3 matrix by horizontally concatenating a 3-by-1 matrix with a 3-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), rhs.at(0, 0), rhs.at(0, 1) + , lhs.at(1, 0), rhs.at(1, 0), rhs.at(1, 1) + , lhs.at(2, 0), rhs.at(2, 0), rhs.at(2, 1)); + } + + /// Forms a 3-by-3 matrix by horizontally concatenating a 3-by-2 matrix with a 3-by-1 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0) + , lhs.at(1, 0), lhs.at(1, 1), rhs.at(1, 0) + , lhs.at(2, 0), lhs.at(2, 1), rhs.at(2, 0)); + } + + /// Concatenates this matrix with a a 3-by-1 matrix to form a 3-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Forms a 3-by-3 matrix by vertically concatenating a 1-by-3 matrix with a 2-by-3 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1), upper.at(0, 2) + , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2) + , lower.at(1, 0), lower.at(1, 1), lower.at(1, 2)); + } + + /// Forms a 3-by-3 matrix by vertically concatenating a 2-by-3 matrix with a 1-by-3 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1), upper.at(0, 2) + , upper.at(1, 0), upper.at(1, 1), upper.at(1, 2) + , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2)); + } + + /// Concatenates this matrix with a a 1-by-3 matrix to form a 4-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Forms a 3-by-3 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Element A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A, B.at(0, 0), B.at(0, 1) + , C.at(0, 0), D.at(0, 0), D.at(0, 1) + , C.at(1, 0), D.at(1, 0), D.at(1, 1) + ); + } + + /// Forms a 3-by-3 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Element B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), A.at(0, 1), B + , C.at(0, 0), C.at(0, 1), D.at(0, 0) + , C.at(1, 0), C.at(1, 1), D.at(1, 0) + ); + } + + /// Forms a 3-by-3 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Element C, Matrix const & D) { + return Matrix( + A.at(0, 0), B.at(0, 0), B.at(0, 1) + , A.at(1, 0), B.at(1, 0), B.at(1, 1) + , C, D.at(0, 0), D.at(0, 1) + ); + } + + /// Forms a 3-by-3 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Element D) { + return Matrix( + A.at(0, 0), A.at(0, 1), B.at(0, 0) + , A.at(1, 0), A.at(1, 1), B.at(1, 0) + , C.at(0, 0), C.at(0, 1), D + ); + } + + /// Elementwise add operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + result.data[1] = data[1] + rhs.data[1]; + result.data[2] = data[2] + rhs.data[2]; + + result.data[3] = data[3] + rhs.data[3]; + result.data[4] = data[4] + rhs.data[4]; + result.data[5] = data[5] + rhs.data[5]; + + result.data[6] = data[6] + rhs.data[6]; + result.data[7] = data[7] + rhs.data[7]; + result.data[8] = data[8] + rhs.data[8]; + + return result; + } + + /// Elementwise add operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + data[1] += rhs.data[1]; + data[2] += rhs.data[2]; + + data[3] += rhs.data[3]; + data[4] += rhs.data[4]; + data[5] += rhs.data[5]; + + data[6] += rhs.data[6]; + data[7] += rhs.data[7]; + data[8] += rhs.data[8]; + + return *this; + } + + /// Elementwise subtract operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + result.data[1] = data[1] - rhs.data[1]; + result.data[2] = data[2] - rhs.data[2]; + + result.data[3] = data[3] - rhs.data[3]; + result.data[4] = data[4] - rhs.data[4]; + result.data[5] = data[5] - rhs.data[5]; + + result.data[6] = data[6] - rhs.data[6]; + result.data[7] = data[7] - rhs.data[7]; + result.data[8] = data[8] - rhs.data[8]; + + return result; + } + + /// Elementwise subtract operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + data[1] -= rhs.data[1]; + data[2] -= rhs.data[2]; + + data[3] -= rhs.data[3]; + data[4] -= rhs.data[4]; + data[5] -= rhs.data[5]; + + data[6] -= rhs.data[6]; + data[7] -= rhs.data[7]; + data[8] -= rhs.data[8]; + + return *this; + } + + /// Elementwise multiply operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + result.data[1] = data[1] * rhs.data[1]; + result.data[2] = data[2] * rhs.data[2]; + + result.data[3] = data[3] * rhs.data[3]; + result.data[4] = data[4] * rhs.data[4]; + result.data[5] = data[5] * rhs.data[5]; + + result.data[6] = data[6] * rhs.data[6]; + result.data[7] = data[7] * rhs.data[7]; + result.data[8] = data[8] * rhs.data[8]; + + return result; + } + + /// Scalar multiply operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + result.data[1] = data[1] * s; + result.data[2] = data[2] * s; + + result.data[3] = data[3] * s; + result.data[4] = data[4] * s; + result.data[5] = data[5] * s; + + result.data[6] = data[6] * s; + result.data[7] = data[7] * s; + result.data[8] = data[8] * s; + + return result; + } + + /// Scalar multiply operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + data[1] *= s; + data[2] *= s; + + data[3] *= s; + data[4] *= s; + data[5] *= s; + + data[6] *= s; + data[7] *= s; + data[8] *= s; + + return *this; + } + + /// Elementwise divide operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + result.data[1] = data[1] / rhs.data[1]; + result.data[2] = data[2] / rhs.data[2]; + + result.data[3] = data[3] / rhs.data[3]; + result.data[4] = data[4] / rhs.data[4]; + result.data[5] = data[5] / rhs.data[5]; + + result.data[6] = data[6] / rhs.data[6]; + result.data[7] = data[7] / rhs.data[7]; + result.data[8] = data[8] / rhs.data[8]; + + return result; + } + + /// Scalar divide operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + result.data[1] = data[1] / s; + result.data[2] = data[2] / s; + + result.data[3] = data[3] / s; + result.data[4] = data[4] / s; + result.data[5] = data[5] / s; + + result.data[6] = data[6] / s; + result.data[7] = data[7] / s; + result.data[8] = data[8] / s; + + return result; + } + + /// Scalar divide operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + data[1] /= s; + data[2] /= s; + + data[3] /= s; + data[4] /= s; + data[5] /= s; + + data[6] /= s; + data[7] /= s; + data[8] /= s; + + return *this; + } + + /// Elementwise divide operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (3-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + data[1] /= rhs.data[1]; + data[2] /= rhs.data[2]; + + data[3] /= rhs.data[3]; + data[4] /= rhs.data[4]; + data[5] /= rhs.data[5]; + + data[6] /= rhs.data[6]; + data[7] /= rhs.data[7]; + data[8] /= rhs.data[8]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + m.data[2] = -m.data[2]; + m.data[3] = -m.data[3]; + m.data[4] = -m.data[4]; + m.data[5] = -m.data[5]; + m.data[6] = -m.data[6]; + m.data[7] = -m.data[7]; + m.data[8] = -m.data[8]; + + return m; + } + + /// Matrix product of size 3-by-1-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[3] * rhs.data[0]; + accum.data[2] += data[6] * rhs.data[0]; + + // k=1 + accum.data[0] += data[1] * rhs.data[1]; + accum.data[1] += data[4] * rhs.data[1]; + accum.data[2] += data[7] * rhs.data[1]; + + // k=2 + accum.data[0] += data[2] * rhs.data[2]; + accum.data[1] += data[5] * rhs.data[2]; + accum.data[2] += data[8] * rhs.data[2]; + + return accum; + } + + /// Matrix product of size 3-by-1-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 3-by-2-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[3] * rhs.data[0]; + accum.data[3] += data[3] * rhs.data[1]; + accum.data[4] += data[6] * rhs.data[0]; + accum.data[5] += data[6] * rhs.data[1]; + + // k=1 + accum.data[0] += data[1] * rhs.data[2]; + accum.data[1] += data[1] * rhs.data[3]; + accum.data[2] += data[4] * rhs.data[2]; + accum.data[3] += data[4] * rhs.data[3]; + accum.data[4] += data[7] * rhs.data[2]; + accum.data[5] += data[7] * rhs.data[3]; + + // k=2 + accum.data[0] += data[2] * rhs.data[4]; + accum.data[1] += data[2] * rhs.data[5]; + accum.data[2] += data[5] * rhs.data[4]; + accum.data[3] += data[5] * rhs.data[5]; + accum.data[4] += data[8] * rhs.data[4]; + accum.data[5] += data[8] * rhs.data[5]; + + return accum; + } + + /// Matrix product of size 3-by-2-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 3-by-3-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[3] * rhs.data[0]; + accum.data[4] += data[3] * rhs.data[1]; + accum.data[5] += data[3] * rhs.data[2]; + accum.data[6] += data[6] * rhs.data[0]; + accum.data[7] += data[6] * rhs.data[1]; + accum.data[8] += data[6] * rhs.data[2]; + + // k=1 + accum.data[0] += data[1] * rhs.data[3]; + accum.data[1] += data[1] * rhs.data[4]; + accum.data[2] += data[1] * rhs.data[5]; + accum.data[3] += data[4] * rhs.data[3]; + accum.data[4] += data[4] * rhs.data[4]; + accum.data[5] += data[4] * rhs.data[5]; + accum.data[6] += data[7] * rhs.data[3]; + accum.data[7] += data[7] * rhs.data[4]; + accum.data[8] += data[7] * rhs.data[5]; + + // k=2 + accum.data[0] += data[2] * rhs.data[6]; + accum.data[1] += data[2] * rhs.data[7]; + accum.data[2] += data[2] * rhs.data[8]; + accum.data[3] += data[5] * rhs.data[6]; + accum.data[4] += data[5] * rhs.data[7]; + accum.data[5] += data[5] * rhs.data[8]; + accum.data[6] += data[8] * rhs.data[6]; + accum.data[7] += data[8] * rhs.data[7]; + accum.data[8] += data[8] * rhs.data[8]; + + return accum; + } + + /// Matrix product of size 3-by-3-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 3-by-3-by-3 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Matrix product of size 3-by-4-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + accum.data[4] += data[3] * rhs.data[0]; + accum.data[5] += data[3] * rhs.data[1]; + accum.data[6] += data[3] * rhs.data[2]; + accum.data[7] += data[3] * rhs.data[3]; + accum.data[8] += data[6] * rhs.data[0]; + accum.data[9] += data[6] * rhs.data[1]; + accum.data[10] += data[6] * rhs.data[2]; + accum.data[11] += data[6] * rhs.data[3]; + + // k=1 + accum.data[0] += data[1] * rhs.data[4]; + accum.data[1] += data[1] * rhs.data[5]; + accum.data[2] += data[1] * rhs.data[6]; + accum.data[3] += data[1] * rhs.data[7]; + accum.data[4] += data[4] * rhs.data[4]; + accum.data[5] += data[4] * rhs.data[5]; + accum.data[6] += data[4] * rhs.data[6]; + accum.data[7] += data[4] * rhs.data[7]; + accum.data[8] += data[7] * rhs.data[4]; + accum.data[9] += data[7] * rhs.data[5]; + accum.data[10] += data[7] * rhs.data[6]; + accum.data[11] += data[7] * rhs.data[7]; + + // k=2 + accum.data[0] += data[2] * rhs.data[8]; + accum.data[1] += data[2] * rhs.data[9]; + accum.data[2] += data[2] * rhs.data[10]; + accum.data[3] += data[2] * rhs.data[11]; + accum.data[4] += data[5] * rhs.data[8]; + accum.data[5] += data[5] * rhs.data[9]; + accum.data[6] += data[5] * rhs.data[10]; + accum.data[7] += data[5] * rhs.data[11]; + accum.data[8] += data[8] * rhs.data[8]; + accum.data[9] += data[8] * rhs.data[9]; + accum.data[10] += data[8] * rhs.data[10]; + accum.data[11] += data[8] * rhs.data[11]; + + return accum; + } + + /// Matrix product of size 3-by-4-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + accum += data[2]; + accum += data[3]; + accum += data[4]; + accum += data[5]; + accum += data[6]; + accum += data[7]; + accum += data[8]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + accum += data[2] * data[2]; + accum += data[3] * data[3]; + accum += data[4] * data[4]; + accum += data[5] * data[5]; + accum += data[6] * data[6]; + accum += data[7] * data[7]; + accum += data[8] * data[8]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + accum += data[4]; + accum += data[8]; + + return accum; + } + + /// Returns 3-by-3 rotation matrix around the X axis + MCTLASS_HOST_DEVICE + static Matrix rotation_X(Element theta) { + Matrix m = identity(); + + Element c = fast_cos(theta); + Element s = fast_sin(theta); + + m.at(1, 1) = c; + m.at(1, 2) = -s; + m.at(2, 1) = s; + m.at(2, 2) = c; + + return m; + } + + /// Returns 3-by-3 rotation matrix around the Y axis + MCTLASS_HOST_DEVICE + static Matrix rotation_Y(Element theta) { + Matrix m = identity(); + + Element c = fast_cos(theta); + Element s = fast_sin(theta); + + m.at(0, 0) = c; + m.at(2, 0) = -s; + m.at(0, 2) = s; + m.at(2, 2) = c; + + return m; + } + + /// Returns 3-by-3 rotation matrix around the Z axis + MCTLASS_HOST_DEVICE + static Matrix rotation_Z(Element theta) { + Matrix m = Matrix::identity(); + + Element c = fast_cos(theta); + Element s = fast_sin(theta); + + m.at(0, 0) = c; + m.at(0, 1) = -s; + m.at(1, 0) = s; + m.at(1, 1) = c; + + return m; + } + + /// Returns a 3-by-3 rotation matrix around a unit-length axis + MCTLASS_HOST_DEVICE + static Matrix rotation(Element theta, Matrix const &u) { + Element x = u.data[0]; + Element y = u.data[1]; + Element z = u.data[2]; + + Element c = fast_cos(theta); + Element s = fast_sin(theta); + + Element one_minus_cos = Element(1) - fast_cos(theta); + + Matrix m; + + m.set_slice_3x3({ + c + x * x * one_minus_cos, x * y * one_minus_cos - z * s, x * z * one_minus_cos + y * s, + y * x * one_minus_cos * z * s, c + y * y * one_minus_cos, y * z * one_minus_cos - x * s, + z * x * one_minus_cos - y * s, z * y * one_minus_cos + x * s, c + z * z * one_minus_cos + }); + + return m; + } + + /// Returns a 3-by-3 reflection about the plane specified by the + /// unit-length normal vector n_unit + MCTLASS_HOST_DEVICE + static Matrix reflection(Matrix const &n_unit) { + + Element a = n_unit.data[0]; + Element b = n_unit.data[1]; + Element c = n_unit.data[2]; + + Matrix m = Matrix::identity(); + + m.set_slice_3x3({ + Element(1) - Element(2) * a * a, Element(-2) * a * b, Element(-2) * a * c, + Element(-2) * a * b, Element(1) - Element(2) * b * b, Element(-2) * b * c, + Element(-2) * a * c, Element(-2) * b * c, Element(1) - Element(2) * c * c + }); + + return m; + } + + /// Computes the determinant of a 3-by-3 matrix + MCTLASS_HOST_DEVICE + Element determinant(Element accum = Element()) const { + + accum += at(0, 0) * Matrix({ at(1, 1), at(1, 2), at(2, 1), at(2, 2) }).determinant(); + accum -= at(0, 1) * Matrix({ at(1, 0), at(1, 2), at(2, 0), at(2, 2) }).determinant(); + accum += at(0, 2) * Matrix({ at(1, 0), at(1, 1), at(2, 0), at(2, 1) }).determinant(); + + return accum; + } + + /// Computes the inverse of a 3-by-3 matrix given + /// the matrix's determinant + MCTLASS_HOST_DEVICE + Matrix inverse(Element det) const { + return Matrix( + at(1, 1) * at(2, 2) - at(1, 2) * at(2, 1), + at(0, 2) * at(2, 1) - at(0, 1) * at(2, 2), + at(0, 1) * at(1, 2) - at(0, 2) * at(1, 1), + + at(1, 2) * at(2, 0) - at(1, 0) * at(2, 2), + at(0, 0) * at(2, 2) - at(0, 2) * at(2, 0), + at(0, 2) * at(1, 0) - at(0, 0) * at(1, 2), + + at(1, 0) * at(2, 1) - at(1, 1) * at(2, 0), + at(0, 1) * at(2, 0) - at(0, 0) * at(2, 1), + at(0, 0) * at(1, 1) - at(0, 1) * at(1, 0) + ) * (Element(1) / det); + } + /// Computes the inverse of a 3-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix inverse() const { + return inverse(determinant()); + } + +}; + +/// Template alias for 3-by-3 matrix +template +using Matrix3x3 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix3x3 make_Matrix3x3( + Element _0_0, Element _0_1, Element _0_2, + Element _1_0, Element _1_1, Element _1_2, + Element _2_0, Element _2_1, Element _2_2 +) { + return Matrix3x3( + _0_0, _0_1, _0_2, + _1_0, _1_1, _1_2, + _2_0, _2_1, _2_2 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 3-by-4 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 3; + + /// Number of columns in matrix + static int const kColumns = 4; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 12; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 3-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 3-by-4 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, Element _0_1, Element _0_2, Element _0_3, + Element _1_0, Element _1_1, Element _1_2, Element _1_3, + Element _2_0, Element _2_1, Element _2_2, Element _2_3 + ) { + + data[0] = _0_0; data[1] = _0_1; data[2] = _0_2; data[3] = _0_3; + data[4] = _1_0; data[5] = _1_1; data[6] = _1_2; data[7] = _1_3; + data[8] = _2_0; data[9] = _2_1; data[10] = _2_2; data[11] = _2_3; + } + + /// Constucts a 3-by-4 matrix from row vectors + MCTLASS_HOST_DEVICE + Matrix( + Matrix const &row_0, + Matrix const &row_1, + Matrix const &row_2 + ) { + data[0] = row_0.data[0]; + data[1] = row_0.data[1]; + data[2] = row_0.data[2]; + data[3] = row_0.data[3]; + data[4] = row_1.data[0]; + data[5] = row_1.data[1]; + data[6] = row_1.data[2]; + data[7] = row_1.data[3]; + data[8] = row_2.data[0]; + data[9] = row_2.data[1]; + data[10] = row_2.data[2]; + data[11] = row_2.data[3]; + } + + /// Static method to construct a 3-by-4 matrix from column vectors + MCTLASS_HOST_DEVICE + static Matrix from_columns( + Matrix const &column_0, + Matrix const &column_1, + Matrix const &column_2, + Matrix const &column_3 + ) { + Matrix result; + + result.data[0] = column_0.data[0]; + result.data[1] = column_1.data[0]; + result.data[2] = column_2.data[0]; + result.data[3] = column_3.data[0]; + result.data[4] = column_0.data[1]; + result.data[5] = column_1.data[1]; + result.data[6] = column_2.data[1]; + result.data[7] = column_3.data[1]; + result.data[8] = column_0.data[2]; + result.data[9] = column_1.data[2]; + result.data[10] = column_2.data[2]; + result.data[11] = column_3.data[2]; + return result; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + m.data[2] = s; + m.data[3] = s; + m.data[4] = s; + m.data[5] = s; + m.data[6] = s; + m.data[7] = s; + m.data[8] = s; + m.data[9] = s; + m.data[10] = s; + m.data[11] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[4] = diag.data[1]; + m.data[8] = diag.data[2]; + + return m; + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[4] = diag.data[1]; + m.data[8] = diag.data[2]; + + return m; + } + + /// Gets an array of diagonal elements + MCTLASS_HOST_DEVICE + Matrix diagonal() const { + Matrix diag; + + diag.data[0] = data[0]; + diag.data[1] = data[4]; + diag.data[2] = data[8]; + + return diag; + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[3] = data[1]; + mt.data[6] = data[2]; + mt.data[9] = data[3]; + mt.data[1] = data[4]; + mt.data[4] = data[5]; + mt.data[7] = data[6]; + mt.data[10] = data[7]; + mt.data[2] = data[8]; + mt.data[5] = data[9]; + mt.data[8] = data[10]; + mt.data[11] = data[11]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 3 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 3 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x4(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 3]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x4(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 3] = m.data[3]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix row(int i) const { + return slice_1x4(i, 0); + } + + MCTLASS_HOST_DEVICE + Matrix &set_row(Matrix const &v, int i = 0) { + return set_slice_1x4(v, i, 0); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 4]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 4] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 4]; + m.data[3] = data[i * 4 + j + 5]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 4] = m.data[2]; + data[i * 4 + j + 5] = m.data[3]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 4]; + m.data[4] = data[i * 4 + j + 5]; + m.data[5] = data[i * 4 + j + 6]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 4] = m.data[3]; + data[i * 4 + j + 5] = m.data[4]; + data[i * 4 + j + 6] = m.data[5]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x4(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 3]; + m.data[4] = data[i * 4 + j + 4]; + m.data[5] = data[i * 4 + j + 5]; + m.data[6] = data[i * 4 + j + 6]; + m.data[7] = data[i * 4 + j + 7]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x4(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 3] = m.data[3]; + data[i * 4 + j + 4] = m.data[4]; + data[i * 4 + j + 5] = m.data[5]; + data[i * 4 + j + 6] = m.data[6]; + data[i * 4 + j + 7] = m.data[7]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 4]; + m.data[2] = data[i * 4 + j + 8]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 4] = m.data[1]; + data[i * 4 + j + 8] = m.data[2]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix column(int j) const { + return slice_3x1(0, j); + } + + MCTLASS_HOST_DEVICE + Matrix &set_column(Matrix const &v, int j =0) { + return set_slice_3x1(v, 0, j); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 4]; + m.data[3] = data[i * 4 + j + 5]; + m.data[4] = data[i * 4 + j + 8]; + m.data[5] = data[i * 4 + j + 9]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 4] = m.data[2]; + data[i * 4 + j + 5] = m.data[3]; + data[i * 4 + j + 8] = m.data[4]; + data[i * 4 + j + 9] = m.data[5]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 4]; + m.data[4] = data[i * 4 + j + 5]; + m.data[5] = data[i * 4 + j + 6]; + m.data[6] = data[i * 4 + j + 8]; + m.data[7] = data[i * 4 + j + 9]; + m.data[8] = data[i * 4 + j + 10]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 4] = m.data[3]; + data[i * 4 + j + 5] = m.data[4]; + data[i * 4 + j + 6] = m.data[5]; + data[i * 4 + j + 8] = m.data[6]; + data[i * 4 + j + 9] = m.data[7]; + data[i * 4 + j + 10] = m.data[8]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x4(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 3]; + m.data[4] = data[i * 4 + j + 4]; + m.data[5] = data[i * 4 + j + 5]; + m.data[6] = data[i * 4 + j + 6]; + m.data[7] = data[i * 4 + j + 7]; + m.data[8] = data[i * 4 + j + 8]; + m.data[9] = data[i * 4 + j + 9]; + m.data[10] = data[i * 4 + j + 10]; + m.data[11] = data[i * 4 + j + 11]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x4(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 3] = m.data[3]; + data[i * 4 + j + 4] = m.data[4]; + data[i * 4 + j + 5] = m.data[5]; + data[i * 4 + j + 6] = m.data[6]; + data[i * 4 + j + 7] = m.data[7]; + data[i * 4 + j + 8] = m.data[8]; + data[i * 4 + j + 9] = m.data[9]; + data[i * 4 + j + 10] = m.data[10]; + data[i * 4 + j + 11] = m.data[11]; + + return *this; + } + + /// Forms a 3-by-4 matrix by horizontally concatenating a 3-by-1 matrix with a 3-by-3 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), rhs.at(0, 0), rhs.at(0, 1), rhs.at(0, 2) + , lhs.at(1, 0), rhs.at(1, 0), rhs.at(1, 1), rhs.at(1, 2) + , lhs.at(2, 0), rhs.at(2, 0), rhs.at(2, 1), rhs.at(2, 2)); + } + + /// Forms a 3-by-4 matrix by horizontally concatenating a 3-by-2 matrix with a 3-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0), rhs.at(0, 1) + , lhs.at(1, 0), lhs.at(1, 1), rhs.at(1, 0), rhs.at(1, 1) + , lhs.at(2, 0), lhs.at(2, 1), rhs.at(2, 0), rhs.at(2, 1)); + } + + /// Forms a 3-by-4 matrix by horizontally concatenating a 3-by-3 matrix with a 3-by-1 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), lhs.at(0, 1), lhs.at(0, 2), rhs.at(0, 0) + , lhs.at(1, 0), lhs.at(1, 1), lhs.at(1, 2), rhs.at(1, 0) + , lhs.at(2, 0), lhs.at(2, 1), lhs.at(2, 2), rhs.at(2, 0)); + } + + /// Forms a 3-by-4 matrix by vertically concatenating a 1-by-4 matrix with a 2-by-4 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1), upper.at(0, 2), upper.at(0, 3) + , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2), lower.at(0, 3) + , lower.at(1, 0), lower.at(1, 1), lower.at(1, 2), lower.at(1, 3)); + } + + /// Forms a 3-by-4 matrix by vertically concatenating a 2-by-4 matrix with a 1-by-4 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1), upper.at(0, 2), upper.at(0, 3) + , upper.at(1, 0), upper.at(1, 1), upper.at(1, 2), upper.at(1, 3) + , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2), lower.at(0, 3)); + } + + /// Concatenates this matrix with a a 1-by-4 matrix to form a 4-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix vcat(Matrix const & rhs) const { + return Matrix::vcat(*this, rhs); + } + + /// Forms a 3-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Element A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A, B.at(0, 0), B.at(0, 1), B.at(0, 2) + , C.at(0, 0), D.at(0, 0), D.at(0, 1), D.at(0, 2) + , C.at(1, 0), D.at(1, 0), D.at(1, 1), D.at(1, 2) + ); + } + + /// Forms a 3-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), A.at(0, 1), B.at(0, 0), B.at(0, 1) + , C.at(0, 0), C.at(0, 1), D.at(0, 0), D.at(0, 1) + , C.at(1, 0), C.at(1, 1), D.at(1, 0), D.at(1, 1) + ); + } + + /// Forms a 3-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Element B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), A.at(0, 1), A.at(0, 2), B + , C.at(0, 0), C.at(0, 1), C.at(0, 2), D.at(0, 0) + , C.at(1, 0), C.at(1, 1), C.at(1, 2), D.at(1, 0) + ); + } + + /// Forms a 3-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Element C, Matrix const & D) { + return Matrix( + A.at(0, 0), B.at(0, 0), B.at(0, 1), B.at(0, 2) + , A.at(1, 0), B.at(1, 0), B.at(1, 1), B.at(1, 2) + , C, D.at(0, 0), D.at(0, 1), D.at(0, 2) + ); + } + + /// Forms a 3-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), A.at(0, 1), B.at(0, 0), B.at(0, 1) + , A.at(1, 0), A.at(1, 1), B.at(1, 0), B.at(1, 1) + , C.at(0, 0), C.at(0, 1), D.at(0, 0), D.at(0, 1) + ); + } + + /// Forms a 3-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Element D) { + return Matrix( + A.at(0, 0), A.at(0, 1), A.at(0, 2), B.at(0, 0) + , A.at(1, 0), A.at(1, 1), A.at(1, 2), B.at(1, 0) + , C.at(0, 0), C.at(0, 1), C.at(0, 2), D + ); + } + + /// Elementwise add operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + result.data[1] = data[1] + rhs.data[1]; + result.data[2] = data[2] + rhs.data[2]; + result.data[3] = data[3] + rhs.data[3]; + + result.data[4] = data[4] + rhs.data[4]; + result.data[5] = data[5] + rhs.data[5]; + result.data[6] = data[6] + rhs.data[6]; + result.data[7] = data[7] + rhs.data[7]; + + result.data[8] = data[8] + rhs.data[8]; + result.data[9] = data[9] + rhs.data[9]; + result.data[10] = data[10] + rhs.data[10]; + result.data[11] = data[11] + rhs.data[11]; + + return result; + } + + /// Elementwise add operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + data[1] += rhs.data[1]; + data[2] += rhs.data[2]; + data[3] += rhs.data[3]; + + data[4] += rhs.data[4]; + data[5] += rhs.data[5]; + data[6] += rhs.data[6]; + data[7] += rhs.data[7]; + + data[8] += rhs.data[8]; + data[9] += rhs.data[9]; + data[10] += rhs.data[10]; + data[11] += rhs.data[11]; + + return *this; + } + + /// Elementwise subtract operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + result.data[1] = data[1] - rhs.data[1]; + result.data[2] = data[2] - rhs.data[2]; + result.data[3] = data[3] - rhs.data[3]; + + result.data[4] = data[4] - rhs.data[4]; + result.data[5] = data[5] - rhs.data[5]; + result.data[6] = data[6] - rhs.data[6]; + result.data[7] = data[7] - rhs.data[7]; + + result.data[8] = data[8] - rhs.data[8]; + result.data[9] = data[9] - rhs.data[9]; + result.data[10] = data[10] - rhs.data[10]; + result.data[11] = data[11] - rhs.data[11]; + + return result; + } + + /// Elementwise subtract operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + data[1] -= rhs.data[1]; + data[2] -= rhs.data[2]; + data[3] -= rhs.data[3]; + + data[4] -= rhs.data[4]; + data[5] -= rhs.data[5]; + data[6] -= rhs.data[6]; + data[7] -= rhs.data[7]; + + data[8] -= rhs.data[8]; + data[9] -= rhs.data[9]; + data[10] -= rhs.data[10]; + data[11] -= rhs.data[11]; + + return *this; + } + + /// Elementwise multiply operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + result.data[1] = data[1] * rhs.data[1]; + result.data[2] = data[2] * rhs.data[2]; + result.data[3] = data[3] * rhs.data[3]; + + result.data[4] = data[4] * rhs.data[4]; + result.data[5] = data[5] * rhs.data[5]; + result.data[6] = data[6] * rhs.data[6]; + result.data[7] = data[7] * rhs.data[7]; + + result.data[8] = data[8] * rhs.data[8]; + result.data[9] = data[9] * rhs.data[9]; + result.data[10] = data[10] * rhs.data[10]; + result.data[11] = data[11] * rhs.data[11]; + + return result; + } + + /// Scalar multiply operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + result.data[1] = data[1] * s; + result.data[2] = data[2] * s; + result.data[3] = data[3] * s; + + result.data[4] = data[4] * s; + result.data[5] = data[5] * s; + result.data[6] = data[6] * s; + result.data[7] = data[7] * s; + + result.data[8] = data[8] * s; + result.data[9] = data[9] * s; + result.data[10] = data[10] * s; + result.data[11] = data[11] * s; + + return result; + } + + /// Scalar multiply operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + data[1] *= s; + data[2] *= s; + data[3] *= s; + + data[4] *= s; + data[5] *= s; + data[6] *= s; + data[7] *= s; + + data[8] *= s; + data[9] *= s; + data[10] *= s; + data[11] *= s; + + return *this; + } + + /// Elementwise divide operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + result.data[1] = data[1] / rhs.data[1]; + result.data[2] = data[2] / rhs.data[2]; + result.data[3] = data[3] / rhs.data[3]; + + result.data[4] = data[4] / rhs.data[4]; + result.data[5] = data[5] / rhs.data[5]; + result.data[6] = data[6] / rhs.data[6]; + result.data[7] = data[7] / rhs.data[7]; + + result.data[8] = data[8] / rhs.data[8]; + result.data[9] = data[9] / rhs.data[9]; + result.data[10] = data[10] / rhs.data[10]; + result.data[11] = data[11] / rhs.data[11]; + + return result; + } + + /// Scalar divide operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + result.data[1] = data[1] / s; + result.data[2] = data[2] / s; + result.data[3] = data[3] / s; + + result.data[4] = data[4] / s; + result.data[5] = data[5] / s; + result.data[6] = data[6] / s; + result.data[7] = data[7] / s; + + result.data[8] = data[8] / s; + result.data[9] = data[9] / s; + result.data[10] = data[10] / s; + result.data[11] = data[11] / s; + + return result; + } + + /// Scalar divide operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + data[1] /= s; + data[2] /= s; + data[3] /= s; + + data[4] /= s; + data[5] /= s; + data[6] /= s; + data[7] /= s; + + data[8] /= s; + data[9] /= s; + data[10] /= s; + data[11] /= s; + + return *this; + } + + /// Elementwise divide operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (3-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + data[1] /= rhs.data[1]; + data[2] /= rhs.data[2]; + data[3] /= rhs.data[3]; + + data[4] /= rhs.data[4]; + data[5] /= rhs.data[5]; + data[6] /= rhs.data[6]; + data[7] /= rhs.data[7]; + + data[8] /= rhs.data[8]; + data[9] /= rhs.data[9]; + data[10] /= rhs.data[10]; + data[11] /= rhs.data[11]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + m.data[2] = -m.data[2]; + m.data[3] = -m.data[3]; + m.data[4] = -m.data[4]; + m.data[5] = -m.data[5]; + m.data[6] = -m.data[6]; + m.data[7] = -m.data[7]; + m.data[8] = -m.data[8]; + m.data[9] = -m.data[9]; + m.data[10] = -m.data[10]; + m.data[11] = -m.data[11]; + + return m; + } + + /// Matrix product of size 3-by-1-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[4] * rhs.data[0]; + accum.data[2] += data[8] * rhs.data[0]; + + // k=1 + accum.data[0] += data[1] * rhs.data[1]; + accum.data[1] += data[5] * rhs.data[1]; + accum.data[2] += data[9] * rhs.data[1]; + + // k=2 + accum.data[0] += data[2] * rhs.data[2]; + accum.data[1] += data[6] * rhs.data[2]; + accum.data[2] += data[10] * rhs.data[2]; + + // k=3 + accum.data[0] += data[3] * rhs.data[3]; + accum.data[1] += data[7] * rhs.data[3]; + accum.data[2] += data[11] * rhs.data[3]; + + return accum; + } + + /// Matrix product of size 3-by-1-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 3-by-2-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[4] * rhs.data[0]; + accum.data[3] += data[4] * rhs.data[1]; + accum.data[4] += data[8] * rhs.data[0]; + accum.data[5] += data[8] * rhs.data[1]; + + // k=1 + accum.data[0] += data[1] * rhs.data[2]; + accum.data[1] += data[1] * rhs.data[3]; + accum.data[2] += data[5] * rhs.data[2]; + accum.data[3] += data[5] * rhs.data[3]; + accum.data[4] += data[9] * rhs.data[2]; + accum.data[5] += data[9] * rhs.data[3]; + + // k=2 + accum.data[0] += data[2] * rhs.data[4]; + accum.data[1] += data[2] * rhs.data[5]; + accum.data[2] += data[6] * rhs.data[4]; + accum.data[3] += data[6] * rhs.data[5]; + accum.data[4] += data[10] * rhs.data[4]; + accum.data[5] += data[10] * rhs.data[5]; + + // k=3 + accum.data[0] += data[3] * rhs.data[6]; + accum.data[1] += data[3] * rhs.data[7]; + accum.data[2] += data[7] * rhs.data[6]; + accum.data[3] += data[7] * rhs.data[7]; + accum.data[4] += data[11] * rhs.data[6]; + accum.data[5] += data[11] * rhs.data[7]; + + return accum; + } + + /// Matrix product of size 3-by-2-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 3-by-3-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[4] * rhs.data[0]; + accum.data[4] += data[4] * rhs.data[1]; + accum.data[5] += data[4] * rhs.data[2]; + accum.data[6] += data[8] * rhs.data[0]; + accum.data[7] += data[8] * rhs.data[1]; + accum.data[8] += data[8] * rhs.data[2]; + + // k=1 + accum.data[0] += data[1] * rhs.data[3]; + accum.data[1] += data[1] * rhs.data[4]; + accum.data[2] += data[1] * rhs.data[5]; + accum.data[3] += data[5] * rhs.data[3]; + accum.data[4] += data[5] * rhs.data[4]; + accum.data[5] += data[5] * rhs.data[5]; + accum.data[6] += data[9] * rhs.data[3]; + accum.data[7] += data[9] * rhs.data[4]; + accum.data[8] += data[9] * rhs.data[5]; + + // k=2 + accum.data[0] += data[2] * rhs.data[6]; + accum.data[1] += data[2] * rhs.data[7]; + accum.data[2] += data[2] * rhs.data[8]; + accum.data[3] += data[6] * rhs.data[6]; + accum.data[4] += data[6] * rhs.data[7]; + accum.data[5] += data[6] * rhs.data[8]; + accum.data[6] += data[10] * rhs.data[6]; + accum.data[7] += data[10] * rhs.data[7]; + accum.data[8] += data[10] * rhs.data[8]; + + // k=3 + accum.data[0] += data[3] * rhs.data[9]; + accum.data[1] += data[3] * rhs.data[10]; + accum.data[2] += data[3] * rhs.data[11]; + accum.data[3] += data[7] * rhs.data[9]; + accum.data[4] += data[7] * rhs.data[10]; + accum.data[5] += data[7] * rhs.data[11]; + accum.data[6] += data[11] * rhs.data[9]; + accum.data[7] += data[11] * rhs.data[10]; + accum.data[8] += data[11] * rhs.data[11]; + + return accum; + } + + /// Matrix product of size 3-by-3-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 3-by-4-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + accum.data[4] += data[4] * rhs.data[0]; + accum.data[5] += data[4] * rhs.data[1]; + accum.data[6] += data[4] * rhs.data[2]; + accum.data[7] += data[4] * rhs.data[3]; + accum.data[8] += data[8] * rhs.data[0]; + accum.data[9] += data[8] * rhs.data[1]; + accum.data[10] += data[8] * rhs.data[2]; + accum.data[11] += data[8] * rhs.data[3]; + + // k=1 + accum.data[0] += data[1] * rhs.data[4]; + accum.data[1] += data[1] * rhs.data[5]; + accum.data[2] += data[1] * rhs.data[6]; + accum.data[3] += data[1] * rhs.data[7]; + accum.data[4] += data[5] * rhs.data[4]; + accum.data[5] += data[5] * rhs.data[5]; + accum.data[6] += data[5] * rhs.data[6]; + accum.data[7] += data[5] * rhs.data[7]; + accum.data[8] += data[9] * rhs.data[4]; + accum.data[9] += data[9] * rhs.data[5]; + accum.data[10] += data[9] * rhs.data[6]; + accum.data[11] += data[9] * rhs.data[7]; + + // k=2 + accum.data[0] += data[2] * rhs.data[8]; + accum.data[1] += data[2] * rhs.data[9]; + accum.data[2] += data[2] * rhs.data[10]; + accum.data[3] += data[2] * rhs.data[11]; + accum.data[4] += data[6] * rhs.data[8]; + accum.data[5] += data[6] * rhs.data[9]; + accum.data[6] += data[6] * rhs.data[10]; + accum.data[7] += data[6] * rhs.data[11]; + accum.data[8] += data[10] * rhs.data[8]; + accum.data[9] += data[10] * rhs.data[9]; + accum.data[10] += data[10] * rhs.data[10]; + accum.data[11] += data[10] * rhs.data[11]; + + // k=3 + accum.data[0] += data[3] * rhs.data[12]; + accum.data[1] += data[3] * rhs.data[13]; + accum.data[2] += data[3] * rhs.data[14]; + accum.data[3] += data[3] * rhs.data[15]; + accum.data[4] += data[7] * rhs.data[12]; + accum.data[5] += data[7] * rhs.data[13]; + accum.data[6] += data[7] * rhs.data[14]; + accum.data[7] += data[7] * rhs.data[15]; + accum.data[8] += data[11] * rhs.data[12]; + accum.data[9] += data[11] * rhs.data[13]; + accum.data[10] += data[11] * rhs.data[14]; + accum.data[11] += data[11] * rhs.data[15]; + + return accum; + } + + /// Matrix product of size 3-by-4-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 3-by-4-by-4 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + accum += data[2]; + accum += data[3]; + accum += data[4]; + accum += data[5]; + accum += data[6]; + accum += data[7]; + accum += data[8]; + accum += data[9]; + accum += data[10]; + accum += data[11]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + accum += data[2] * data[2]; + accum += data[3] * data[3]; + accum += data[4] * data[4]; + accum += data[5] * data[5]; + accum += data[6] * data[6]; + accum += data[7] * data[7]; + accum += data[8] * data[8]; + accum += data[9] * data[9]; + accum += data[10] * data[10]; + accum += data[11] * data[11]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + accum += data[5]; + accum += data[10]; + + return accum; + } + +}; + +/// Template alias for 3-by-4 matrix +template +using Matrix3x4 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix3x4 make_Matrix3x4( + Element _0_0, Element _0_1, Element _0_2, Element _0_3, + Element _1_0, Element _1_1, Element _1_2, Element _1_3, + Element _2_0, Element _2_1, Element _2_2, Element _2_3 +) { + return Matrix3x4( + _0_0, _0_1, _0_2, _0_3, + _1_0, _1_1, _1_2, _1_3, + _2_0, _2_1, _2_2, _2_3 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 4-by-1 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 4; + + /// Number of columns in matrix + static int const kColumns = 1; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 4; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 4-by-1 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 4-by-1 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, + Element _1_0, + Element _2_0, + Element _3_0 + ) { + + data[0] = _0_0; + data[1] = _1_0; + data[2] = _2_0; + data[3] = _3_0; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + m.data[2] = s; + m.data[3] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[1] = data[1]; + mt.data[2] = data[2]; + mt.data[3] = data[3]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 4 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 4 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 1 + j + 0]; + m.data[1] = data[i * 1 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 1 + j + 0] = m.data[0]; + data[i * 1 + j + 1] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 1 + j + 0]; + m.data[1] = data[i * 1 + j + 1]; + m.data[2] = data[i * 1 + j + 2]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 1 + j + 0] = m.data[0]; + data[i * 1 + j + 1] = m.data[1]; + data[i * 1 + j + 2] = m.data[2]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_4x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 1 + j + 0]; + m.data[1] = data[i * 1 + j + 1]; + m.data[2] = data[i * 1 + j + 2]; + m.data[3] = data[i * 1 + j + 3]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_4x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 1 + j + 0] = m.data[0]; + data[i * 1 + j + 1] = m.data[1]; + data[i * 1 + j + 2] = m.data[2]; + data[i * 1 + j + 3] = m.data[3]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix column(int j) const { + return slice_4x1(0, j); + } + + MCTLASS_HOST_DEVICE + Matrix &set_column(Matrix const &v, int j =0) { + return set_slice_4x1(v, 0, j); + } + + /// Concatenates this matrix with a a 4-by-1 matrix to form a 4-by-2 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Concatenates this matrix with a a 4-by-2 matrix to form a 4-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Concatenates this matrix with a a 4-by-3 matrix to form a 4-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Forms a 4-by-1 matrix by vertically concatenating an Element with a 3-by-1 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Element upper, Matrix const & lower) { + return Matrix( + upper + , lower.at(0, 0) + , lower.at(1, 0) + , lower.at(2, 0)); + } + + /// Forms a 4-by-1 matrix by vertically concatenating a 2-by-1 matrix with a 2-by-1 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0) + , upper.at(1, 0) + , lower.at(0, 0) + , lower.at(1, 0)); + } + + /// Forms a 4-by-1 matrix by vertically concatenating a 3-by-1 matrix with an Element + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Element lower) { + return Matrix( + upper.at(0, 0) + , upper.at(1, 0) + , upper.at(2, 0) + , lower); + } + + /// Elementwise add operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + + result.data[1] = data[1] + rhs.data[1]; + + result.data[2] = data[2] + rhs.data[2]; + + result.data[3] = data[3] + rhs.data[3]; + + return result; + } + + /// Elementwise add operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + + data[1] += rhs.data[1]; + + data[2] += rhs.data[2]; + + data[3] += rhs.data[3]; + + return *this; + } + + /// Elementwise subtract operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + + result.data[1] = data[1] - rhs.data[1]; + + result.data[2] = data[2] - rhs.data[2]; + + result.data[3] = data[3] - rhs.data[3]; + + return result; + } + + /// Elementwise subtract operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + + data[1] -= rhs.data[1]; + + data[2] -= rhs.data[2]; + + data[3] -= rhs.data[3]; + + return *this; + } + + /// Elementwise multiply operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + + result.data[1] = data[1] * rhs.data[1]; + + result.data[2] = data[2] * rhs.data[2]; + + result.data[3] = data[3] * rhs.data[3]; + + return result; + } + + /// Scalar multiply operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + + result.data[1] = data[1] * s; + + result.data[2] = data[2] * s; + + result.data[3] = data[3] * s; + + return result; + } + + /// Scalar multiply operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + + data[1] *= s; + + data[2] *= s; + + data[3] *= s; + + return *this; + } + + /// Elementwise divide operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + + result.data[1] = data[1] / rhs.data[1]; + + result.data[2] = data[2] / rhs.data[2]; + + result.data[3] = data[3] / rhs.data[3]; + + return result; + } + + /// Scalar divide operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + + result.data[1] = data[1] / s; + + result.data[2] = data[2] / s; + + result.data[3] = data[3] / s; + + return result; + } + + /// Scalar divide operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + + data[1] /= s; + + data[2] /= s; + + data[3] /= s; + + return *this; + } + + /// Elementwise divide operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (4-by-1) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + + data[1] /= rhs.data[1]; + + data[2] /= rhs.data[2]; + + data[3] /= rhs.data[3]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + m.data[2] = -m.data[2]; + m.data[3] = -m.data[3]; + + return m; + } + + /// Matrix product of size 4-by-1-by-1 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[1] * rhs.data[0]; + accum.data[2] += data[2] * rhs.data[0]; + accum.data[3] += data[3] * rhs.data[0]; + + return accum; + } + + /// Matrix product of size 4-by-1-by-1 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 4-by-1-by-1 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Matrix product of size 4-by-2-by-1 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[1] * rhs.data[0]; + accum.data[3] += data[1] * rhs.data[1]; + accum.data[4] += data[2] * rhs.data[0]; + accum.data[5] += data[2] * rhs.data[1]; + accum.data[6] += data[3] * rhs.data[0]; + accum.data[7] += data[3] * rhs.data[1]; + + return accum; + } + + /// Matrix product of size 4-by-2-by-1 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 4-by-3-by-1 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[1] * rhs.data[0]; + accum.data[4] += data[1] * rhs.data[1]; + accum.data[5] += data[1] * rhs.data[2]; + accum.data[6] += data[2] * rhs.data[0]; + accum.data[7] += data[2] * rhs.data[1]; + accum.data[8] += data[2] * rhs.data[2]; + accum.data[9] += data[3] * rhs.data[0]; + accum.data[10] += data[3] * rhs.data[1]; + accum.data[11] += data[3] * rhs.data[2]; + + return accum; + } + + /// Matrix product of size 4-by-3-by-1 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 4-by-4-by-1 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + accum.data[4] += data[1] * rhs.data[0]; + accum.data[5] += data[1] * rhs.data[1]; + accum.data[6] += data[1] * rhs.data[2]; + accum.data[7] += data[1] * rhs.data[3]; + accum.data[8] += data[2] * rhs.data[0]; + accum.data[9] += data[2] * rhs.data[1]; + accum.data[10] += data[2] * rhs.data[2]; + accum.data[11] += data[2] * rhs.data[3]; + accum.data[12] += data[3] * rhs.data[0]; + accum.data[13] += data[3] * rhs.data[1]; + accum.data[14] += data[3] * rhs.data[2]; + accum.data[15] += data[3] * rhs.data[3]; + + return accum; + } + + /// Matrix product of size 4-by-4-by-1 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Dot product of vectors with extent 4 + MCTLASS_HOST_DEVICE + Element dot(Matrix const &rhs, Element accum = Element()) const { + + accum += data[0] * rhs.data[0]; + accum += data[1] * rhs.data[1]; + accum += data[2] * rhs.data[2]; + accum += data[3] * rhs.data[3]; + return accum; + } + + /// Dot product of vectors with extent 4 + MCTLASS_HOST_DEVICE + Element dot(Matrix const &rhs, Element accum = Element()) const { + + accum += data[0] * rhs.data[0]; + accum += data[1] * rhs.data[1]; + accum += data[2] * rhs.data[2]; + accum += data[3] * rhs.data[3]; + return accum; + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + accum += data[2]; + accum += data[3]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + accum += data[2] * data[2]; + accum += data[3] * data[3]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + + return accum; + } + +}; + +/// Template alias for 4-by-1 matrix +template +using Matrix4x1 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix4x1 make_Matrix4x1( + Element _0_0, + Element _1_0, + Element _2_0, + Element _3_0 +) { + return Matrix4x1( + _0_0, + _1_0, + _2_0, + _3_0 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 4-by-2 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 4; + + /// Number of columns in matrix + static int const kColumns = 2; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 8; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 4-by-2 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 4-by-2 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, Element _0_1, + Element _1_0, Element _1_1, + Element _2_0, Element _2_1, + Element _3_0, Element _3_1 + ) { + + data[0] = _0_0; data[1] = _0_1; + data[2] = _1_0; data[3] = _1_1; + data[4] = _2_0; data[5] = _2_1; + data[6] = _3_0; data[7] = _3_1; + } + + /// Constucts a 4-by-2 matrix from row vectors + MCTLASS_HOST_DEVICE + Matrix( + Matrix const &row_0, + Matrix const &row_1, + Matrix const &row_2, + Matrix const &row_3 + ) { + data[0] = row_0.data[0]; + data[1] = row_0.data[1]; + data[2] = row_1.data[0]; + data[3] = row_1.data[1]; + data[4] = row_2.data[0]; + data[5] = row_2.data[1]; + data[6] = row_3.data[0]; + data[7] = row_3.data[1]; + } + + /// Static method to construct a 4-by-2 matrix from column vectors + MCTLASS_HOST_DEVICE + static Matrix from_columns( + Matrix const &column_0, + Matrix const &column_1 + ) { + Matrix result; + + result.data[0] = column_0.data[0]; + result.data[1] = column_1.data[0]; + result.data[2] = column_0.data[1]; + result.data[3] = column_1.data[1]; + result.data[4] = column_0.data[2]; + result.data[5] = column_1.data[2]; + result.data[6] = column_0.data[3]; + result.data[7] = column_1.data[3]; + return result; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + m.data[2] = s; + m.data[3] = s; + m.data[4] = s; + m.data[5] = s; + m.data[6] = s; + m.data[7] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[5] = diag.data[1]; + m.data[10] = diag.data[2]; + m.data[15] = diag.data[3]; + + return m; + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[5] = diag.data[1]; + m.data[10] = diag.data[2]; + m.data[15] = diag.data[3]; + + return m; + } + + /// Gets an array of diagonal elements + MCTLASS_HOST_DEVICE + Matrix diagonal() const { + Matrix diag; + + diag.data[0] = data[0]; + diag.data[1] = data[5]; + diag.data[2] = data[10]; + diag.data[3] = data[15]; + + return diag; + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[4] = data[1]; + mt.data[1] = data[2]; + mt.data[5] = data[3]; + mt.data[2] = data[4]; + mt.data[6] = data[5]; + mt.data[3] = data[6]; + mt.data[7] = data[7]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 4 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 4 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 1] = m.data[1]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix row(int i) const { + return slice_1x2(i, 0); + } + + MCTLASS_HOST_DEVICE + Matrix &set_row(Matrix const &v, int i = 0) { + return set_slice_1x2(v, i, 0); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 2]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 2] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 1]; + m.data[2] = data[i * 2 + j + 2]; + m.data[3] = data[i * 2 + j + 3]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 1] = m.data[1]; + data[i * 2 + j + 2] = m.data[2]; + data[i * 2 + j + 3] = m.data[3]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 2]; + m.data[2] = data[i * 2 + j + 4]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 2] = m.data[1]; + data[i * 2 + j + 4] = m.data[2]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 1]; + m.data[2] = data[i * 2 + j + 2]; + m.data[3] = data[i * 2 + j + 3]; + m.data[4] = data[i * 2 + j + 4]; + m.data[5] = data[i * 2 + j + 5]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 1] = m.data[1]; + data[i * 2 + j + 2] = m.data[2]; + data[i * 2 + j + 3] = m.data[3]; + data[i * 2 + j + 4] = m.data[4]; + data[i * 2 + j + 5] = m.data[5]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_4x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 2]; + m.data[2] = data[i * 2 + j + 4]; + m.data[3] = data[i * 2 + j + 6]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_4x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 2] = m.data[1]; + data[i * 2 + j + 4] = m.data[2]; + data[i * 2 + j + 6] = m.data[3]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix column(int j) const { + return slice_4x1(0, j); + } + + MCTLASS_HOST_DEVICE + Matrix &set_column(Matrix const &v, int j =0) { + return set_slice_4x1(v, 0, j); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_4x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 2 + j + 0]; + m.data[1] = data[i * 2 + j + 1]; + m.data[2] = data[i * 2 + j + 2]; + m.data[3] = data[i * 2 + j + 3]; + m.data[4] = data[i * 2 + j + 4]; + m.data[5] = data[i * 2 + j + 5]; + m.data[6] = data[i * 2 + j + 6]; + m.data[7] = data[i * 2 + j + 7]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_4x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 2 + j + 0] = m.data[0]; + data[i * 2 + j + 1] = m.data[1]; + data[i * 2 + j + 2] = m.data[2]; + data[i * 2 + j + 3] = m.data[3]; + data[i * 2 + j + 4] = m.data[4]; + data[i * 2 + j + 5] = m.data[5]; + data[i * 2 + j + 6] = m.data[6]; + data[i * 2 + j + 7] = m.data[7]; + + return *this; + } + + /// Forms a 4-by-2 matrix by horizontally concatenating a 4-by-1 matrix with a 4-by-1 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), rhs.at(0, 0) + , lhs.at(1, 0), rhs.at(1, 0) + , lhs.at(2, 0), rhs.at(2, 0) + , lhs.at(3, 0), rhs.at(3, 0)); + } + + /// Concatenates this matrix with a a 4-by-1 matrix to form a 4-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Concatenates this matrix with a a 4-by-2 matrix to form a 4-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Forms a 4-by-2 matrix by vertically concatenating a 1-by-2 matrix with a 3-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1) + , lower.at(0, 0), lower.at(0, 1) + , lower.at(1, 0), lower.at(1, 1) + , lower.at(2, 0), lower.at(2, 1)); + } + + /// Forms a 4-by-2 matrix by vertically concatenating a 2-by-2 matrix with a 2-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1) + , upper.at(1, 0), upper.at(1, 1) + , lower.at(0, 0), lower.at(0, 1) + , lower.at(1, 0), lower.at(1, 1)); + } + + /// Forms a 4-by-2 matrix by vertically concatenating a 3-by-2 matrix with a 1-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1) + , upper.at(1, 0), upper.at(1, 1) + , upper.at(2, 0), upper.at(2, 1) + , lower.at(0, 0), lower.at(0, 1)); + } + + /// Forms a 4-by-2 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Element A, Element B, + Matrix const & C, Matrix const & D) { + return Matrix( + A, B + , C.at(0, 0), D.at(0, 0) + , C.at(1, 0), D.at(1, 0) + , C.at(2, 0), D.at(2, 0) + ); + } + + /// Forms a 4-by-2 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), B.at(0, 0) + , A.at(1, 0), B.at(1, 0) + , C.at(0, 0), D.at(0, 0) + , C.at(1, 0), D.at(1, 0) + ); + } + + /// Forms a 4-by-2 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Element C, Element D) { + return Matrix( + A.at(0, 0), B.at(0, 0) + , A.at(1, 0), B.at(1, 0) + , A.at(2, 0), B.at(2, 0) + , C, D + ); + } + + /// Elementwise add operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + result.data[1] = data[1] + rhs.data[1]; + + result.data[2] = data[2] + rhs.data[2]; + result.data[3] = data[3] + rhs.data[3]; + + result.data[4] = data[4] + rhs.data[4]; + result.data[5] = data[5] + rhs.data[5]; + + result.data[6] = data[6] + rhs.data[6]; + result.data[7] = data[7] + rhs.data[7]; + + return result; + } + + /// Elementwise add operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + data[1] += rhs.data[1]; + + data[2] += rhs.data[2]; + data[3] += rhs.data[3]; + + data[4] += rhs.data[4]; + data[5] += rhs.data[5]; + + data[6] += rhs.data[6]; + data[7] += rhs.data[7]; + + return *this; + } + + /// Elementwise subtract operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + result.data[1] = data[1] - rhs.data[1]; + + result.data[2] = data[2] - rhs.data[2]; + result.data[3] = data[3] - rhs.data[3]; + + result.data[4] = data[4] - rhs.data[4]; + result.data[5] = data[5] - rhs.data[5]; + + result.data[6] = data[6] - rhs.data[6]; + result.data[7] = data[7] - rhs.data[7]; + + return result; + } + + /// Elementwise subtract operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + data[1] -= rhs.data[1]; + + data[2] -= rhs.data[2]; + data[3] -= rhs.data[3]; + + data[4] -= rhs.data[4]; + data[5] -= rhs.data[5]; + + data[6] -= rhs.data[6]; + data[7] -= rhs.data[7]; + + return *this; + } + + /// Elementwise multiply operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + result.data[1] = data[1] * rhs.data[1]; + + result.data[2] = data[2] * rhs.data[2]; + result.data[3] = data[3] * rhs.data[3]; + + result.data[4] = data[4] * rhs.data[4]; + result.data[5] = data[5] * rhs.data[5]; + + result.data[6] = data[6] * rhs.data[6]; + result.data[7] = data[7] * rhs.data[7]; + + return result; + } + + /// Scalar multiply operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + result.data[1] = data[1] * s; + + result.data[2] = data[2] * s; + result.data[3] = data[3] * s; + + result.data[4] = data[4] * s; + result.data[5] = data[5] * s; + + result.data[6] = data[6] * s; + result.data[7] = data[7] * s; + + return result; + } + + /// Scalar multiply operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + data[1] *= s; + + data[2] *= s; + data[3] *= s; + + data[4] *= s; + data[5] *= s; + + data[6] *= s; + data[7] *= s; + + return *this; + } + + /// Elementwise divide operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + result.data[1] = data[1] / rhs.data[1]; + + result.data[2] = data[2] / rhs.data[2]; + result.data[3] = data[3] / rhs.data[3]; + + result.data[4] = data[4] / rhs.data[4]; + result.data[5] = data[5] / rhs.data[5]; + + result.data[6] = data[6] / rhs.data[6]; + result.data[7] = data[7] / rhs.data[7]; + + return result; + } + + /// Scalar divide operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + result.data[1] = data[1] / s; + + result.data[2] = data[2] / s; + result.data[3] = data[3] / s; + + result.data[4] = data[4] / s; + result.data[5] = data[5] / s; + + result.data[6] = data[6] / s; + result.data[7] = data[7] / s; + + return result; + } + + /// Scalar divide operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + data[1] /= s; + + data[2] /= s; + data[3] /= s; + + data[4] /= s; + data[5] /= s; + + data[6] /= s; + data[7] /= s; + + return *this; + } + + /// Elementwise divide operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (4-by-2) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + data[1] /= rhs.data[1]; + + data[2] /= rhs.data[2]; + data[3] /= rhs.data[3]; + + data[4] /= rhs.data[4]; + data[5] /= rhs.data[5]; + + data[6] /= rhs.data[6]; + data[7] /= rhs.data[7]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + m.data[2] = -m.data[2]; + m.data[3] = -m.data[3]; + m.data[4] = -m.data[4]; + m.data[5] = -m.data[5]; + m.data[6] = -m.data[6]; + m.data[7] = -m.data[7]; + + return m; + } + + /// Matrix product of size 4-by-1-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[2] * rhs.data[0]; + accum.data[2] += data[4] * rhs.data[0]; + accum.data[3] += data[6] * rhs.data[0]; + + // k=1 + accum.data[0] += data[1] * rhs.data[1]; + accum.data[1] += data[3] * rhs.data[1]; + accum.data[2] += data[5] * rhs.data[1]; + accum.data[3] += data[7] * rhs.data[1]; + + return accum; + } + + /// Matrix product of size 4-by-1-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 4-by-2-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[2] * rhs.data[0]; + accum.data[3] += data[2] * rhs.data[1]; + accum.data[4] += data[4] * rhs.data[0]; + accum.data[5] += data[4] * rhs.data[1]; + accum.data[6] += data[6] * rhs.data[0]; + accum.data[7] += data[6] * rhs.data[1]; + + // k=1 + accum.data[0] += data[1] * rhs.data[2]; + accum.data[1] += data[1] * rhs.data[3]; + accum.data[2] += data[3] * rhs.data[2]; + accum.data[3] += data[3] * rhs.data[3]; + accum.data[4] += data[5] * rhs.data[2]; + accum.data[5] += data[5] * rhs.data[3]; + accum.data[6] += data[7] * rhs.data[2]; + accum.data[7] += data[7] * rhs.data[3]; + + return accum; + } + + /// Matrix product of size 4-by-2-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 4-by-2-by-2 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Matrix product of size 4-by-3-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[2] * rhs.data[0]; + accum.data[4] += data[2] * rhs.data[1]; + accum.data[5] += data[2] * rhs.data[2]; + accum.data[6] += data[4] * rhs.data[0]; + accum.data[7] += data[4] * rhs.data[1]; + accum.data[8] += data[4] * rhs.data[2]; + accum.data[9] += data[6] * rhs.data[0]; + accum.data[10] += data[6] * rhs.data[1]; + accum.data[11] += data[6] * rhs.data[2]; + + // k=1 + accum.data[0] += data[1] * rhs.data[3]; + accum.data[1] += data[1] * rhs.data[4]; + accum.data[2] += data[1] * rhs.data[5]; + accum.data[3] += data[3] * rhs.data[3]; + accum.data[4] += data[3] * rhs.data[4]; + accum.data[5] += data[3] * rhs.data[5]; + accum.data[6] += data[5] * rhs.data[3]; + accum.data[7] += data[5] * rhs.data[4]; + accum.data[8] += data[5] * rhs.data[5]; + accum.data[9] += data[7] * rhs.data[3]; + accum.data[10] += data[7] * rhs.data[4]; + accum.data[11] += data[7] * rhs.data[5]; + + return accum; + } + + /// Matrix product of size 4-by-3-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 4-by-4-by-2 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + accum.data[4] += data[2] * rhs.data[0]; + accum.data[5] += data[2] * rhs.data[1]; + accum.data[6] += data[2] * rhs.data[2]; + accum.data[7] += data[2] * rhs.data[3]; + accum.data[8] += data[4] * rhs.data[0]; + accum.data[9] += data[4] * rhs.data[1]; + accum.data[10] += data[4] * rhs.data[2]; + accum.data[11] += data[4] * rhs.data[3]; + accum.data[12] += data[6] * rhs.data[0]; + accum.data[13] += data[6] * rhs.data[1]; + accum.data[14] += data[6] * rhs.data[2]; + accum.data[15] += data[6] * rhs.data[3]; + + // k=1 + accum.data[0] += data[1] * rhs.data[4]; + accum.data[1] += data[1] * rhs.data[5]; + accum.data[2] += data[1] * rhs.data[6]; + accum.data[3] += data[1] * rhs.data[7]; + accum.data[4] += data[3] * rhs.data[4]; + accum.data[5] += data[3] * rhs.data[5]; + accum.data[6] += data[3] * rhs.data[6]; + accum.data[7] += data[3] * rhs.data[7]; + accum.data[8] += data[5] * rhs.data[4]; + accum.data[9] += data[5] * rhs.data[5]; + accum.data[10] += data[5] * rhs.data[6]; + accum.data[11] += data[5] * rhs.data[7]; + accum.data[12] += data[7] * rhs.data[4]; + accum.data[13] += data[7] * rhs.data[5]; + accum.data[14] += data[7] * rhs.data[6]; + accum.data[15] += data[7] * rhs.data[7]; + + return accum; + } + + /// Matrix product of size 4-by-4-by-2 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + accum += data[2]; + accum += data[3]; + accum += data[4]; + accum += data[5]; + accum += data[6]; + accum += data[7]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + accum += data[2] * data[2]; + accum += data[3] * data[3]; + accum += data[4] * data[4]; + accum += data[5] * data[5]; + accum += data[6] * data[6]; + accum += data[7] * data[7]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + accum += data[3]; + + return accum; + } + +}; + +/// Template alias for 4-by-2 matrix +template +using Matrix4x2 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix4x2 make_Matrix4x2( + Element _0_0, Element _0_1, + Element _1_0, Element _1_1, + Element _2_0, Element _2_1, + Element _3_0, Element _3_1 +) { + return Matrix4x2( + _0_0, _0_1, + _1_0, _1_1, + _2_0, _2_1, + _3_0, _3_1 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 4-by-3 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 4; + + /// Number of columns in matrix + static int const kColumns = 3; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 12; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 4-by-3 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 4-by-3 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, Element _0_1, Element _0_2, + Element _1_0, Element _1_1, Element _1_2, + Element _2_0, Element _2_1, Element _2_2, + Element _3_0, Element _3_1, Element _3_2 + ) { + + data[0] = _0_0; data[1] = _0_1; data[2] = _0_2; + data[3] = _1_0; data[4] = _1_1; data[5] = _1_2; + data[6] = _2_0; data[7] = _2_1; data[8] = _2_2; + data[9] = _3_0; data[10] = _3_1; data[11] = _3_2; + } + + /// Constucts a 4-by-3 matrix from row vectors + MCTLASS_HOST_DEVICE + Matrix( + Matrix const &row_0, + Matrix const &row_1, + Matrix const &row_2, + Matrix const &row_3 + ) { + data[0] = row_0.data[0]; + data[1] = row_0.data[1]; + data[2] = row_0.data[2]; + data[3] = row_1.data[0]; + data[4] = row_1.data[1]; + data[5] = row_1.data[2]; + data[6] = row_2.data[0]; + data[7] = row_2.data[1]; + data[8] = row_2.data[2]; + data[9] = row_3.data[0]; + data[10] = row_3.data[1]; + data[11] = row_3.data[2]; + } + + /// Static method to construct a 4-by-3 matrix from column vectors + MCTLASS_HOST_DEVICE + static Matrix from_columns( + Matrix const &column_0, + Matrix const &column_1, + Matrix const &column_2 + ) { + Matrix result; + + result.data[0] = column_0.data[0]; + result.data[1] = column_1.data[0]; + result.data[2] = column_2.data[0]; + result.data[3] = column_0.data[1]; + result.data[4] = column_1.data[1]; + result.data[5] = column_2.data[1]; + result.data[6] = column_0.data[2]; + result.data[7] = column_1.data[2]; + result.data[8] = column_2.data[2]; + result.data[9] = column_0.data[3]; + result.data[10] = column_1.data[3]; + result.data[11] = column_2.data[3]; + return result; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + m.data[2] = s; + m.data[3] = s; + m.data[4] = s; + m.data[5] = s; + m.data[6] = s; + m.data[7] = s; + m.data[8] = s; + m.data[9] = s; + m.data[10] = s; + m.data[11] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[5] = diag.data[1]; + m.data[10] = diag.data[2]; + m.data[15] = diag.data[3]; + + return m; + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[5] = diag.data[1]; + m.data[10] = diag.data[2]; + m.data[15] = diag.data[3]; + + return m; + } + + /// Gets an array of diagonal elements + MCTLASS_HOST_DEVICE + Matrix diagonal() const { + Matrix diag; + + diag.data[0] = data[0]; + diag.data[1] = data[5]; + diag.data[2] = data[10]; + diag.data[3] = data[15]; + + return diag; + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[4] = data[1]; + mt.data[8] = data[2]; + mt.data[1] = data[3]; + mt.data[5] = data[4]; + mt.data[9] = data[5]; + mt.data[2] = data[6]; + mt.data[6] = data[7]; + mt.data[10] = data[8]; + mt.data[3] = data[9]; + mt.data[7] = data[10]; + mt.data[11] = data[11]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 4 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 4 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 2]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 2] = m.data[2]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix row(int i) const { + return slice_1x3(i, 0); + } + + MCTLASS_HOST_DEVICE + Matrix &set_row(Matrix const &v, int i = 0) { + return set_slice_1x3(v, i, 0); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 3]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 3] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 3]; + m.data[3] = data[i * 3 + j + 4]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 3] = m.data[2]; + data[i * 3 + j + 4] = m.data[3]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 2]; + m.data[3] = data[i * 3 + j + 3]; + m.data[4] = data[i * 3 + j + 4]; + m.data[5] = data[i * 3 + j + 5]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 2] = m.data[2]; + data[i * 3 + j + 3] = m.data[3]; + data[i * 3 + j + 4] = m.data[4]; + data[i * 3 + j + 5] = m.data[5]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 3]; + m.data[2] = data[i * 3 + j + 6]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 3] = m.data[1]; + data[i * 3 + j + 6] = m.data[2]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 3]; + m.data[3] = data[i * 3 + j + 4]; + m.data[4] = data[i * 3 + j + 6]; + m.data[5] = data[i * 3 + j + 7]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 3] = m.data[2]; + data[i * 3 + j + 4] = m.data[3]; + data[i * 3 + j + 6] = m.data[4]; + data[i * 3 + j + 7] = m.data[5]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 2]; + m.data[3] = data[i * 3 + j + 3]; + m.data[4] = data[i * 3 + j + 4]; + m.data[5] = data[i * 3 + j + 5]; + m.data[6] = data[i * 3 + j + 6]; + m.data[7] = data[i * 3 + j + 7]; + m.data[8] = data[i * 3 + j + 8]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 2] = m.data[2]; + data[i * 3 + j + 3] = m.data[3]; + data[i * 3 + j + 4] = m.data[4]; + data[i * 3 + j + 5] = m.data[5]; + data[i * 3 + j + 6] = m.data[6]; + data[i * 3 + j + 7] = m.data[7]; + data[i * 3 + j + 8] = m.data[8]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_4x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 3]; + m.data[2] = data[i * 3 + j + 6]; + m.data[3] = data[i * 3 + j + 9]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_4x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 3] = m.data[1]; + data[i * 3 + j + 6] = m.data[2]; + data[i * 3 + j + 9] = m.data[3]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix column(int j) const { + return slice_4x1(0, j); + } + + MCTLASS_HOST_DEVICE + Matrix &set_column(Matrix const &v, int j =0) { + return set_slice_4x1(v, 0, j); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_4x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 3]; + m.data[3] = data[i * 3 + j + 4]; + m.data[4] = data[i * 3 + j + 6]; + m.data[5] = data[i * 3 + j + 7]; + m.data[6] = data[i * 3 + j + 9]; + m.data[7] = data[i * 3 + j + 10]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_4x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 3] = m.data[2]; + data[i * 3 + j + 4] = m.data[3]; + data[i * 3 + j + 6] = m.data[4]; + data[i * 3 + j + 7] = m.data[5]; + data[i * 3 + j + 9] = m.data[6]; + data[i * 3 + j + 10] = m.data[7]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_4x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 3 + j + 0]; + m.data[1] = data[i * 3 + j + 1]; + m.data[2] = data[i * 3 + j + 2]; + m.data[3] = data[i * 3 + j + 3]; + m.data[4] = data[i * 3 + j + 4]; + m.data[5] = data[i * 3 + j + 5]; + m.data[6] = data[i * 3 + j + 6]; + m.data[7] = data[i * 3 + j + 7]; + m.data[8] = data[i * 3 + j + 8]; + m.data[9] = data[i * 3 + j + 9]; + m.data[10] = data[i * 3 + j + 10]; + m.data[11] = data[i * 3 + j + 11]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_4x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 3 + j + 0] = m.data[0]; + data[i * 3 + j + 1] = m.data[1]; + data[i * 3 + j + 2] = m.data[2]; + data[i * 3 + j + 3] = m.data[3]; + data[i * 3 + j + 4] = m.data[4]; + data[i * 3 + j + 5] = m.data[5]; + data[i * 3 + j + 6] = m.data[6]; + data[i * 3 + j + 7] = m.data[7]; + data[i * 3 + j + 8] = m.data[8]; + data[i * 3 + j + 9] = m.data[9]; + data[i * 3 + j + 10] = m.data[10]; + data[i * 3 + j + 11] = m.data[11]; + + return *this; + } + + /// Forms a 4-by-3 matrix by horizontally concatenating a 4-by-1 matrix with a 4-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), rhs.at(0, 0), rhs.at(0, 1) + , lhs.at(1, 0), rhs.at(1, 0), rhs.at(1, 1) + , lhs.at(2, 0), rhs.at(2, 0), rhs.at(2, 1) + , lhs.at(3, 0), rhs.at(3, 0), rhs.at(3, 1)); + } + + /// Forms a 4-by-3 matrix by horizontally concatenating a 4-by-2 matrix with a 4-by-1 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0) + , lhs.at(1, 0), lhs.at(1, 1), rhs.at(1, 0) + , lhs.at(2, 0), lhs.at(2, 1), rhs.at(2, 0) + , lhs.at(3, 0), lhs.at(3, 1), rhs.at(3, 0)); + } + + /// Concatenates this matrix with a a 4-by-1 matrix to form a 4-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix hcat(Matrix const & rhs) const { + return Matrix::hcat(*this, rhs); + } + + /// Forms a 4-by-3 matrix by vertically concatenating a 1-by-3 matrix with a 3-by-3 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1), upper.at(0, 2) + , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2) + , lower.at(1, 0), lower.at(1, 1), lower.at(1, 2) + , lower.at(2, 0), lower.at(2, 1), lower.at(2, 2)); + } + + /// Forms a 4-by-3 matrix by vertically concatenating a 2-by-3 matrix with a 2-by-3 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1), upper.at(0, 2) + , upper.at(1, 0), upper.at(1, 1), upper.at(1, 2) + , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2) + , lower.at(1, 0), lower.at(1, 1), lower.at(1, 2)); + } + + /// Forms a 4-by-3 matrix by vertically concatenating a 3-by-3 matrix with a 1-by-3 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1), upper.at(0, 2) + , upper.at(1, 0), upper.at(1, 1), upper.at(1, 2) + , upper.at(2, 0), upper.at(2, 1), upper.at(2, 2) + , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2)); + } + + /// Forms a 4-by-3 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Element A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A, B.at(0, 0), B.at(0, 1) + , C.at(0, 0), D.at(0, 0), D.at(0, 1) + , C.at(1, 0), D.at(1, 0), D.at(1, 1) + , C.at(2, 0), D.at(2, 0), D.at(2, 1) + ); + } + + /// Forms a 4-by-3 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Element B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), A.at(0, 1), B + , C.at(0, 0), C.at(0, 1), D.at(0, 0) + , C.at(1, 0), C.at(1, 1), D.at(1, 0) + , C.at(2, 0), C.at(2, 1), D.at(2, 0) + ); + } + + /// Forms a 4-by-3 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), B.at(0, 0), B.at(0, 1) + , A.at(1, 0), B.at(1, 0), B.at(1, 1) + , C.at(0, 0), D.at(0, 0), D.at(0, 1) + , C.at(1, 0), D.at(1, 0), D.at(1, 1) + ); + } + + /// Forms a 4-by-3 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), A.at(0, 1), B.at(0, 0) + , A.at(1, 0), A.at(1, 1), B.at(1, 0) + , C.at(0, 0), C.at(0, 1), D.at(0, 0) + , C.at(1, 0), C.at(1, 1), D.at(1, 0) + ); + } + + /// Forms a 4-by-3 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Element C, Matrix const & D) { + return Matrix( + A.at(0, 0), B.at(0, 0), B.at(0, 1) + , A.at(1, 0), B.at(1, 0), B.at(1, 1) + , A.at(2, 0), B.at(2, 0), B.at(2, 1) + , C, D.at(0, 0), D.at(0, 1) + ); + } + + /// Forms a 4-by-3 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Element D) { + return Matrix( + A.at(0, 0), A.at(0, 1), B.at(0, 0) + , A.at(1, 0), A.at(1, 1), B.at(1, 0) + , A.at(2, 0), A.at(2, 1), B.at(2, 0) + , C.at(0, 0), C.at(0, 1), D + ); + } + + /// Elementwise add operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + result.data[1] = data[1] + rhs.data[1]; + result.data[2] = data[2] + rhs.data[2]; + + result.data[3] = data[3] + rhs.data[3]; + result.data[4] = data[4] + rhs.data[4]; + result.data[5] = data[5] + rhs.data[5]; + + result.data[6] = data[6] + rhs.data[6]; + result.data[7] = data[7] + rhs.data[7]; + result.data[8] = data[8] + rhs.data[8]; + + result.data[9] = data[9] + rhs.data[9]; + result.data[10] = data[10] + rhs.data[10]; + result.data[11] = data[11] + rhs.data[11]; + + return result; + } + + /// Elementwise add operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + data[1] += rhs.data[1]; + data[2] += rhs.data[2]; + + data[3] += rhs.data[3]; + data[4] += rhs.data[4]; + data[5] += rhs.data[5]; + + data[6] += rhs.data[6]; + data[7] += rhs.data[7]; + data[8] += rhs.data[8]; + + data[9] += rhs.data[9]; + data[10] += rhs.data[10]; + data[11] += rhs.data[11]; + + return *this; + } + + /// Elementwise subtract operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + result.data[1] = data[1] - rhs.data[1]; + result.data[2] = data[2] - rhs.data[2]; + + result.data[3] = data[3] - rhs.data[3]; + result.data[4] = data[4] - rhs.data[4]; + result.data[5] = data[5] - rhs.data[5]; + + result.data[6] = data[6] - rhs.data[6]; + result.data[7] = data[7] - rhs.data[7]; + result.data[8] = data[8] - rhs.data[8]; + + result.data[9] = data[9] - rhs.data[9]; + result.data[10] = data[10] - rhs.data[10]; + result.data[11] = data[11] - rhs.data[11]; + + return result; + } + + /// Elementwise subtract operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + data[1] -= rhs.data[1]; + data[2] -= rhs.data[2]; + + data[3] -= rhs.data[3]; + data[4] -= rhs.data[4]; + data[5] -= rhs.data[5]; + + data[6] -= rhs.data[6]; + data[7] -= rhs.data[7]; + data[8] -= rhs.data[8]; + + data[9] -= rhs.data[9]; + data[10] -= rhs.data[10]; + data[11] -= rhs.data[11]; + + return *this; + } + + /// Elementwise multiply operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + result.data[1] = data[1] * rhs.data[1]; + result.data[2] = data[2] * rhs.data[2]; + + result.data[3] = data[3] * rhs.data[3]; + result.data[4] = data[4] * rhs.data[4]; + result.data[5] = data[5] * rhs.data[5]; + + result.data[6] = data[6] * rhs.data[6]; + result.data[7] = data[7] * rhs.data[7]; + result.data[8] = data[8] * rhs.data[8]; + + result.data[9] = data[9] * rhs.data[9]; + result.data[10] = data[10] * rhs.data[10]; + result.data[11] = data[11] * rhs.data[11]; + + return result; + } + + /// Scalar multiply operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + result.data[1] = data[1] * s; + result.data[2] = data[2] * s; + + result.data[3] = data[3] * s; + result.data[4] = data[4] * s; + result.data[5] = data[5] * s; + + result.data[6] = data[6] * s; + result.data[7] = data[7] * s; + result.data[8] = data[8] * s; + + result.data[9] = data[9] * s; + result.data[10] = data[10] * s; + result.data[11] = data[11] * s; + + return result; + } + + /// Scalar multiply operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + data[1] *= s; + data[2] *= s; + + data[3] *= s; + data[4] *= s; + data[5] *= s; + + data[6] *= s; + data[7] *= s; + data[8] *= s; + + data[9] *= s; + data[10] *= s; + data[11] *= s; + + return *this; + } + + /// Elementwise divide operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + result.data[1] = data[1] / rhs.data[1]; + result.data[2] = data[2] / rhs.data[2]; + + result.data[3] = data[3] / rhs.data[3]; + result.data[4] = data[4] / rhs.data[4]; + result.data[5] = data[5] / rhs.data[5]; + + result.data[6] = data[6] / rhs.data[6]; + result.data[7] = data[7] / rhs.data[7]; + result.data[8] = data[8] / rhs.data[8]; + + result.data[9] = data[9] / rhs.data[9]; + result.data[10] = data[10] / rhs.data[10]; + result.data[11] = data[11] / rhs.data[11]; + + return result; + } + + /// Scalar divide operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + result.data[1] = data[1] / s; + result.data[2] = data[2] / s; + + result.data[3] = data[3] / s; + result.data[4] = data[4] / s; + result.data[5] = data[5] / s; + + result.data[6] = data[6] / s; + result.data[7] = data[7] / s; + result.data[8] = data[8] / s; + + result.data[9] = data[9] / s; + result.data[10] = data[10] / s; + result.data[11] = data[11] / s; + + return result; + } + + /// Scalar divide operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + data[1] /= s; + data[2] /= s; + + data[3] /= s; + data[4] /= s; + data[5] /= s; + + data[6] /= s; + data[7] /= s; + data[8] /= s; + + data[9] /= s; + data[10] /= s; + data[11] /= s; + + return *this; + } + + /// Elementwise divide operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (4-by-3) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + data[1] /= rhs.data[1]; + data[2] /= rhs.data[2]; + + data[3] /= rhs.data[3]; + data[4] /= rhs.data[4]; + data[5] /= rhs.data[5]; + + data[6] /= rhs.data[6]; + data[7] /= rhs.data[7]; + data[8] /= rhs.data[8]; + + data[9] /= rhs.data[9]; + data[10] /= rhs.data[10]; + data[11] /= rhs.data[11]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + m.data[2] = -m.data[2]; + m.data[3] = -m.data[3]; + m.data[4] = -m.data[4]; + m.data[5] = -m.data[5]; + m.data[6] = -m.data[6]; + m.data[7] = -m.data[7]; + m.data[8] = -m.data[8]; + m.data[9] = -m.data[9]; + m.data[10] = -m.data[10]; + m.data[11] = -m.data[11]; + + return m; + } + + /// Matrix product of size 4-by-1-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[3] * rhs.data[0]; + accum.data[2] += data[6] * rhs.data[0]; + accum.data[3] += data[9] * rhs.data[0]; + + // k=1 + accum.data[0] += data[1] * rhs.data[1]; + accum.data[1] += data[4] * rhs.data[1]; + accum.data[2] += data[7] * rhs.data[1]; + accum.data[3] += data[10] * rhs.data[1]; + + // k=2 + accum.data[0] += data[2] * rhs.data[2]; + accum.data[1] += data[5] * rhs.data[2]; + accum.data[2] += data[8] * rhs.data[2]; + accum.data[3] += data[11] * rhs.data[2]; + + return accum; + } + + /// Matrix product of size 4-by-1-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 4-by-2-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[3] * rhs.data[0]; + accum.data[3] += data[3] * rhs.data[1]; + accum.data[4] += data[6] * rhs.data[0]; + accum.data[5] += data[6] * rhs.data[1]; + accum.data[6] += data[9] * rhs.data[0]; + accum.data[7] += data[9] * rhs.data[1]; + + // k=1 + accum.data[0] += data[1] * rhs.data[2]; + accum.data[1] += data[1] * rhs.data[3]; + accum.data[2] += data[4] * rhs.data[2]; + accum.data[3] += data[4] * rhs.data[3]; + accum.data[4] += data[7] * rhs.data[2]; + accum.data[5] += data[7] * rhs.data[3]; + accum.data[6] += data[10] * rhs.data[2]; + accum.data[7] += data[10] * rhs.data[3]; + + // k=2 + accum.data[0] += data[2] * rhs.data[4]; + accum.data[1] += data[2] * rhs.data[5]; + accum.data[2] += data[5] * rhs.data[4]; + accum.data[3] += data[5] * rhs.data[5]; + accum.data[4] += data[8] * rhs.data[4]; + accum.data[5] += data[8] * rhs.data[5]; + accum.data[6] += data[11] * rhs.data[4]; + accum.data[7] += data[11] * rhs.data[5]; + + return accum; + } + + /// Matrix product of size 4-by-2-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 4-by-3-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[3] * rhs.data[0]; + accum.data[4] += data[3] * rhs.data[1]; + accum.data[5] += data[3] * rhs.data[2]; + accum.data[6] += data[6] * rhs.data[0]; + accum.data[7] += data[6] * rhs.data[1]; + accum.data[8] += data[6] * rhs.data[2]; + accum.data[9] += data[9] * rhs.data[0]; + accum.data[10] += data[9] * rhs.data[1]; + accum.data[11] += data[9] * rhs.data[2]; + + // k=1 + accum.data[0] += data[1] * rhs.data[3]; + accum.data[1] += data[1] * rhs.data[4]; + accum.data[2] += data[1] * rhs.data[5]; + accum.data[3] += data[4] * rhs.data[3]; + accum.data[4] += data[4] * rhs.data[4]; + accum.data[5] += data[4] * rhs.data[5]; + accum.data[6] += data[7] * rhs.data[3]; + accum.data[7] += data[7] * rhs.data[4]; + accum.data[8] += data[7] * rhs.data[5]; + accum.data[9] += data[10] * rhs.data[3]; + accum.data[10] += data[10] * rhs.data[4]; + accum.data[11] += data[10] * rhs.data[5]; + + // k=2 + accum.data[0] += data[2] * rhs.data[6]; + accum.data[1] += data[2] * rhs.data[7]; + accum.data[2] += data[2] * rhs.data[8]; + accum.data[3] += data[5] * rhs.data[6]; + accum.data[4] += data[5] * rhs.data[7]; + accum.data[5] += data[5] * rhs.data[8]; + accum.data[6] += data[8] * rhs.data[6]; + accum.data[7] += data[8] * rhs.data[7]; + accum.data[8] += data[8] * rhs.data[8]; + accum.data[9] += data[11] * rhs.data[6]; + accum.data[10] += data[11] * rhs.data[7]; + accum.data[11] += data[11] * rhs.data[8]; + + return accum; + } + + /// Matrix product of size 4-by-3-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 4-by-3-by-3 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Matrix product of size 4-by-4-by-3 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + accum.data[4] += data[3] * rhs.data[0]; + accum.data[5] += data[3] * rhs.data[1]; + accum.data[6] += data[3] * rhs.data[2]; + accum.data[7] += data[3] * rhs.data[3]; + accum.data[8] += data[6] * rhs.data[0]; + accum.data[9] += data[6] * rhs.data[1]; + accum.data[10] += data[6] * rhs.data[2]; + accum.data[11] += data[6] * rhs.data[3]; + accum.data[12] += data[9] * rhs.data[0]; + accum.data[13] += data[9] * rhs.data[1]; + accum.data[14] += data[9] * rhs.data[2]; + accum.data[15] += data[9] * rhs.data[3]; + + // k=1 + accum.data[0] += data[1] * rhs.data[4]; + accum.data[1] += data[1] * rhs.data[5]; + accum.data[2] += data[1] * rhs.data[6]; + accum.data[3] += data[1] * rhs.data[7]; + accum.data[4] += data[4] * rhs.data[4]; + accum.data[5] += data[4] * rhs.data[5]; + accum.data[6] += data[4] * rhs.data[6]; + accum.data[7] += data[4] * rhs.data[7]; + accum.data[8] += data[7] * rhs.data[4]; + accum.data[9] += data[7] * rhs.data[5]; + accum.data[10] += data[7] * rhs.data[6]; + accum.data[11] += data[7] * rhs.data[7]; + accum.data[12] += data[10] * rhs.data[4]; + accum.data[13] += data[10] * rhs.data[5]; + accum.data[14] += data[10] * rhs.data[6]; + accum.data[15] += data[10] * rhs.data[7]; + + // k=2 + accum.data[0] += data[2] * rhs.data[8]; + accum.data[1] += data[2] * rhs.data[9]; + accum.data[2] += data[2] * rhs.data[10]; + accum.data[3] += data[2] * rhs.data[11]; + accum.data[4] += data[5] * rhs.data[8]; + accum.data[5] += data[5] * rhs.data[9]; + accum.data[6] += data[5] * rhs.data[10]; + accum.data[7] += data[5] * rhs.data[11]; + accum.data[8] += data[8] * rhs.data[8]; + accum.data[9] += data[8] * rhs.data[9]; + accum.data[10] += data[8] * rhs.data[10]; + accum.data[11] += data[8] * rhs.data[11]; + accum.data[12] += data[11] * rhs.data[8]; + accum.data[13] += data[11] * rhs.data[9]; + accum.data[14] += data[11] * rhs.data[10]; + accum.data[15] += data[11] * rhs.data[11]; + + return accum; + } + + /// Matrix product of size 4-by-4-by-3 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + accum += data[2]; + accum += data[3]; + accum += data[4]; + accum += data[5]; + accum += data[6]; + accum += data[7]; + accum += data[8]; + accum += data[9]; + accum += data[10]; + accum += data[11]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + accum += data[2] * data[2]; + accum += data[3] * data[3]; + accum += data[4] * data[4]; + accum += data[5] * data[5]; + accum += data[6] * data[6]; + accum += data[7] * data[7]; + accum += data[8] * data[8]; + accum += data[9] * data[9]; + accum += data[10] * data[10]; + accum += data[11] * data[11]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + accum += data[4]; + accum += data[8]; + + return accum; + } + +}; + +/// Template alias for 4-by-3 matrix +template +using Matrix4x3 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix4x3 make_Matrix4x3( + Element _0_0, Element _0_1, Element _0_2, + Element _1_0, Element _1_1, Element _1_2, + Element _2_0, Element _2_1, Element _2_2, + Element _3_0, Element _3_1, Element _3_2 +) { + return Matrix4x3( + _0_0, _0_1, _0_2, + _1_0, _1_1, _1_2, + _2_0, _2_1, _2_2, + _3_0, _3_1, _3_2 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// 4-by-4 matrix template class definition +template +struct Matrix { + + // + // Type definitions + // + + /// Element data type + using Element = Element_; + + /// Number of rows in matrix + static int const kRows = 4; + + /// Number of columns in matrix + static int const kColumns = 4; + + /// Layout of matrix in underlying array + using Layout = layout::RowMajor; + + /// Number of elements in matrix + static int const kCount = 16; + + // + // Data members + // + + /// Elements of the matrix in row-major layout + Array data; + + // + // Methods + // + + /// Constructs a zero matrix + MCTLASS_HOST_DEVICE + Matrix() { + data.clear(); + } + + /// Copy constructor for a 4-by-4 matrix + MCTLASS_HOST_DEVICE + Matrix(Matrix const &rhs) { + data = rhs.data; + } + + /// Constucts a 4-by-4 matrix from scalar elements + MCTLASS_HOST_DEVICE + Matrix( + Element _0_0, Element _0_1, Element _0_2, Element _0_3, + Element _1_0, Element _1_1, Element _1_2, Element _1_3, + Element _2_0, Element _2_1, Element _2_2, Element _2_3, + Element _3_0, Element _3_1, Element _3_2, Element _3_3 + ) { + + data[0] = _0_0; data[1] = _0_1; data[2] = _0_2; data[3] = _0_3; + data[4] = _1_0; data[5] = _1_1; data[6] = _1_2; data[7] = _1_3; + data[8] = _2_0; data[9] = _2_1; data[10] = _2_2; data[11] = _2_3; + data[12] = _3_0; data[13] = _3_1; data[14] = _3_2; data[15] = _3_3; + } + + /// Constucts a 4-by-4 matrix from row vectors + MCTLASS_HOST_DEVICE + Matrix( + Matrix const &row_0, + Matrix const &row_1, + Matrix const &row_2, + Matrix const &row_3 + ) { + data[0] = row_0.data[0]; + data[1] = row_0.data[1]; + data[2] = row_0.data[2]; + data[3] = row_0.data[3]; + data[4] = row_1.data[0]; + data[5] = row_1.data[1]; + data[6] = row_1.data[2]; + data[7] = row_1.data[3]; + data[8] = row_2.data[0]; + data[9] = row_2.data[1]; + data[10] = row_2.data[2]; + data[11] = row_2.data[3]; + data[12] = row_3.data[0]; + data[13] = row_3.data[1]; + data[14] = row_3.data[2]; + data[15] = row_3.data[3]; + } + + /// Static method to construct a 4-by-4 matrix from column vectors + MCTLASS_HOST_DEVICE + static Matrix from_columns( + Matrix const &column_0, + Matrix const &column_1, + Matrix const &column_2, + Matrix const &column_3 + ) { + Matrix result; + + result.data[0] = column_0.data[0]; + result.data[1] = column_1.data[0]; + result.data[2] = column_2.data[0]; + result.data[3] = column_3.data[0]; + result.data[4] = column_0.data[1]; + result.data[5] = column_1.data[1]; + result.data[6] = column_2.data[1]; + result.data[7] = column_3.data[1]; + result.data[8] = column_0.data[2]; + result.data[9] = column_1.data[2]; + result.data[10] = column_2.data[2]; + result.data[11] = column_3.data[2]; + result.data[12] = column_0.data[3]; + result.data[13] = column_1.data[3]; + result.data[14] = column_2.data[3]; + result.data[15] = column_3.data[3]; + return result; + } + + /// Constructs an identity matrix + MCTLASS_HOST_DEVICE + static Matrix identity() { + Matrix m; + + m.data[0] = Element(1); + m.data[5] = Element(1); + m.data[10] = Element(1); + m.data[15] = Element(1); + + return m; + } + + /// Constructs a matrix from a uniform element + MCTLASS_HOST_DEVICE + static Matrix uniform(Element s) { + Matrix m; + + m.data[0] = s; + m.data[1] = s; + m.data[2] = s; + m.data[3] = s; + m.data[4] = s; + m.data[5] = s; + m.data[6] = s; + m.data[7] = s; + m.data[8] = s; + m.data[9] = s; + m.data[10] = s; + m.data[11] = s; + m.data[12] = s; + m.data[13] = s; + m.data[14] = s; + m.data[15] = s; + + return m; + } + + /// Constructs a matrix from a uniform element 1 + MCTLASS_HOST_DEVICE + static Matrix ones() { + return uniform(Element(1)); + } + + /// Constructs a matrix from a uniform element 0 + MCTLASS_HOST_DEVICE + static Matrix zero() { + return Matrix(); + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[5] = diag.data[1]; + m.data[10] = diag.data[2]; + m.data[15] = diag.data[3]; + + return m; + } + + /// Constructs a matrix from elements along its diagonal + MCTLASS_HOST_DEVICE + static Matrix from_diagonal(Matrix const &diag) { + Matrix m; + + m.data[0] = diag.data[0]; + m.data[5] = diag.data[1]; + m.data[10] = diag.data[2]; + m.data[15] = diag.data[3]; + + return m; + } + + /// Gets an array of diagonal elements + MCTLASS_HOST_DEVICE + Matrix diagonal() const { + Matrix diag; + + diag.data[0] = data[0]; + diag.data[1] = data[5]; + diag.data[2] = data[10]; + diag.data[3] = data[15]; + + return diag; + } + + /// Returns a transposed matrix + MCTLASS_HOST_DEVICE + Matrix transpose() const { + Matrix mt; + + mt.data[0] = data[0]; + mt.data[4] = data[1]; + mt.data[8] = data[2]; + mt.data[12] = data[3]; + mt.data[1] = data[4]; + mt.data[5] = data[5]; + mt.data[9] = data[6]; + mt.data[13] = data[7]; + mt.data[2] = data[8]; + mt.data[6] = data[9]; + mt.data[10] = data[10]; + mt.data[14] = data[11]; + mt.data[3] = data[12]; + mt.data[7] = data[13]; + mt.data[11] = data[14]; + mt.data[15] = data[15]; + + return mt; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(int i, int j) const { + return data[i * 4 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(int i, int j) { + return data[i * 4 + j]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element at(Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & at(Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element &at(int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element at(int offset) const { + return data[offset]; + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element operator[](Coord<2> const &coord) const { + return at(coord[0], coord[1]); + } + + /// Accesses an element by coordinate + MCTLASS_HOST_DEVICE + Element & operator[](Coord<2> const &coord) { + return at(coord[0], coord[1]); + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element & operator[](int offset) { + return data[offset]; + } + + /// Accesses an element by offset + MCTLASS_HOST_DEVICE + Element operator[](int offset) const { + return data[offset]; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_1x4(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 3]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_1x4(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 3] = m.data[3]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix row(int i) const { + return slice_1x4(i, 0); + } + + MCTLASS_HOST_DEVICE + Matrix &set_row(Matrix const &v, int i = 0) { + return set_slice_1x4(v, i, 0); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 4]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 4] = m.data[1]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 4]; + m.data[3] = data[i * 4 + j + 5]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 4] = m.data[2]; + data[i * 4 + j + 5] = m.data[3]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 4]; + m.data[4] = data[i * 4 + j + 5]; + m.data[5] = data[i * 4 + j + 6]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 4] = m.data[3]; + data[i * 4 + j + 5] = m.data[4]; + data[i * 4 + j + 6] = m.data[5]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_2x4(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 3]; + m.data[4] = data[i * 4 + j + 4]; + m.data[5] = data[i * 4 + j + 5]; + m.data[6] = data[i * 4 + j + 6]; + m.data[7] = data[i * 4 + j + 7]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_2x4(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 3] = m.data[3]; + data[i * 4 + j + 4] = m.data[4]; + data[i * 4 + j + 5] = m.data[5]; + data[i * 4 + j + 6] = m.data[6]; + data[i * 4 + j + 7] = m.data[7]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 4]; + m.data[2] = data[i * 4 + j + 8]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 4] = m.data[1]; + data[i * 4 + j + 8] = m.data[2]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 4]; + m.data[3] = data[i * 4 + j + 5]; + m.data[4] = data[i * 4 + j + 8]; + m.data[5] = data[i * 4 + j + 9]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 4] = m.data[2]; + data[i * 4 + j + 5] = m.data[3]; + data[i * 4 + j + 8] = m.data[4]; + data[i * 4 + j + 9] = m.data[5]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 4]; + m.data[4] = data[i * 4 + j + 5]; + m.data[5] = data[i * 4 + j + 6]; + m.data[6] = data[i * 4 + j + 8]; + m.data[7] = data[i * 4 + j + 9]; + m.data[8] = data[i * 4 + j + 10]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 4] = m.data[3]; + data[i * 4 + j + 5] = m.data[4]; + data[i * 4 + j + 6] = m.data[5]; + data[i * 4 + j + 8] = m.data[6]; + data[i * 4 + j + 9] = m.data[7]; + data[i * 4 + j + 10] = m.data[8]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_3x4(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 3]; + m.data[4] = data[i * 4 + j + 4]; + m.data[5] = data[i * 4 + j + 5]; + m.data[6] = data[i * 4 + j + 6]; + m.data[7] = data[i * 4 + j + 7]; + m.data[8] = data[i * 4 + j + 8]; + m.data[9] = data[i * 4 + j + 9]; + m.data[10] = data[i * 4 + j + 10]; + m.data[11] = data[i * 4 + j + 11]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_3x4(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 3] = m.data[3]; + data[i * 4 + j + 4] = m.data[4]; + data[i * 4 + j + 5] = m.data[5]; + data[i * 4 + j + 6] = m.data[6]; + data[i * 4 + j + 7] = m.data[7]; + data[i * 4 + j + 8] = m.data[8]; + data[i * 4 + j + 9] = m.data[9]; + data[i * 4 + j + 10] = m.data[10]; + data[i * 4 + j + 11] = m.data[11]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_4x1(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 4]; + m.data[2] = data[i * 4 + j + 8]; + m.data[3] = data[i * 4 + j + 12]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_4x1(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 4] = m.data[1]; + data[i * 4 + j + 8] = m.data[2]; + data[i * 4 + j + 12] = m.data[3]; + + return *this; + } + + MCTLASS_HOST_DEVICE + Matrix column(int j) const { + return slice_4x1(0, j); + } + + MCTLASS_HOST_DEVICE + Matrix &set_column(Matrix const &v, int j =0) { + return set_slice_4x1(v, 0, j); + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_4x2(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 4]; + m.data[3] = data[i * 4 + j + 5]; + m.data[4] = data[i * 4 + j + 8]; + m.data[5] = data[i * 4 + j + 9]; + m.data[6] = data[i * 4 + j + 12]; + m.data[7] = data[i * 4 + j + 13]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_4x2(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 4] = m.data[2]; + data[i * 4 + j + 5] = m.data[3]; + data[i * 4 + j + 8] = m.data[4]; + data[i * 4 + j + 9] = m.data[5]; + data[i * 4 + j + 12] = m.data[6]; + data[i * 4 + j + 13] = m.data[7]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_4x3(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 4]; + m.data[4] = data[i * 4 + j + 5]; + m.data[5] = data[i * 4 + j + 6]; + m.data[6] = data[i * 4 + j + 8]; + m.data[7] = data[i * 4 + j + 9]; + m.data[8] = data[i * 4 + j + 10]; + m.data[9] = data[i * 4 + j + 12]; + m.data[10] = data[i * 4 + j + 13]; + m.data[11] = data[i * 4 + j + 14]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_4x3(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 4] = m.data[3]; + data[i * 4 + j + 5] = m.data[4]; + data[i * 4 + j + 6] = m.data[5]; + data[i * 4 + j + 8] = m.data[6]; + data[i * 4 + j + 9] = m.data[7]; + data[i * 4 + j + 10] = m.data[8]; + data[i * 4 + j + 12] = m.data[9]; + data[i * 4 + j + 13] = m.data[10]; + data[i * 4 + j + 14] = m.data[11]; + + return *this; + } + + /// Gets a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix slice_4x4(int i = 0, int j = 0) const { + Matrix m; + + m.data[0] = data[i * 4 + j + 0]; + m.data[1] = data[i * 4 + j + 1]; + m.data[2] = data[i * 4 + j + 2]; + m.data[3] = data[i * 4 + j + 3]; + m.data[4] = data[i * 4 + j + 4]; + m.data[5] = data[i * 4 + j + 5]; + m.data[6] = data[i * 4 + j + 6]; + m.data[7] = data[i * 4 + j + 7]; + m.data[8] = data[i * 4 + j + 8]; + m.data[9] = data[i * 4 + j + 9]; + m.data[10] = data[i * 4 + j + 10]; + m.data[11] = data[i * 4 + j + 11]; + m.data[12] = data[i * 4 + j + 12]; + m.data[13] = data[i * 4 + j + 13]; + m.data[14] = data[i * 4 + j + 14]; + m.data[15] = data[i * 4 + j + 15]; + + return m; + } + + /// Overwrites a submatrix with optional offset + MCTLASS_HOST_DEVICE + Matrix & set_slice_4x4(Matrix const &m, int i = 0, int j = 0) { + + data[i * 4 + j + 0] = m.data[0]; + data[i * 4 + j + 1] = m.data[1]; + data[i * 4 + j + 2] = m.data[2]; + data[i * 4 + j + 3] = m.data[3]; + data[i * 4 + j + 4] = m.data[4]; + data[i * 4 + j + 5] = m.data[5]; + data[i * 4 + j + 6] = m.data[6]; + data[i * 4 + j + 7] = m.data[7]; + data[i * 4 + j + 8] = m.data[8]; + data[i * 4 + j + 9] = m.data[9]; + data[i * 4 + j + 10] = m.data[10]; + data[i * 4 + j + 11] = m.data[11]; + data[i * 4 + j + 12] = m.data[12]; + data[i * 4 + j + 13] = m.data[13]; + data[i * 4 + j + 14] = m.data[14]; + data[i * 4 + j + 15] = m.data[15]; + + return *this; + } + + /// Forms a 4-by-4 matrix by horizontally concatenating a 4-by-1 matrix with a 4-by-3 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), rhs.at(0, 0), rhs.at(0, 1), rhs.at(0, 2) + , lhs.at(1, 0), rhs.at(1, 0), rhs.at(1, 1), rhs.at(1, 2) + , lhs.at(2, 0), rhs.at(2, 0), rhs.at(2, 1), rhs.at(2, 2) + , lhs.at(3, 0), rhs.at(3, 0), rhs.at(3, 1), rhs.at(3, 2)); + } + + /// Forms a 4-by-4 matrix by horizontally concatenating a 4-by-2 matrix with a 4-by-2 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0), rhs.at(0, 1) + , lhs.at(1, 0), lhs.at(1, 1), rhs.at(1, 0), rhs.at(1, 1) + , lhs.at(2, 0), lhs.at(2, 1), rhs.at(2, 0), rhs.at(2, 1) + , lhs.at(3, 0), lhs.at(3, 1), rhs.at(3, 0), rhs.at(3, 1)); + } + + /// Forms a 4-by-4 matrix by horizontally concatenating a 4-by-3 matrix with a 4-by-1 matrix + MCTLASS_HOST_DEVICE + static Matrix hcat(Matrix const & lhs, Matrix const & rhs) { + return Matrix( + lhs.at(0, 0), lhs.at(0, 1), lhs.at(0, 2), rhs.at(0, 0) + , lhs.at(1, 0), lhs.at(1, 1), lhs.at(1, 2), rhs.at(1, 0) + , lhs.at(2, 0), lhs.at(2, 1), lhs.at(2, 2), rhs.at(2, 0) + , lhs.at(3, 0), lhs.at(3, 1), lhs.at(3, 2), rhs.at(3, 0)); + } + + /// Forms a 4-by-4 matrix by vertically concatenating a 1-by-4 matrix with a 3-by-4 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1), upper.at(0, 2), upper.at(0, 3) + , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2), lower.at(0, 3) + , lower.at(1, 0), lower.at(1, 1), lower.at(1, 2), lower.at(1, 3) + , lower.at(2, 0), lower.at(2, 1), lower.at(2, 2), lower.at(2, 3)); + } + + /// Forms a 4-by-4 matrix by vertically concatenating a 2-by-4 matrix with a 2-by-4 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1), upper.at(0, 2), upper.at(0, 3) + , upper.at(1, 0), upper.at(1, 1), upper.at(1, 2), upper.at(1, 3) + , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2), lower.at(0, 3) + , lower.at(1, 0), lower.at(1, 1), lower.at(1, 2), lower.at(1, 3)); + } + + /// Forms a 4-by-4 matrix by vertically concatenating a 3-by-4 matrix with a 1-by-4 matrix + MCTLASS_HOST_DEVICE + static Matrix vcat(Matrix const & upper, Matrix const & lower) { + return Matrix( + upper.at(0, 0), upper.at(0, 1), upper.at(0, 2), upper.at(0, 3) + , upper.at(1, 0), upper.at(1, 1), upper.at(1, 2), upper.at(1, 3) + , upper.at(2, 0), upper.at(2, 1), upper.at(2, 2), upper.at(2, 3) + , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2), lower.at(0, 3)); + } + + /// Forms a 4-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Element A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A, B.at(0, 0), B.at(0, 1), B.at(0, 2) + , C.at(0, 0), D.at(0, 0), D.at(0, 1), D.at(0, 2) + , C.at(1, 0), D.at(1, 0), D.at(1, 1), D.at(1, 2) + , C.at(2, 0), D.at(2, 0), D.at(2, 1), D.at(2, 2) + ); + } + + /// Forms a 4-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), A.at(0, 1), B.at(0, 0), B.at(0, 1) + , C.at(0, 0), C.at(0, 1), D.at(0, 0), D.at(0, 1) + , C.at(1, 0), C.at(1, 1), D.at(1, 0), D.at(1, 1) + , C.at(2, 0), C.at(2, 1), D.at(2, 0), D.at(2, 1) + ); + } + + /// Forms a 4-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Element B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), A.at(0, 1), A.at(0, 2), B + , C.at(0, 0), C.at(0, 1), C.at(0, 2), D.at(0, 0) + , C.at(1, 0), C.at(1, 1), C.at(1, 2), D.at(1, 0) + , C.at(2, 0), C.at(2, 1), C.at(2, 2), D.at(2, 0) + ); + } + + /// Forms a 4-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), B.at(0, 0), B.at(0, 1), B.at(0, 2) + , A.at(1, 0), B.at(1, 0), B.at(1, 1), B.at(1, 2) + , C.at(0, 0), D.at(0, 0), D.at(0, 1), D.at(0, 2) + , C.at(1, 0), D.at(1, 0), D.at(1, 1), D.at(1, 2) + ); + } + + /// Forms a 4-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), A.at(0, 1), B.at(0, 0), B.at(0, 1) + , A.at(1, 0), A.at(1, 1), B.at(1, 0), B.at(1, 1) + , C.at(0, 0), C.at(0, 1), D.at(0, 0), D.at(0, 1) + , C.at(1, 0), C.at(1, 1), D.at(1, 0), D.at(1, 1) + ); + } + + /// Forms a 4-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), A.at(0, 1), A.at(0, 2), B.at(0, 0) + , A.at(1, 0), A.at(1, 1), A.at(1, 2), B.at(1, 0) + , C.at(0, 0), C.at(0, 1), C.at(0, 2), D.at(0, 0) + , C.at(1, 0), C.at(1, 1), C.at(1, 2), D.at(1, 0) + ); + } + + /// Forms a 4-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Element C, Matrix const & D) { + return Matrix( + A.at(0, 0), B.at(0, 0), B.at(0, 1), B.at(0, 2) + , A.at(1, 0), B.at(1, 0), B.at(1, 1), B.at(1, 2) + , A.at(2, 0), B.at(2, 0), B.at(2, 1), B.at(2, 2) + , C, D.at(0, 0), D.at(0, 1), D.at(0, 2) + ); + } + + /// Forms a 4-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Matrix const & D) { + return Matrix( + A.at(0, 0), A.at(0, 1), B.at(0, 0), B.at(0, 1) + , A.at(1, 0), A.at(1, 1), B.at(1, 0), B.at(1, 1) + , A.at(2, 0), A.at(2, 1), B.at(2, 0), B.at(2, 1) + , C.at(0, 0), C.at(0, 1), D.at(0, 0), D.at(0, 1) + ); + } + + /// Forms a 4-by-4 matrix by concatenating four components + MCTLASS_HOST_DEVICE + static Matrix block( + Matrix const & A, Matrix const & B, + Matrix const & C, Element D) { + return Matrix( + A.at(0, 0), A.at(0, 1), A.at(0, 2), B.at(0, 0) + , A.at(1, 0), A.at(1, 1), A.at(1, 2), B.at(1, 0) + , A.at(2, 0), A.at(2, 1), A.at(2, 2), B.at(2, 0) + , C.at(0, 0), C.at(0, 1), C.at(0, 2), D + ); + } + + /// Elementwise add operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix add(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] + rhs.data[0]; + result.data[1] = data[1] + rhs.data[1]; + result.data[2] = data[2] + rhs.data[2]; + result.data[3] = data[3] + rhs.data[3]; + + result.data[4] = data[4] + rhs.data[4]; + result.data[5] = data[5] + rhs.data[5]; + result.data[6] = data[6] + rhs.data[6]; + result.data[7] = data[7] + rhs.data[7]; + + result.data[8] = data[8] + rhs.data[8]; + result.data[9] = data[9] + rhs.data[9]; + result.data[10] = data[10] + rhs.data[10]; + result.data[11] = data[11] + rhs.data[11]; + + result.data[12] = data[12] + rhs.data[12]; + result.data[13] = data[13] + rhs.data[13]; + result.data[14] = data[14] + rhs.data[14]; + result.data[15] = data[15] + rhs.data[15]; + + return result; + } + + /// Elementwise add operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix operator +(Matrix const &rhs) const { + return add(rhs); + } + + /// Elementwise add operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator +=(Matrix const &rhs) { + + data[0] += rhs.data[0]; + data[1] += rhs.data[1]; + data[2] += rhs.data[2]; + data[3] += rhs.data[3]; + + data[4] += rhs.data[4]; + data[5] += rhs.data[5]; + data[6] += rhs.data[6]; + data[7] += rhs.data[7]; + + data[8] += rhs.data[8]; + data[9] += rhs.data[9]; + data[10] += rhs.data[10]; + data[11] += rhs.data[11]; + + data[12] += rhs.data[12]; + data[13] += rhs.data[13]; + data[14] += rhs.data[14]; + data[15] += rhs.data[15]; + + return *this; + } + + /// Elementwise subtract operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix subtract(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] - rhs.data[0]; + result.data[1] = data[1] - rhs.data[1]; + result.data[2] = data[2] - rhs.data[2]; + result.data[3] = data[3] - rhs.data[3]; + + result.data[4] = data[4] - rhs.data[4]; + result.data[5] = data[5] - rhs.data[5]; + result.data[6] = data[6] - rhs.data[6]; + result.data[7] = data[7] - rhs.data[7]; + + result.data[8] = data[8] - rhs.data[8]; + result.data[9] = data[9] - rhs.data[9]; + result.data[10] = data[10] - rhs.data[10]; + result.data[11] = data[11] - rhs.data[11]; + + result.data[12] = data[12] - rhs.data[12]; + result.data[13] = data[13] - rhs.data[13]; + result.data[14] = data[14] - rhs.data[14]; + result.data[15] = data[15] - rhs.data[15]; + + return result; + } + + /// Elementwise subtract operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix operator -(Matrix const &rhs) const { + return subtract(rhs); + } + + /// Elementwise subtract operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator -=(Matrix const &rhs) { + + data[0] -= rhs.data[0]; + data[1] -= rhs.data[1]; + data[2] -= rhs.data[2]; + data[3] -= rhs.data[3]; + + data[4] -= rhs.data[4]; + data[5] -= rhs.data[5]; + data[6] -= rhs.data[6]; + data[7] -= rhs.data[7]; + + data[8] -= rhs.data[8]; + data[9] -= rhs.data[9]; + data[10] -= rhs.data[10]; + data[11] -= rhs.data[11]; + + data[12] -= rhs.data[12]; + data[13] -= rhs.data[13]; + data[14] -= rhs.data[14]; + data[15] -= rhs.data[15]; + + return *this; + } + + /// Elementwise multiply operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix multiply(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] * rhs.data[0]; + result.data[1] = data[1] * rhs.data[1]; + result.data[2] = data[2] * rhs.data[2]; + result.data[3] = data[3] * rhs.data[3]; + + result.data[4] = data[4] * rhs.data[4]; + result.data[5] = data[5] * rhs.data[5]; + result.data[6] = data[6] * rhs.data[6]; + result.data[7] = data[7] * rhs.data[7]; + + result.data[8] = data[8] * rhs.data[8]; + result.data[9] = data[9] * rhs.data[9]; + result.data[10] = data[10] * rhs.data[10]; + result.data[11] = data[11] * rhs.data[11]; + + result.data[12] = data[12] * rhs.data[12]; + result.data[13] = data[13] * rhs.data[13]; + result.data[14] = data[14] * rhs.data[14]; + result.data[15] = data[15] * rhs.data[15]; + + return result; + } + + /// Scalar multiply operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix multiply(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] * s; + result.data[1] = data[1] * s; + result.data[2] = data[2] * s; + result.data[3] = data[3] * s; + + result.data[4] = data[4] * s; + result.data[5] = data[5] * s; + result.data[6] = data[6] * s; + result.data[7] = data[7] * s; + + result.data[8] = data[8] * s; + result.data[9] = data[9] * s; + result.data[10] = data[10] * s; + result.data[11] = data[11] * s; + + result.data[12] = data[12] * s; + result.data[13] = data[13] * s; + result.data[14] = data[14] * s; + result.data[15] = data[15] * s; + + return result; + } + + /// Scalar multiply operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix operator *(Element const &s) const { + return multiply(s); + } + + /// Scalar multiply operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator *=(Element const &s) { + + data[0] *= s; + data[1] *= s; + data[2] *= s; + data[3] *= s; + + data[4] *= s; + data[5] *= s; + data[6] *= s; + data[7] *= s; + + data[8] *= s; + data[9] *= s; + data[10] *= s; + data[11] *= s; + + data[12] *= s; + data[13] *= s; + data[14] *= s; + data[15] *= s; + + return *this; + } + + /// Elementwise divide operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix divide(Matrix const &rhs) const { + + Matrix result; + + result.data[0] = data[0] / rhs.data[0]; + result.data[1] = data[1] / rhs.data[1]; + result.data[2] = data[2] / rhs.data[2]; + result.data[3] = data[3] / rhs.data[3]; + + result.data[4] = data[4] / rhs.data[4]; + result.data[5] = data[5] / rhs.data[5]; + result.data[6] = data[6] / rhs.data[6]; + result.data[7] = data[7] / rhs.data[7]; + + result.data[8] = data[8] / rhs.data[8]; + result.data[9] = data[9] / rhs.data[9]; + result.data[10] = data[10] / rhs.data[10]; + result.data[11] = data[11] / rhs.data[11]; + + result.data[12] = data[12] / rhs.data[12]; + result.data[13] = data[13] / rhs.data[13]; + result.data[14] = data[14] / rhs.data[14]; + result.data[15] = data[15] / rhs.data[15]; + + return result; + } + + /// Scalar divide operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix divide(Element const &s) const { + + Matrix result; + + result.data[0] = data[0] / s; + result.data[1] = data[1] / s; + result.data[2] = data[2] / s; + result.data[3] = data[3] / s; + + result.data[4] = data[4] / s; + result.data[5] = data[5] / s; + result.data[6] = data[6] / s; + result.data[7] = data[7] / s; + + result.data[8] = data[8] / s; + result.data[9] = data[9] / s; + result.data[10] = data[10] / s; + result.data[11] = data[11] / s; + + result.data[12] = data[12] / s; + result.data[13] = data[13] / s; + result.data[14] = data[14] / s; + result.data[15] = data[15] / s; + + return result; + } + + /// Scalar divide operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix operator /(Element const &s) const { + return divide(s); + } + + /// Scalar divide operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Element const &s) { + + data[0] /= s; + data[1] /= s; + data[2] /= s; + data[3] /= s; + + data[4] /= s; + data[5] /= s; + data[6] /= s; + data[7] /= s; + + data[8] /= s; + data[9] /= s; + data[10] /= s; + data[11] /= s; + + data[12] /= s; + data[13] /= s; + data[14] /= s; + data[15] /= s; + + return *this; + } + + /// Elementwise divide operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix operator /(Matrix const &rhs) const { + return divide(rhs); + } + + /// Elementwise divide operator (4-by-4) + MCTLASS_HOST_DEVICE + Matrix & operator /=(Matrix const &rhs) { + + data[0] /= rhs.data[0]; + data[1] /= rhs.data[1]; + data[2] /= rhs.data[2]; + data[3] /= rhs.data[3]; + + data[4] /= rhs.data[4]; + data[5] /= rhs.data[5]; + data[6] /= rhs.data[6]; + data[7] /= rhs.data[7]; + + data[8] /= rhs.data[8]; + data[9] /= rhs.data[9]; + data[10] /= rhs.data[10]; + data[11] /= rhs.data[11]; + + data[12] /= rhs.data[12]; + data[13] /= rhs.data[13]; + data[14] /= rhs.data[14]; + data[15] /= rhs.data[15]; + + return *this; + } + + /// Negates each element of the matrix + MCTLASS_HOST_DEVICE + Matrix operator-() const { + Matrix m; + + m.data[0] = -m.data[0]; + m.data[1] = -m.data[1]; + m.data[2] = -m.data[2]; + m.data[3] = -m.data[3]; + m.data[4] = -m.data[4]; + m.data[5] = -m.data[5]; + m.data[6] = -m.data[6]; + m.data[7] = -m.data[7]; + m.data[8] = -m.data[8]; + m.data[9] = -m.data[9]; + m.data[10] = -m.data[10]; + m.data[11] = -m.data[11]; + m.data[12] = -m.data[12]; + m.data[13] = -m.data[13]; + m.data[14] = -m.data[14]; + m.data[15] = -m.data[15]; + + return m; + } + + /// Matrix product of size 4-by-1-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[4] * rhs.data[0]; + accum.data[2] += data[8] * rhs.data[0]; + accum.data[3] += data[12] * rhs.data[0]; + + // k=1 + accum.data[0] += data[1] * rhs.data[1]; + accum.data[1] += data[5] * rhs.data[1]; + accum.data[2] += data[9] * rhs.data[1]; + accum.data[3] += data[13] * rhs.data[1]; + + // k=2 + accum.data[0] += data[2] * rhs.data[2]; + accum.data[1] += data[6] * rhs.data[2]; + accum.data[2] += data[10] * rhs.data[2]; + accum.data[3] += data[14] * rhs.data[2]; + + // k=3 + accum.data[0] += data[3] * rhs.data[3]; + accum.data[1] += data[7] * rhs.data[3]; + accum.data[2] += data[11] * rhs.data[3]; + accum.data[3] += data[15] * rhs.data[3]; + + return accum; + } + + /// Matrix product of size 4-by-1-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 4-by-2-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[4] * rhs.data[0]; + accum.data[3] += data[4] * rhs.data[1]; + accum.data[4] += data[8] * rhs.data[0]; + accum.data[5] += data[8] * rhs.data[1]; + accum.data[6] += data[12] * rhs.data[0]; + accum.data[7] += data[12] * rhs.data[1]; + + // k=1 + accum.data[0] += data[1] * rhs.data[2]; + accum.data[1] += data[1] * rhs.data[3]; + accum.data[2] += data[5] * rhs.data[2]; + accum.data[3] += data[5] * rhs.data[3]; + accum.data[4] += data[9] * rhs.data[2]; + accum.data[5] += data[9] * rhs.data[3]; + accum.data[6] += data[13] * rhs.data[2]; + accum.data[7] += data[13] * rhs.data[3]; + + // k=2 + accum.data[0] += data[2] * rhs.data[4]; + accum.data[1] += data[2] * rhs.data[5]; + accum.data[2] += data[6] * rhs.data[4]; + accum.data[3] += data[6] * rhs.data[5]; + accum.data[4] += data[10] * rhs.data[4]; + accum.data[5] += data[10] * rhs.data[5]; + accum.data[6] += data[14] * rhs.data[4]; + accum.data[7] += data[14] * rhs.data[5]; + + // k=3 + accum.data[0] += data[3] * rhs.data[6]; + accum.data[1] += data[3] * rhs.data[7]; + accum.data[2] += data[7] * rhs.data[6]; + accum.data[3] += data[7] * rhs.data[7]; + accum.data[4] += data[11] * rhs.data[6]; + accum.data[5] += data[11] * rhs.data[7]; + accum.data[6] += data[15] * rhs.data[6]; + accum.data[7] += data[15] * rhs.data[7]; + + return accum; + } + + /// Matrix product of size 4-by-2-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 4-by-3-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[4] * rhs.data[0]; + accum.data[4] += data[4] * rhs.data[1]; + accum.data[5] += data[4] * rhs.data[2]; + accum.data[6] += data[8] * rhs.data[0]; + accum.data[7] += data[8] * rhs.data[1]; + accum.data[8] += data[8] * rhs.data[2]; + accum.data[9] += data[12] * rhs.data[0]; + accum.data[10] += data[12] * rhs.data[1]; + accum.data[11] += data[12] * rhs.data[2]; + + // k=1 + accum.data[0] += data[1] * rhs.data[3]; + accum.data[1] += data[1] * rhs.data[4]; + accum.data[2] += data[1] * rhs.data[5]; + accum.data[3] += data[5] * rhs.data[3]; + accum.data[4] += data[5] * rhs.data[4]; + accum.data[5] += data[5] * rhs.data[5]; + accum.data[6] += data[9] * rhs.data[3]; + accum.data[7] += data[9] * rhs.data[4]; + accum.data[8] += data[9] * rhs.data[5]; + accum.data[9] += data[13] * rhs.data[3]; + accum.data[10] += data[13] * rhs.data[4]; + accum.data[11] += data[13] * rhs.data[5]; + + // k=2 + accum.data[0] += data[2] * rhs.data[6]; + accum.data[1] += data[2] * rhs.data[7]; + accum.data[2] += data[2] * rhs.data[8]; + accum.data[3] += data[6] * rhs.data[6]; + accum.data[4] += data[6] * rhs.data[7]; + accum.data[5] += data[6] * rhs.data[8]; + accum.data[6] += data[10] * rhs.data[6]; + accum.data[7] += data[10] * rhs.data[7]; + accum.data[8] += data[10] * rhs.data[8]; + accum.data[9] += data[14] * rhs.data[6]; + accum.data[10] += data[14] * rhs.data[7]; + accum.data[11] += data[14] * rhs.data[8]; + + // k=3 + accum.data[0] += data[3] * rhs.data[9]; + accum.data[1] += data[3] * rhs.data[10]; + accum.data[2] += data[3] * rhs.data[11]; + accum.data[3] += data[7] * rhs.data[9]; + accum.data[4] += data[7] * rhs.data[10]; + accum.data[5] += data[7] * rhs.data[11]; + accum.data[6] += data[11] * rhs.data[9]; + accum.data[7] += data[11] * rhs.data[10]; + accum.data[8] += data[11] * rhs.data[11]; + accum.data[9] += data[15] * rhs.data[9]; + accum.data[10] += data[15] * rhs.data[10]; + accum.data[11] += data[15] * rhs.data[11]; + + return accum; + } + + /// Matrix product of size 4-by-3-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 4-by-4-by-4 + MCTLASS_HOST_DEVICE + Matrix product( + Matrix const &rhs, + Matrix accum = Matrix() + ) const { + + // k=0 + accum.data[0] += data[0] * rhs.data[0]; + accum.data[1] += data[0] * rhs.data[1]; + accum.data[2] += data[0] * rhs.data[2]; + accum.data[3] += data[0] * rhs.data[3]; + accum.data[4] += data[4] * rhs.data[0]; + accum.data[5] += data[4] * rhs.data[1]; + accum.data[6] += data[4] * rhs.data[2]; + accum.data[7] += data[4] * rhs.data[3]; + accum.data[8] += data[8] * rhs.data[0]; + accum.data[9] += data[8] * rhs.data[1]; + accum.data[10] += data[8] * rhs.data[2]; + accum.data[11] += data[8] * rhs.data[3]; + accum.data[12] += data[12] * rhs.data[0]; + accum.data[13] += data[12] * rhs.data[1]; + accum.data[14] += data[12] * rhs.data[2]; + accum.data[15] += data[12] * rhs.data[3]; + + // k=1 + accum.data[0] += data[1] * rhs.data[4]; + accum.data[1] += data[1] * rhs.data[5]; + accum.data[2] += data[1] * rhs.data[6]; + accum.data[3] += data[1] * rhs.data[7]; + accum.data[4] += data[5] * rhs.data[4]; + accum.data[5] += data[5] * rhs.data[5]; + accum.data[6] += data[5] * rhs.data[6]; + accum.data[7] += data[5] * rhs.data[7]; + accum.data[8] += data[9] * rhs.data[4]; + accum.data[9] += data[9] * rhs.data[5]; + accum.data[10] += data[9] * rhs.data[6]; + accum.data[11] += data[9] * rhs.data[7]; + accum.data[12] += data[13] * rhs.data[4]; + accum.data[13] += data[13] * rhs.data[5]; + accum.data[14] += data[13] * rhs.data[6]; + accum.data[15] += data[13] * rhs.data[7]; + + // k=2 + accum.data[0] += data[2] * rhs.data[8]; + accum.data[1] += data[2] * rhs.data[9]; + accum.data[2] += data[2] * rhs.data[10]; + accum.data[3] += data[2] * rhs.data[11]; + accum.data[4] += data[6] * rhs.data[8]; + accum.data[5] += data[6] * rhs.data[9]; + accum.data[6] += data[6] * rhs.data[10]; + accum.data[7] += data[6] * rhs.data[11]; + accum.data[8] += data[10] * rhs.data[8]; + accum.data[9] += data[10] * rhs.data[9]; + accum.data[10] += data[10] * rhs.data[10]; + accum.data[11] += data[10] * rhs.data[11]; + accum.data[12] += data[14] * rhs.data[8]; + accum.data[13] += data[14] * rhs.data[9]; + accum.data[14] += data[14] * rhs.data[10]; + accum.data[15] += data[14] * rhs.data[11]; + + // k=3 + accum.data[0] += data[3] * rhs.data[12]; + accum.data[1] += data[3] * rhs.data[13]; + accum.data[2] += data[3] * rhs.data[14]; + accum.data[3] += data[3] * rhs.data[15]; + accum.data[4] += data[7] * rhs.data[12]; + accum.data[5] += data[7] * rhs.data[13]; + accum.data[6] += data[7] * rhs.data[14]; + accum.data[7] += data[7] * rhs.data[15]; + accum.data[8] += data[11] * rhs.data[12]; + accum.data[9] += data[11] * rhs.data[13]; + accum.data[10] += data[11] * rhs.data[14]; + accum.data[11] += data[11] * rhs.data[15]; + accum.data[12] += data[15] * rhs.data[12]; + accum.data[13] += data[15] * rhs.data[13]; + accum.data[14] += data[15] * rhs.data[14]; + accum.data[15] += data[15] * rhs.data[15]; + + return accum; + } + + /// Matrix product of size 4-by-4-by-4 + MCTLASS_HOST_DEVICE + Matrix operator*(Matrix const &rhs) const { + return product(rhs); + } + + /// Matrix product of size 4-by-4-by-4 + MCTLASS_HOST_DEVICE + Matrix & operator*=(Matrix const &rhs) { + *this = product(rhs); + return *this; + } + + /// Returns the sum of elements + MCTLASS_HOST_DEVICE + Element sum(Element accum = Element()) const { + + accum += data[0]; + accum += data[1]; + accum += data[2]; + accum += data[3]; + accum += data[4]; + accum += data[5]; + accum += data[6]; + accum += data[7]; + accum += data[8]; + accum += data[9]; + accum += data[10]; + accum += data[11]; + accum += data[12]; + accum += data[13]; + accum += data[14]; + accum += data[15]; + + return accum; + } + + /// Returns the sum of squared elements + MCTLASS_HOST_DEVICE + Element norm(Element accum = Element()) const { + + accum += data[0] * data[0]; + accum += data[1] * data[1]; + accum += data[2] * data[2]; + accum += data[3] * data[3]; + accum += data[4] * data[4]; + accum += data[5] * data[5]; + accum += data[6] * data[6]; + accum += data[7] * data[7]; + accum += data[8] * data[8]; + accum += data[9] * data[9]; + accum += data[10] * data[10]; + accum += data[11] * data[11]; + accum += data[12] * data[12]; + accum += data[13] * data[13]; + accum += data[14] * data[14]; + accum += data[15] * data[15]; + + return accum; + } + + /// Returns square root of the norm + MCTLASS_HOST_DEVICE + Element magnitude() const { + return fast_sqrt(norm()); + } + + /// Returns the sum of diagonal elements + MCTLASS_HOST_DEVICE + Element trace(Element accum = Element()) const { + + accum += data[0]; + accum += data[5]; + accum += data[10]; + accum += data[15]; + + return accum; + } + + /// Returns 4-by-4 rotation matrix around the X axis + MCTLASS_HOST_DEVICE + static Matrix rotation_X(Element theta) { + Matrix m = identity(); + + Element c = fast_cos(theta); + Element s = fast_sin(theta); + + m.at(1, 1) = c; + m.at(1, 2) = -s; + m.at(2, 1) = s; + m.at(2, 2) = c; + + return m; + } + + /// Returns 4-by-4 rotation matrix around the Y axis + MCTLASS_HOST_DEVICE + static Matrix rotation_Y(Element theta) { + Matrix m = identity(); + + Element c = fast_cos(theta); + Element s = fast_sin(theta); + + m.at(0, 0) = c; + m.at(2, 0) = -s; + m.at(0, 2) = s; + m.at(2, 2) = c; + + return m; + } + + /// Returns 4-by-4 rotation matrix around the Z axis + MCTLASS_HOST_DEVICE + static Matrix rotation_Z(Element theta) { + Matrix m = Matrix::identity(); + + Element c = fast_cos(theta); + Element s = fast_sin(theta); + + m.at(0, 0) = c; + m.at(0, 1) = -s; + m.at(1, 0) = s; + m.at(1, 1) = c; + + return m; + } + + /// Returns a 4-by-4 rotation matrix around a unit-length axis + MCTLASS_HOST_DEVICE + static Matrix rotation(Element theta, Matrix const &u) { + Element x = u.data[0]; + Element y = u.data[1]; + Element z = u.data[2]; + + Element c = fast_cos(theta); + Element s = fast_sin(theta); + + Element one_minus_cos = Element(1) - fast_cos(theta); + + Matrix m; + + m.set_slice_3x3({ + c + x * x * one_minus_cos, x * y * one_minus_cos - z * s, x * z * one_minus_cos + y * s, + y * x * one_minus_cos * z * s, c + y * y * one_minus_cos, y * z * one_minus_cos - x * s, + z * x * one_minus_cos - y * s, z * y * one_minus_cos + x * s, c + z * z * one_minus_cos + }); + + return m; + } + + /// Returns a 4-by-4 reflection about the plane specified by the + /// unit-length normal vector n_unit + MCTLASS_HOST_DEVICE + static Matrix reflection(Matrix const &n_unit) { + + Element a = n_unit.data[0]; + Element b = n_unit.data[1]; + Element c = n_unit.data[2]; + + Matrix m = Matrix::identity(); + + m.set_slice_3x3({ + Element(1) - Element(2) * a * a, Element(-2) * a * b, Element(-2) * a * c, + Element(-2) * a * b, Element(1) - Element(2) * b * b, Element(-2) * b * c, + Element(-2) * a * c, Element(-2) * b * c, Element(1) - Element(2) * c * c + }); + + return m; + } + + /// Returns a perspective projection matrix typical of OpenGL applications + MCTLASS_HOST_DEVICE + static Matrix perspective(Element near_plane, Element far_plane, Element fovH, Element fovV) { + Element aspect = fovH / fovV; + Element f = Element(cos(fovV)) / Element(fovH); + Element Q = near_plane - far_plane; + + return Matrix( + f / aspect, 0, 0, 0, + 0, f, 0, 0, + 0, 0, (near_plane + far_plane) / Q, Element(2) * far_plane * near_plane / Q, + 0, 0, -1, 0 + ); + } + + MCTLASS_HOST_DEVICE + static Matrix translation(Matrix const &v) { + return Matrix( + 1, 0, 0, v.data[0], + 0, 1, 0, v.data[1], + 0, 0, 1, v.data[2], + 0, 0, 0, 1 + ); + } + + /// Computes the determinant of a 4-by-4 matrix + MCTLASS_HOST_DEVICE + Element determinant(Element accum = Element()) const { + + accum += at(0, 0) * Matrix({ at(1, 1), at(1, 2), at(1, 3), at(2, 1), at(2, 2), at(2, 3), at(3, 1), at(3, 2), at(3, 3) }).determinant(); + accum -= at(0, 1) * Matrix({ at(1, 0), at(1, 2), at(1, 3), at(2, 0), at(2, 2), at(2, 3), at(3, 0), at(3, 2), at(3, 3) }).determinant(); + accum += at(0, 2) * Matrix({ at(1, 0), at(1, 1), at(1, 3), at(2, 0), at(2, 1), at(2, 3), at(3, 0), at(3, 1), at(3, 3) }).determinant(); + accum -= at(0, 3) * Matrix({ at(1, 0), at(1, 1), at(1, 2), at(2, 0), at(2, 1), at(2, 2), at(3, 0), at(3, 1), at(3, 2) }).determinant(); + + return accum; + } + + /// Computes the inverse of a 4-by-4 matrix (ignores the optional argument) + MCTLASS_HOST_DEVICE + Matrix inverse(Element ignore = 1) const { + Matrix B = slice_2x2(0, 2); + Matrix A = slice_2x2(0, 0); + Matrix C = slice_2x2(2, 0); + Matrix D = slice_2x2(2, 2); + + Matrix D_inv = D.inverse(); + + Matrix E = (A - B * D_inv * C).inverse(); + + return Matrix::block( + E, -E * B * D_inv, + -D_inv * C * E, D_inv + D_inv * C * E * B * D_inv + ); + } + +}; + +/// Template alias for 4-by-4 matrix +template +using Matrix4x4 = Matrix; + + +/// Free funciton to infer element type from template arguments +template +MCTLASS_HOST_DEVICE Matrix4x4 make_Matrix4x4( + Element _0_0, Element _0_1, Element _0_2, Element _0_3, + Element _1_0, Element _1_1, Element _1_2, Element _1_3, + Element _2_0, Element _2_1, Element _2_2, Element _2_3, + Element _3_0, Element _3_1, Element _3_2, Element _3_3 +) { + return Matrix4x4( + _0_0, _0_1, _0_2, _0_3, + _1_0, _1_1, _1_2, _1_3, + _2_0, _2_1, _2_2, _2_3, + _3_0, _3_1, _3_2, _3_3 + ); +} + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Elementwise scalar multiplication +template +MCTLASS_HOST_DEVICE +Matrix operator*(Element s, Matrix const &rhs) { + return rhs.multiply(s); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/matrix_coord.h b/csrc/mctlass/include/mctlass/matrix_coord.h new file mode 100644 index 0000000..2aea6e6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/matrix_coord.h @@ -0,0 +1,164 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines a canonical coordinate for rank=2 matrices offering named indices. +*/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/coord.h" + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// MatrixCoord wraps Coord<2, int> to provide a helper for accessing named dimensions. Classes +/// expecting a coordinate in the rank=2 index space of a matrix should use MatrixCoord. +struct MatrixCoord : public Coord<2, int> { + +public: + + /// Integer-valued index + using Index = int; + + /// Base type is a Coord of rank=2 + using Base = Coord<2, Index>; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + +private: + + /// Rows dimension + static int const kRow = 0; + + /// Columns dimension + static int const kColumn = 1; + +public: + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + MatrixCoord() { } + + /// Constructs from Coord<2> + MCTLASS_HOST_DEVICE + MatrixCoord(Coord<2, Index> const &coord): Base(coord) { } + + /// Helper to construct from a row and column + MCTLASS_HOST_DEVICE + MatrixCoord(Index row, Index column): Base(make_Coord(row, column)) { } + + /// Helper to construct from a row and column, which are LongIndex based + MCTLASS_HOST_DEVICE + MatrixCoord(LongIndex row, LongIndex column): Base(make_Coord(Index(row), Index(column))) { } + + /// Returns the row of the coordinate + MCTLASS_HOST_DEVICE + Index const & row() const { return this->at(kRow); } + + /// Returns the row of the coordinate + MCTLASS_HOST_DEVICE + Index & row() { return this->at(kRow); } + + /// Returns the column of the coordinate + MCTLASS_HOST_DEVICE + Index const & column() const { return this->at(kColumn); } + + /// Returns the column of the coordinate + MCTLASS_HOST_DEVICE + Index & column() { return this->at(kColumn); } + + // + // Coord operators + // + + /// Element-wise addition + MCTLASS_HOST_DEVICE + MatrixCoord operator+(Base const& b) const { + return MatrixCoord(Base::operator+(b)); + } + + /// Element-wise subtraction + MCTLASS_HOST_DEVICE + MatrixCoord operator-(Base const& b) const { + return MatrixCoord(Base::operator-(b)); + } + + /// Element-wise multiplication + MCTLASS_HOST_DEVICE + MatrixCoord operator*(Base const& b) const { + return MatrixCoord(Base::operator*(b)); + } + + /// Element-wise division + MCTLASS_HOST_DEVICE + MatrixCoord operator/(Base const& b) const { + return MatrixCoord(Base::operator/(b)); + } + + /// In-place addition + MCTLASS_HOST_DEVICE + MatrixCoord& operator+=(Base const& b) { + Base::operator+=(b); + return *this; + } + + /// In-place subtraction + MCTLASS_HOST_DEVICE + MatrixCoord& operator-=(Base const& b) { + Base::operator-=(b); + return *this; + } + + /// In-place multiplication + MCTLASS_HOST_DEVICE + MatrixCoord& operator*=(Base const& b) { + Base::operator*=(b); + return *this; + } + + /// In-place division + MCTLASS_HOST_DEVICE + MatrixCoord& operator/=(Base const& b) { + Base::operator/=(b); + return *this; + } + +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/matrix_shape.h b/csrc/mctlass/include/mctlass/matrix_shape.h new file mode 100644 index 0000000..bc3f98f --- /dev/null +++ b/csrc/mctlass/include/mctlass/matrix_shape.h @@ -0,0 +1,65 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines a Shape template for matrix tiles +*/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/coord.h" + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Describes the size of a matrix tile +template < + int Row_, ///< rows of a matrix + int Column_ ///< columns of a matrix +> +struct MatrixShape { + static int const kRow = Row_; ///< rows of a matrix + static int const kColumn = Column_; ///< columns of a matrix + static int const kCount = Row_ * Column_; ///< total number of elements in a matrix + + // + // Static member functions + // + + MCTLASS_HOST_DEVICE + static Coord<2> toCoord() { + return make_Coord(kRow, kColumn); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/mctlass.h b/csrc/mctlass/include/mctlass/mctlass.h new file mode 100644 index 0000000..3c0e64d --- /dev/null +++ b/csrc/mctlass/include/mctlass/mctlass.h @@ -0,0 +1,245 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Basic include for MCTLASS. +*/ + +#pragma once +#include "mc_runtime_types.h" +#include "mc_runtime_api.h" +//////////////////////////////////////////////////////////////////////////////////////////////////// +#ifdef MCTLASS_NAMESPACE +#define concat_tok(a, b) a ## b +#define mkmctlassnamespace(pre, ns) concat_tok(pre, ns) +#define mctlass mkmctlassnamespace(mctlass_, MCTLASS_NAMESPACE) +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(__MACACC__) || (defined(__clang__) && defined(__MACA__)) + #define MCTLASS_HOST_DEVICE __forceinline__ __device__ __host__ + #define MCTLASS_DEVICE __forceinline__ __device__ +#elif defined(__MACACC_RTC__) + #define MCTLASS_HOST_DEVICE __forceinline__ __device__ + #define MCTLASS_DEVICE __forceinline__ __device__ +#else + #define MCTLASS_HOST_DEVICE inline + #define MCTLASS_DEVICE inline +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +MCTLASS_HOST_DEVICE void __MCTLASS_UNUSED(T const &) +{ } + +#if defined(__GNUC__) + #define MCTLASS_UNUSED(expr) __MCTLASS_UNUSED(expr) +#else + #define MCTLASS_UNUSED(expr) do { ; } while (&expr != &expr) +#endif + +#ifdef _MSC_VER +// Provides support for alternative operators 'and', 'or', and 'not' +#include +#endif // _MSC_VER + +#if !defined(__MACACC_RTC__) +#include +#endif + +#if defined(__MACA_ARCH__) + #if defined(_MSC_VER) + #define MCTLASS_NOT_IMPLEMENTED() { printf("%s not implemented\n", __FUNCSIG__); asm (";maca not implemented;\n"); } + #else + #define MCTLASS_NOT_IMPLEMENTED() { printf("%s not implemented\n", __PRETTY_FUNCTION__); asm (";maca not implemented;\n"); } + #endif +#else + #if defined(_MSC_VER) + #define MCTLASS_NOT_IMPLEMENTED() assert(0 && __FUNCSIG__) + #else + #define MCTLASS_NOT_IMPLEMENTED() assert(0 && __PRETTY_FUNCTION__) + #endif +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +/// Status code returned by MCTLASS operations +enum class Status { + kSuccess, ///< Operation was successful. + kErrorMisalignedOperand, ///< operands fail alignment requirements. + kErrorInvalidDataType, ///< DataType fails requirement. + kErrorInvalidLayout, ///< Layout fails alignment requirement. + kErrorInvalidProblem, ///< Specified problem size is not supported by operator. + kErrorNotSupported, ///< Operation is not supported on current device. + kErrorWorkspaceNull, ///< The given workspace is null when it is required to be non-null. + kErrorInternal, ///< An error within MCTLASS occurred. + kErrorArchMismatch, ///< MCTLASS runs on a device that it was not compiled for. + kErrorInsufficientDriver, ///< MCTLASS runs with a driver that is too old. + kErrorMemoryAllocation, ///< Kernel launch failed due to insufficient device memory. + kInvalid ///< Status is unspecified. +}; + +/// Convert mctlass status to status strings +MCTLASS_HOST_DEVICE +static char const* mctlassGetStatusString(mctlass::Status status) { + switch (status) { + case mctlass::Status::kSuccess: + return "Success"; + case mctlass::Status::kErrorMisalignedOperand: + return "Error Misaligned Operand"; + case mctlass::Status::kErrorInvalidDataType: + return "Error Invalid Data Type"; + case mctlass::Status::kErrorInvalidLayout: + return "Error Invalid Layout"; + case mctlass::Status::kErrorInvalidProblem: + return "Error Invalid Problem"; + case mctlass::Status::kErrorNotSupported: + return "Error Not Supported"; + case mctlass::Status::kErrorWorkspaceNull: + return "Error Workspace Null"; + case mctlass::Status::kErrorInternal: + return "Error Internal"; + case mctlass::Status::kErrorInsufficientDriver: + return "Error Insufficient Driver"; + case mctlass::Status::kErrorArchMismatch: + return "Error Architecture Mismatch"; + case mctlass::Status::kErrorMemoryAllocation: + return "Error Memory Allocation failed"; + case mctlass::Status::kInvalid: break; + } + + return "Invalid status"; +} + +template +static __inline__ __host__ mcError_t mcFuncSetAttribute( + T *entry, + mcFuncAttribute attr, + int value +) +{ + return mcFuncSetAttribute((const void *)entry,attr,value); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + + +#ifndef MCTLASS_CONV_UNIT_TEST_RIGOROUS_SIZE_ENABLED +#define MCTLASS_CONV_UNIT_TEST_RIGOROUS_SIZE_ENABLED 0 +#endif + + +// CUDA 10.1 introduces the mma instruction +#if !defined(MCTLASS_ENABLE_TENSOR_CORE_MMA) +#define MCTLASS_ENABLE_TENSOR_CORE_MMA 0 +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#define MCTLASS_ASSERT(x) assert(x) + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// MCTLASS_PRAGMA_(UNROLL|NO_UNROLL) optimization directives for the MACA compiler. +#if defined(__MACA_ARCH__) && !defined(__INTELLISENSE__) + #if defined(__MACACC_RTC__) || (defined(__clang__) && defined(__MACA__)) + #define MCTLASS_PRAGMA_UNROLL _Pragma("unroll") + #define MCTLASS_PRAGMA_NO_UNROLL _Pragma("unroll 1") + #else + // We cannot build with this define for mxcc + //#define MCTLASS_PRAGMA_UNROLL #pragma unroll + //#define MCTLASS_PRAGMA_NO_UNROLL #pragma unroll 1 + //#define MCTLASS_PRAGMA_UNROLL + //#define MCTLASS_PRAGMA_NO_UNROLL + #define MCTLASS_PRAGMA_UNROLL _Pragma("unroll") + #define MCTLASS_PRAGMA_NO_UNROLL _Pragma("unroll 1") + #endif + + #define MCTLASS_GEMM_LOOP MCTLASS_PRAGMA_NO_UNROLL + +#else + + #define MCTLASS_PRAGMA_UNROLL + #define MCTLASS_PRAGMA_NO_UNROLL + #define MCTLASS_GEMM_LOOP + +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +static const int NumThreadsPerWarp = 64; +static const int NumThreadsPerWarpGroup = 128; +static const int NumThreadsPerHalfWarp = NumThreadsPerWarp / 2; +static const int NumThreadsPerQuad = 4; +static const int NumThreadsPerQuadPair = NumThreadsPerQuad * 2; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Helper function to return true when called by thread 0 of threadblock 0. +MCTLASS_HOST_DEVICE bool thread0() { + #if defined(__MACA_ARCH__) + return (!threadIdx.x && !threadIdx.y && !threadIdx.z) && (!blockIdx.x && !blockIdx.y && !blockIdx.z); + #else + return false; + #endif +} + +/// Returns a warp-uniform value indicating the canonical warp index of the calling threads. +/// Threads within the warp must be converged. +MCTLASS_DEVICE +int canonical_warp_idx() { + #if defined(__MACA_ARCH__) + return __shfl_sync(0xffffffff, threadIdx.x / NumThreadsPerWarp, 0); + #else + return 0; + #endif +} + +/// Returns a warp-uniform value indicating the canonical warp group index of the calling threads. +/// Threads within the warp must be converged. +MCTLASS_DEVICE +int canonical_warp_group_idx() { + #if defined(__MACA_ARCH__) + return __shfl_sync(0xffffffff, threadIdx.x / NumThreadsPerWarpGroup, 0); + #else + return 0; + #endif +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/numeric_conversion.h b/csrc/mctlass/include/mctlass/numeric_conversion.h new file mode 100644 index 0000000..d1474c6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/numeric_conversion.h @@ -0,0 +1,2649 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Boost-like numeric conversion operator for MCTLASS numeric types +*/ +#pragma once + +#if !defined(__MACACC_RTC__) +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/transform/thread/unary_op.h" + +#include "mctlass/array.h" +#include "mctlass/half.h" + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Floating-point rounding style similare to Standard Library's formats but supporting +/// additional rounding options. +enum class FloatRoundStyle { + round_indeterminate, ///< rounding mode unknown + round_toward_zero, ///< round toward zero + round_to_nearest, ///< round to nearest even + round_toward_infinity, ///< round toward infinity + round_toward_neg_infinity, ///< round toward negative infinity + round_half_ulp_truncate, ///< add 0.5ulp to integer representation then round toward zero + round_half_ulp_trunc_dntz ///< like round_half_ulp_truncate, except denorms are rounded *toward* zero +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename T, + typename S, + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest +> +struct NumericConverter { + + using result_type = T; + using source_type = S; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + return static_cast(s); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s)const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations for float => int32_t +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(__MACA_ARCH__) +template <> +struct NumericConverter { + + using result_type = int32_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + #if defined(__MACA_ARCH__) + return __float2int_rn(s); + #endif + + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s)const { + return convert(s); + } +}; + +template <> +struct NumericConverter { + + using result_type = int32_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + #if defined(__MACA_ARCH__) + return __float2int_rz(s); + #else + return static_cast(s); + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s)const { + return convert(s); + } +}; + +#elif !defined(__MACACC_RTC__) + +template <> +struct NumericConverter { + + using result_type = int32_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest; + + static result_type convert(source_type const & s) { + std::fesetround(FE_TONEAREST); + return (result_type)std::nearbyint(s); + } + + result_type operator()(source_type const &s)const { + return convert(s); + } +}; + +template <> +struct NumericConverter { + + using result_type = int32_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero; + + static result_type convert(source_type const & s) { + std::fesetround(FE_TOWARDZERO); + return (result_type)std::nearbyint(s); + } + + result_type operator()(source_type const &s)const { + return convert(s); + } +}; +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations for float => int8_t +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(__MACA_ARCH__) +template <> +struct NumericConverter { + + using result_type = int8_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + #if defined(__MACA_ARCH__) + int32_t intermediate = __float2int_rn(s); + #else + + float r = roundf(s); + float d = r - s; + int32_t intermediate = static_cast(s - d); + if ((d != 0.5f) && (d != -0.5f)) intermediate= static_cast(r); + if(fmod(r,2.0f)==0.0f)intermediate= static_cast(r); + + #endif + return static_cast(intermediate); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s)const { + return convert(s); + } +}; + +template <> +struct NumericConverter { + + using result_type = int8_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + #if defined(__MACA_ARCH__) + int32_t intermediate = __float2int_rz(s); + #else + int32_t intermediate=static_cast(s); + #endif + + return static_cast(intermediate); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s)const { + return convert(s); + } +}; + +#elif !defined(__MACACC_RTC__) + +template <> +struct NumericConverter { + + using result_type = int8_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest; + + static result_type convert(source_type const & s) { + std::fesetround(FE_TONEAREST); + int32_t intermediate = (int32_t)std::nearbyint(s); + + // Low-end saturation + intermediate = std::max(intermediate, (int32_t)std::numeric_limits::lowest()); + + // High-end saturation + intermediate = std::min(intermediate, (int32_t)std::numeric_limits::max()); + + return static_cast(intermediate); + } + + result_type operator()(source_type const &s)const { + return convert(s); + } +}; + +template <> +struct NumericConverter { + + using result_type = int8_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero; + + static result_type convert(source_type const & s) { + std::fesetround(FE_TOWARDZERO); + int32_t intermediate = (int32_t)std::nearbyint(s); + + // Low-end saturation + intermediate = std::max(intermediate, (int32_t)std::numeric_limits::lowest()); + + // High-end saturation + intermediate = std::min(intermediate, (int32_t)std::numeric_limits::max()); + + return static_cast(intermediate); + } + + result_type operator()(source_type const &s)const { + return convert(s); + } +}; + +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for float <= half_t +template +struct NumericConverter { + + using result_type = T; + using source_type = T; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + return s; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s)const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations for float <=> half_t +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for float <= half_t +template +struct NumericConverter { + + using result_type = float; + using source_type = half_t; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + result_type result = static_cast(s); + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s)const { + return convert(s); + } +}; + +/// Specialization for round-to-nearest +template <> +struct NumericConverter { + + using result_type = half_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + result_type result = static_cast(s); + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Specialization for round-toward-zero +template <> +struct NumericConverter { + + using result_type = half_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero; + + /// Round toward zero + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & flt) { + + #if defined(__MACA_ARCH__) + return half_t(__float2half_rz(flt)); + #else + // software implementation rounds toward nearest even + unsigned const& s = reinterpret_cast(flt); + uint16_t sign = uint16_t((s >> 16) & 0x8000); + int16_t exp = uint16_t(((s >> 23) & 0xff) - 127); + int mantissa = s & 0x7fffff; + uint16_t u = 0; + + if ((s & 0x7fffffff) == 0) { + // sign-preserving zero + return half_t::bitcast(sign); + } + + if (exp > 15) { + if (exp == 128 && mantissa) { + // not a number + u = 0x7fff; + } else { + // overflow to infinity + u = sign | 0x7c00; + } + return half_t::bitcast(u); + } + + if (exp >= -14) { + // normal fp32 to normal fp16 + exp = uint16_t(exp + uint16_t(15)); + u = uint16_t(((exp & 0x1f) << 10)); + u = uint16_t(u | (mantissa >> 13)); + } else { + // normal single-precision to subnormal half_t-precision representation + int rshift = (-14 - exp); + if (rshift < 32) { + mantissa |= (1 << 23); + mantissa = (mantissa >> rshift); + u = (uint16_t(mantissa >> 13) & 0x3ff); + } else { + mantissa = 0; + u = 0; + } + } + + u |= sign; + + return half_t::bitcast(u); + + #endif // defined(__MACA_ARCH__) + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations for float <=> bfloat16_t +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for float <= bfloat16_t +template +struct NumericConverter { + + using result_type = float; + using source_type = bfloat16_t; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + return static_cast(s); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +template <> +struct NumericConverter { + using result_type = bfloat16_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + // Note: __builtin_mxc_cvt_f32tobf16_fast is slower in xcore1000 now + // use __builtin_mxc_cvt_f32tobf16_fast in xcore1000 after complier fix this + #if defined(__MACA_ARCH__) && (__MACA_ARCH__ == 1500) + bfloat16_t result; + unsigned short* temp_result = reinterpret_cast(&result); + *temp_result = __builtin_mxc_cvt_f32tobf16_fast(s); + return result; + #else + return static_cast(s); + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +template <> +struct NumericConverter { + using result_type = bfloat16_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_half_ulp_truncate; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + uint32_t x32 = reinterpret_cast(s); + + #if defined(__MACA_ARCH__) + if (::isfinite(s)) { + x32 += 0x8000; + } + #else + if (std::isfinite(s)) { + x32 += 0x8000; + } + #endif + + uint16_t x16 = uint16_t((x32 >> 16) & 0xffff); + return bfloat16_t::bitcast(x16); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +template <> +struct NumericConverter { + using result_type = bfloat16_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + uint32_t x32 = reinterpret_cast(s); + uint16_t x16 = uint16_t(x32 >> 16); + + return bfloat16_t::bitcast(x16); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations for float <=> tfloat32_t +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for float <= tfloat32_t +template +struct NumericConverter { + + using result_type = float; + using source_type = tfloat32_t; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + return static_cast(s); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +template <> +struct NumericConverter { + using result_type = tfloat32_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + unsigned storage = reinterpret_cast(s); + +//#if defined(__MACA_ARCH__) && __MACA_ARCH__ >= 900 +// asm volatile("cvt.rn.tf32.f32 %0, %1;" : "=r"(storage) : "r"(storage)); +//#else +#if defined(__MACA_ARCH__) + if ((storage & 0x7f800000) != 0x7f800000) { + + bool mantissa_bit = ((storage & (1 << 13)) != 0); + bool round_bit = ((storage & (1 << 12)) != 0); + bool sticky_bit = ((storage & ((1 << 12) - 1)) != 0); + + if ((round_bit && sticky_bit) || (round_bit && mantissa_bit)) { + storage += uint32_t(1 << 13); + } + + // Note, the following is intentionally commented out. TF32 + // does not define the low order bits, so they may be left in + // an undefined state. + // + // By not truncating these bit explicitly, we avoid an extra logical + // operation. + // + // TF32 may be implicitly converted to float by performing this + // operation as needed. + // + // storage = (storage & ~0x1fff); + } + else if (storage & ~0xff800000) { + storage = 0x7fffffff; + } +#endif + + return tfloat32_t::bitcast(storage); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +template <> +struct NumericConverter { + using result_type = tfloat32_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_half_ulp_truncate; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + return tfloat32_t::round_half_ulp_truncate(s); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// This rounding operation is similar to half_ulp_truncate except it rounds denorms toward zero. +/// It avoids predicated code, though it requires a temporary register. +template <> +struct NumericConverter { + using result_type = tfloat32_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_half_ulp_trunc_dntz; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + unsigned y = reinterpret_cast(s); + y = y & 0xff800000; + float d = reinterpret_cast(y); + float z = d / float(1 << 11) + s; + + return reinterpret_cast(z); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +template <> +struct NumericConverter { + using result_type = tfloat32_t; + using source_type = float; + static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + uint32_t x = reinterpret_cast(s); + return tfloat32_t::bitcast(x & 0xffffe000); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Conversion operator for float to tfloat32_t big and small values +// +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + FloatRoundStyle RoundBig = FloatRoundStyle::round_toward_zero, + FloatRoundStyle RoundSmall = FloatRoundStyle::round_half_ulp_truncate +> +struct NumericConverterFastF32 { + + // result_type holds big tfloat32_t at idx(0) and small tfloat32_t at idx(1) + using result_type = Array; + + // source data type + using source_type = float; + + // rounding styles for big and small part + static FloatRoundStyle const kRoundBig = RoundBig; + static FloatRoundStyle const kRoundSmall = RoundSmall; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + result_type result; + NumericConverter convert_big_; + NumericConverter convert_small_; + + // convert and fill tfloat32_t big at idx 0 + result[0] = convert_big_(source); + + // convert and fill tfloat32_t small at idx 1 + result[1] = convert_small_(source - static_cast(result[0])); + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Conversion and Clamp operator for Integers +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename T, + typename S +> +struct NumericConverterClamp { + + using result_type = T; + using source_type = S; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + NumericConverter convert_op; + result_type const kClamp_max = platform::numeric_limits::max(); + result_type const kClamp_min = platform::numeric_limits::lowest(); + if (s < (source_type)kClamp_min) + return kClamp_min; + if (s > (source_type)kClamp_max) + return kClamp_max; + return convert_op(s); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +// This converter is needed to enable half_t output types when using int32_t accumulators. +// Since floating-point types do not require a clamp, this converter simply casts from +// the source type to half_t. +template < + typename S +> +struct NumericConverterClamp { + + using result_type = mctlass::half_t; + using source_type = S; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const &source) { + return static_cast(source); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Conversion operator for Array +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Conversion operator for Array +template < + typename T, + typename S, + int N, + FloatRoundStyle Round = FloatRoundStyle::round_to_nearest, + typename Transform = mctlass::transform::thread::UnaryTransform::Identity +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + static_assert(platform::is_same::value || + platform::is_same::value, + "Unary Operator not supported."); + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & s) { + + result_type result; + NumericConverter convert_; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + if( platform::is_same::value ) + { + result[i] = convert_(s[i]); + } else { // conjugate + result[i] = conj(convert_(s[i])); + } + } + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +template < + typename T, + int N, + FloatRoundStyle Round, + typename Transform +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + static_assert(platform::is_same::value || + platform::is_same::value, + "Unary Operator not supported."); + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const &source) { + if (platform::is_same::value) { + return source; + } else { + result_type result; + for (int i = 0; i < N; ++i) { + result[i] = conj(source[i]); + } + return result; + } + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Array <= Array, round to nearest +template <> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + Array result; + + #if !defined(__FAST_HALF_CVT__) && defined(__MACA_ARCH__) + // float22half2_rn not fast now. + reinterpret_cast<__half2 &>(result) = __float22half2_rn(reinterpret_cast(source)); + #else + //NumericConverter convert_; + //result[0] = convert_(source[0]); + //result[1] = convert_(source[1]); + /****** + * reinterpret_cast array to half2 can reduce MTE inst after maca-20240514-965 + ******/ + NumericConverter convert_; + auto res = reinterpret_cast<__half2 *>(&result); + res->x = convert_(source[0]); + res->y = convert_(source[1]); + + // typedef __NATIVE_VECTOR__(2, uint16_t) v2i16; + // *(reinterpret_cast(&result)) = {convert_(source[0]).storage, convert_(source[1]).storage}; + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array, round to nearest +template +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + Array result; + + #if defined(__MACA_ARCH__) + reinterpret_cast(result) = __half22float2(reinterpret_cast<__half2 const &>(source)); + #else + NumericConverter convert_; + result[0] = convert_(source[0]); + result[1] = convert_(source[1]); + #endif + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Array <= Array +template < + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + NumericArrayConverter convert_vector_; + NumericConverter convert_element_; + + result_type result; + + Array *result_ptr = reinterpret_cast *>(&result); + Array const *source_ptr = reinterpret_cast const *>(&source); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = convert_vector_(source_ptr[i]); + } + + if (N % 2) { + result[N - 1] = convert_element_(source[N - 1]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + + +/// Partial specialization for Array <= Array +template < + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + NumericArrayConverter convert_vector_; + NumericConverter convert_element_; + + result_type result; + + Array *result_ptr = reinterpret_cast *>(&result); + Array const *source_ptr = reinterpret_cast const *>(&source); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = convert_vector_(source_ptr[i]); + } + + if (N % 2) { + result[N - 1] = convert_element_(source[N - 1]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#if defined(__MACA_ARCH__) +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Array <= Array, round to nearest +template <> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + // unsigned d; + // asm("cvt.rn.bf16x2.f32 %0, %1, %2;\n" : "=r"(d) : "f"(source[1]), "f"(source[0]) ); + // return reinterpret_cast(d); + + result_type result; +#if defined(__FAST_BF16_CVT__) + typedef __NATIVE_VECTOR__(2, uint16_t) v2i16; + *(reinterpret_cast(&result)) = {(*(v2i16 *)&source[0]).s1, (*(v2i16 *)&source[1]).s1}; +#else + NumericConverter convert_; + result[0] = convert_(source[0]); + result[1] = convert_(source[1]); +#endif + return result; + + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Array <= Array +template < + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + NumericArrayConverter convert_vector_; + NumericConverter convert_element_; + + result_type result; + + Array *result_ptr = reinterpret_cast *>(&result); + Array const *source_ptr = reinterpret_cast const *>(&source); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 2; ++i) { + result_ptr[i] = convert_vector_(source_ptr[i]); + } + + if (N % 2) { + result[N - 1] = convert_element_(source[N - 1]); + } + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template +struct NumericArrayConverter { + using result_type = Array; + using source_type = Array; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + result_type result; + half_t *result_ptr = reinterpret_cast(&result); + const int8_t *source_ptr = reinterpret_cast(&source); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result_ptr[i] = __half(source_ptr[i]); + } + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +#endif // if defined(__MACA_ARCH__) + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Conditional guards to enable partial specialization for packed integers +#if defined(__MACA_ARCH__) && \ + ((__CUDACC_VER_MAJOR__ > 10) || \ + ((__CUDACC_VER_MAJOR__ >= 10) && (__CUDACC_VER_MINOR__ >= 2))) + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + NumericConverter convert_element_; + + result_type result; + + result[0] = convert_element_(source[0]); + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + uint32_t tmp; + #if 0 + asm volatile( + "cvt.pack.sat.s8.s32.b32 %0, %2, %1, 0;\n" + : "=r"(tmp) : "r"(source[0]), "r"(source[1])); + #endif + uint16_t out = (tmp & 0xffff); + printf("numeric_conversion.h L1106 this function cannot working correctly now.\n"); + return reinterpret_cast(out); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + unsigned out; +#if 0 + asm volatile( + "{ .reg .u32 r4;" + "cvt.pack.sat.s8.s32.b32 r4, %4, %3, 0;" + "cvt.pack.sat.s8.s32.b32 %0, %2, %1, r4;" + "}" + : "=r"(out) : "r"(source[0]), "r"(source[1]), "r"(source[2]), "r"(source[3])); +#endif + printf("numeric_conversion.h L1138 this function cannot working correctly now.\n"); + return reinterpret_cast(out); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter { + static_assert(!(N % 4), "N must be multiple of 4."); + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + NumericArrayConverter convert_vector_; + + result_type result; + + Array *result_ptr = reinterpret_cast *>(&result); + Array const *source_ptr = reinterpret_cast const *>(&source); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 4; ++i) { + result_ptr[i] = convert_vector_(source_ptr[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + NumericConverter convert_element_; + + result_type result; + + result[0] = convert_element_(source[0]); + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + uint32_t tmp; +#if 0 + asm volatile( + "cvt.pack.sat.u8.s32.b32 %0, %2, %1, 0;\n" + : "=r"(tmp) : "r"(source[0]), "r"(source[1])); +#endif + uint16_t out = (tmp & 0xffff); + printf("numeric_conversion L1231 this function cannot working correctly now.\n"); + return reinterpret_cast(out); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + unsigned out; +#if 0 + asm volatile( + "{ .reg .u32 r4;" + "cvt.pack.sat.u8.s32.b32 r4, %4, %3, 0;" + "cvt.pack.sat.u8.s32.b32 %0, %2, %1, r4;" + "}" + : "=r"(out) : "r"(source[0]), "r"(source[1]), "r"(source[2]), "r"(source[3])); +#endif + printf("numeric_conversion.h L1263 this function cannot working correctly now.\n"); + return reinterpret_cast(out); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter { + static_assert(!(N % 4), "N must be multiple of 4."); + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + NumericArrayConverter convert_vector_; + + result_type result; + + Array *result_ptr = reinterpret_cast *>(&result); + Array const *source_ptr = reinterpret_cast const *>(&source); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 4; ++i) { + result_ptr[i] = convert_vector_(source_ptr[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations for Array <=> Array +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = float; + using source_element = float_e4m3_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + + #if defined(CUDA_PTX_FP8_CVT_ENABLED) + uint32_t out_fp16[2]; + uint32_t const& src_packed = reinterpret_cast(source); + + asm volatile( \ + "{\n" \ + ".reg .b16 lo, hi;\n" \ + "mov.b32 {lo, hi}, %2;\n" \ + "cvt.rn.f16x2.e4m3x2 %0, lo;\n" \ + "cvt.rn.f16x2.e4m3x2 %1, hi;\n" \ + "}\n" : "=r"(out_fp16[0]), "=r"(out_fp16[1]) : "r"(src_packed)); + + float2 res0 = __half22float2(reinterpret_cast<__half2 &>(out_fp16[0])); + float2 res1 = __half22float2(reinterpret_cast<__half2 &>(out_fp16[1])); + + result_type out; + out[0] = res0.x; + out[1] = res0.y; + out[2] = res1.x; + out[3] = res1.y; + return out; + #else + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = float_e4m3_t; + using source_element = float; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + + #if defined(CUDA_PTX_FP8_CVT_ENABLED) + uint32_t out; + + asm volatile( \ + "{\n" \ + ".reg .b16 lo;\n" \ + ".reg .b16 hi;\n" \ + "cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n" \ + "cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n" \ + "mov.b32 %0, {lo, hi};\n" \ + "}" \ + : "=r"(out) : "f"(source[0]), "f"(source[1]), "f"(source[2]), "f"(source[3])); + + return reinterpret_cast(out); + #else + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = float; + using source_element = float_e5m2_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + + #if defined(CUDA_PTX_FP8_CVT_ENABLED) + uint32_t out_fp16[2]; + uint32_t const& src_packed = reinterpret_cast(source); + + asm volatile( \ + "{\n" \ + ".reg .b16 lo, hi;\n" \ + "mov.b32 {lo, hi}, %2;\n" \ + "cvt.rn.f16x2.e5m2x2 %0, lo;\n" \ + "cvt.rn.f16x2.e5m2x2 %1, hi;\n" \ + "}\n" : "=r"(out_fp16[0]), "=r"(out_fp16[1]) : "r"(src_packed)); + + float2 res0 = __half22float2(reinterpret_cast<__half2 &>(out_fp16[0])); + float2 res1 = __half22float2(reinterpret_cast<__half2 &>(out_fp16[1])); + + result_type out; + out[0] = res0.x; + out[1] = res0.y; + out[2] = res1.x; + out[3] = res1.y; + return out; + #else + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = float_e5m2_t; + using source_element = float; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + + #if defined(CUDA_PTX_FP8_CVT_ENABLED) + uint32_t out; + + asm volatile( \ + "{\n" \ + ".reg .b16 lo;\n" \ + ".reg .b16 hi;\n" \ + "cvt.rn.satfinite.e5m2x2.f32 lo, %2, %1;\n" \ + "cvt.rn.satfinite.e5m2x2.f32 hi, %4, %3;\n" \ + "mov.b32 %0, {lo, hi};\n" \ + "}" \ + : "=r"(out) : "f"(source[0]), "f"(source[1]), "f"(source[2]), "f"(source[3])); + + return reinterpret_cast(out); + #else + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations for Array <=> Array +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = half_t; + using source_element = float_e4m3_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + + #if defined(CUDA_PTX_FP8_CVT_ENABLED) + uint32_t out[2]; + uint32_t const& src_packed = reinterpret_cast(source); + asm volatile( \ + "{\n" \ + ".reg .b16 lo, hi;\n" \ + "mov.b32 {lo, hi}, %2;\n" \ + "cvt.rn.f16x2.e4m3x2 %0, lo;\n" \ + "cvt.rn.f16x2.e4m3x2 %1, hi;\n" \ + "}\n" : "=r"(out[0]), "=r"(out[1]) : "r"(src_packed)); + return reinterpret_cast(out); + #else + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = float_e4m3_t; + using source_element = half_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + + #if defined(CUDA_PTX_FP8_CVT_ENABLED) + uint32_t out; + uint32_t const* src_packed = reinterpret_cast(&source); + + asm volatile( \ + "{\n" \ + ".reg .b16 lo;\n" \ + ".reg .b16 hi;\n" \ + "cvt.rn.satfinite.e4m3x2.f16x2 lo, %1;\n" \ + "cvt.rn.satfinite.e4m3x2.f16x2 hi, %2;\n" \ + "mov.b32 %0, {lo, hi};\n" \ + "}" \ + : "=r"(out) : "r"(src_packed[0]), "r"(src_packed[1])); + + return reinterpret_cast(out); + #else + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = half_t; + using source_element = float_e5m2_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + + #if defined(CUDA_PTX_FP8_CVT_ENABLED) + uint32_t out[2]; + uint32_t const& src_packed = reinterpret_cast(source); + asm volatile( \ + "{\n" \ + ".reg .b16 lo, hi;\n" \ + "mov.b32 {lo, hi}, %2;\n" \ + "cvt.rn.f16x2.e5m2x2 %0, lo;\n" \ + "cvt.rn.f16x2.e5m2x2 %1, hi;\n" \ + "}\n" : "=r"(out[0]), "=r"(out[1]) : "r"(src_packed)); + return reinterpret_cast(out); + #else + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = float_e5m2_t; + using source_element = half_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + + #if defined(CUDA_PTX_FP8_CVT_ENABLED) + uint32_t out; + uint32_t const* src_packed = reinterpret_cast(&source); + + asm volatile( \ + "{\n" \ + ".reg .b16 lo;\n" \ + ".reg .b16 hi;\n" \ + "cvt.rn.satfinite.e5m2x2.f16x2 lo, %1;\n" \ + "cvt.rn.satfinite.e5m2x2.f16x2 hi, %2;\n" \ + "mov.b32 %0, {lo, hi};\n" \ + "}" \ + : "=r"(out) : "r"(src_packed[0]), "r"(src_packed[1])); + + return reinterpret_cast(out); + #else + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations for Array <=> Array +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = bfloat16_t; + using source_element = float_e4m3_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + + #if defined(CUDA_PTX_FP8_CVT_ENABLED) + // Convert f8 to float + NumericArrayConverter src2float; + Array tmp_floats = src2float(source); + + // Convert float to bf16 + result_type out; + Array* packed_tmp = reinterpret_cast*>(&tmp_floats); + Array* packed_out = reinterpret_cast*>(&out); + NumericArrayConverter float2result; + packed_out[0] = float2result(packed_tmp[0]); + packed_out[1] = float2result(packed_tmp[1]); + + return out; + #else + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = float_e4m3_t; + using source_element = bfloat16_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + + #if defined(CUDA_PTX_FP8_CVT_ENABLED) + // Convert bf16 to float + Array tmp; + Array* packed_tmp = reinterpret_cast*>(&tmp); + Array const* packed_source = reinterpret_cast const*>(&source); + NumericArrayConverter src2float; + packed_tmp[0] = src2float(packed_source[0]); + packed_tmp[1] = src2float(packed_source[1]); + + // Convert float to f8 + NumericArrayConverter float2result; + return float2result(tmp); + #else + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = bfloat16_t; + using source_element = float_e5m2_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + + #if defined(CUDA_PTX_FP8_CVT_ENABLED) + // Convert f8 to float + NumericArrayConverter src2float; + Array tmp_floats = src2float(source); + + // Convert float to bf16 + result_type out; + Array* packed_tmp = reinterpret_cast*>(&tmp_floats); + Array* packed_out = reinterpret_cast*>(&out); + NumericArrayConverter float2result; + packed_out[0] = float2result(packed_tmp[0]); + packed_out[1] = float2result(packed_tmp[1]); + + return out; + #else + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = float_e5m2_t; + using source_element = bfloat16_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + + #if defined(CUDA_PTX_FP8_CVT_ENABLED) + // Convert bf16 to float + Array tmp; + Array* packed_tmp = reinterpret_cast*>(&tmp); + Array const* packed_source = reinterpret_cast const*>(&source); + NumericArrayConverter src2float; + packed_tmp[0] = src2float(packed_source[0]); + packed_tmp[1] = src2float(packed_source[1]); + + // Convert float to f8 + NumericArrayConverter float2result; + return float2result(tmp); + #else + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + #endif + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations for Array <=> Array +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = float_e4m3_t; + using source_element = float_e5m2_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = float_e5m2_t; + using source_element = float_e4m3_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + result_type result; + NumericConverter converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + result[i] = converter(source[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations for: +// Array <=> Array +// Array <=> Array +// +// These are needed to avoid multiple-matching-template compilation errors (e.g., when +// compiling float_e4m3_t <=> float_e4m3_t, which among T <= float_e4m3_t and float_e4m3_t <= T +// should be used?) +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = float_e4m3_t; + using source_element = float_e4m3_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const &source) { + return source; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + using result_element = float_e5m2_t; + using source_element = float_e5m2_t; + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const &source) { + return source; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations for: +// Array <=> Array +// Array <=> Array +// using packed converter under the hood +// +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename T, + typename S, + int N, + FloatRoundStyle Round +> +struct PackedNumericArrayConverter { + using result_element = T; + using source_element = S; + + using result_type = Array; + using source_type = Array; + + static FloatRoundStyle const round_style = Round; + +private: + using packed_result_type = Array; + using packed_source_type = Array; + +public: + MCTLASS_DEVICE + static result_type convert(source_type const & source) { + result_type result; + packed_result_type* packed_result = reinterpret_cast(&result); + const packed_source_type* packed_source = reinterpret_cast(&source); + + NumericArrayConverter packed_converter; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 4; ++i) { + packed_result[i] = packed_converter(packed_source[i]); + } + + // Handle leftovers + NumericConverter converter; + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N % 4; ++i) { + int idx = ((N / 4) * 4) + i; + result[idx] = converter(source[idx]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + typename T, + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter : + public PackedNumericArrayConverter {}; + +/// Partial specialization for Array <= Array +template < + typename T, + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter : + public PackedNumericArrayConverter {}; + +/// Partial specialization for Array <= Array +template < + typename S, + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter : + public PackedNumericArrayConverter {}; + +/// Partial specialization for Array <= Array +template < + typename S, + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter : + public PackedNumericArrayConverter {}; + +/// Partial specialization for Array <= Array +template < + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter : + public PackedNumericArrayConverter {}; + +/// Partial specialization for Array <= Array +template < + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter : + public PackedNumericArrayConverter {}; + +/// Partial specialization for Array <= Array +template < + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter : + public PackedNumericArrayConverter {}; + +/// Partial specialization for Array <= Array +template < + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter : + public PackedNumericArrayConverter {}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Array <= Array +/// Conversion is performed with saturation regardless of setting of +/// the `Round` template parameter. +template < + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + // Convert float to int + Array temporary; + + NumericArrayConverter compute_converter; + temporary = compute_converter(source); + + // Convert to int to int8_t + NumericArrayConverter destination_converter; + return destination_converter(temporary); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// #if defined(__MACA_ARCH__) && \ +// ((__CUDACC_VER_MAJOR__ > 10) || \ +// ((__CUDACC_VER_MAJOR__ >= 10) && (__CUDACC_VER_MINOR__ >= 2))) +#if 0 //mcTlass not support int4b_t + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + unsigned out; +#if 0 + asm volatile( + "{ .reg .u32 r4;" + "cvt.pack.sat.s4.s32.b32 r4, %8, %7, 0;" + "cvt.pack.sat.s4.s32.b32 r4, %6, %5, r4;" + "cvt.pack.sat.s4.s32.b32 r4, %4, %3, r4;" + "cvt.pack.sat.s4.s32.b32 %0, %2, %1, r4;" + "}" + : "=r"(out) + : "r"(source[0]), "r"(source[1]), "r"(source[2]), "r"(source[3]), + "r"(source[4]), "r"(source[5]), "r"(source[6]), "r"(source[7])); +#endif + printf("numeric_conversion.h L1378 this function cannot working correctly now.\n"); + return reinterpret_cast(out); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter { + static_assert(!(N % 8), "N must be multiple of 8."); + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + NumericArrayConverter convert_vector_; + + result_type result; + + Array *result_ptr = reinterpret_cast *>(&result); + Array const *source_ptr = reinterpret_cast const *>(&source); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 8; ++i) { + result_ptr[i] = convert_vector_(source_ptr[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + FloatRoundStyle Round +> +struct NumericArrayConverter { + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + unsigned out; +#if 0 + asm volatile( + "{ .reg .u32 r4;" + "cvt.pack.sat.u4.s32.b32 r4, %8, %7, 0;" + "cvt.pack.sat.u4.s32.b32 r4, %6, %5, r4;" + "cvt.pack.sat.u4.s32.b32 r4, %4, %3, r4;" + "cvt.pack.sat.u4.s32.b32 %0, %2, %1, r4;" + "}" + : "=r"(out) + : "r"(source[0]), "r"(source[1]), "r"(source[2]), "r"(source[3]), + "r"(source[4]), "r"(source[5]), "r"(source[6]), "r"(source[7])); +#endif + printf("numeric_conversion.h L1450 this function cannot working correctly now.\n"); + return reinterpret_cast(out); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +/// Partial specialization for Array <= Array +template < + int N, + FloatRoundStyle Round +> +struct NumericArrayConverter { + static_assert(!(N % 8), "N must be multiple of 8."); + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const & source) { + + NumericArrayConverter convert_vector_; + + result_type result; + + Array *result_ptr = reinterpret_cast *>(&result); + Array const *source_ptr = reinterpret_cast const *>(&source); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 8; ++i) { + result_ptr[i] = convert_vector_(source_ptr[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { + return convert(s); + } +}; + +#endif // Conditional guards to enable partial specialization for packed integers + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// FastNumericArrayConverter only works when the source is within center range. +/// Conversion operator for Array. See the comments before +/// FastLinearCombinationClamp. +template +struct FastNumericArrayConverter { + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const &s) { + result_type result; + NumericArrayConverter convert_; + + return convert_(s); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { return convert(s); } +}; + +/// Partial specialization for Array <= Array +template +struct FastNumericArrayConverter { + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const &source) { + result_type result; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + int tmp = source[i] + 1262485504 /*0x4B400000*/; + result[i] = reinterpret_cast(tmp) - 12582912.0f; + } + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { return convert(s); } +}; + +#if !defined(__MACA_ARCH__) + struct ucharHolder + { + union + { + unsigned char c[4]; + unsigned int ui; + }; + }; + + struct uchar2Holder + { + union + { + unsigned int ui[2]; + unsigned char c[8]; + }; + }; + + MCTLASS_HOST_DEVICE + unsigned int __byte_perm_maca(unsigned int x,unsigned int y,unsigned int s) + { + + struct uchar2Holder cHoldVal; + struct ucharHolder cHoldOut; + + cHoldVal.ui[0]=x; + cHoldVal.ui[1]=y; + + cHoldOut.c[0] = cHoldVal.c[((s << (31 - 2)) >> (31 - 2 + 0))]; + cHoldOut.c[1] = cHoldVal.c[((s << (31 - 6)) >> (31 - 6 + 4))]; + cHoldOut.c[2] = cHoldVal.c[((s << (31 - 10)) >> (31 - 10 + 8))]; + cHoldOut.c[3] = cHoldVal.c[((s << (31 - 14)) >> (31 - 14 + 12))]; + + return cHoldOut.ui; + + } + +#endif + +/// Partial specialization for Array <= Array +template +struct FastNumericArrayConverter { + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const &source) { + Array result; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < 4; ++i) { + float tmp = source[i] + 12582912.0f; + result[i] = reinterpret_cast(tmp); + } + #if defined(__MACA_ARCH__) + result[0] = __byte_perm(result[0], result[1], 0x40); + result[2] = __byte_perm(result[2], result[3], 0x40); + result[0] = __byte_perm(result[0], result[2], 0x5410); + #else + result[0] = __byte_perm_maca(result[0], result[1], 0x40); + result[2] = __byte_perm_maca(result[2], result[3], 0x40); + result[0] = __byte_perm_maca(result[0], result[2], 0x5410); + #endif + + return reinterpret_cast(result[0]); + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { return convert(s); } +}; + +/// Partial specialization for Array <= Array +template +struct FastNumericArrayConverter { + static_assert(!(N % 4), "N must be multiple of 4."); + + using result_type = Array; + using source_type = Array; + static FloatRoundStyle const round_style = Round; + + MCTLASS_HOST_DEVICE + static result_type convert(source_type const &source) { + FastNumericArrayConverter convert_vector_; + + result_type result; + + Array *result_ptr = + reinterpret_cast *>(&result); + Array const *source_ptr = + reinterpret_cast const *>(&source); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / 4; ++i) { + result_ptr[i] = convert_vector_(source_ptr[i]); + } + + return result; + } + + MCTLASS_HOST_DEVICE + result_type operator()(source_type const &s) const { return convert(s); } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines preferred rounding mode for a pair of types +template +struct PreferredRoundingMode { + static FloatRoundStyle const kRound = FloatRoundStyle::round_to_nearest; +}; + +#if defined(__MACA_ARCH__) && __MACA_ARCH__ < 900 +/// Defines preferred rounding mode for a pair of types +template <> +struct PreferredRoundingMode { + static FloatRoundStyle const kRound = FloatRoundStyle::round_half_ulp_truncate; +}; +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Packs predicates into an array. +template +struct PackPredicates { + using result_type = Array; + + static_assert(!(N % 4), "Must pack predicates in a count that is a multiple of 4"); + + MCTLASS_HOST_DEVICE + result_type operator()(bool const predicates[]) { + + result_type packed; + packed.clear(); + + int const kWordSize = 8; + uint8_t *bytes = reinterpret_cast(packed.data()); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + int word_idx = (i / kWordSize); + int bit_idx = (i % kWordSize); + + uint8_t mask = ((predicates[i] ? 1u : 0u) << bit_idx); + bytes[word_idx] = (bytes[word_idx] | mask); + } + return packed; + } +}; + +/// Packs predicates into an array +template +struct UnpackPredicates { + using result_type = Array; + + static_assert(!(N % 4), "Must unpack predicates in a count that is a multiple of 4"); + + MCTLASS_HOST_DEVICE + void operator()(bool predicates[], result_type const &packed) { + + int const kWordSize = 8; + uint8_t const *bytes = reinterpret_cast(packed.data()); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + int word_idx = (i / kWordSize); + int bit_idx = (i % kWordSize); + + predicates[i] = bool((bytes[word_idx] >> bit_idx) & 0x1); + } + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/numeric_types.h b/csrc/mctlass/include/mctlass/numeric_types.h new file mode 100644 index 0000000..da33296 --- /dev/null +++ b/csrc/mctlass/include/mctlass/numeric_types.h @@ -0,0 +1,94 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Top-level include for all MCTLASS numeric types. +*/ +#pragma once + +#include "mctlass/mctlass.h" + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines the size of an element in bits +template +struct sizeof_bits { + static int const value = int(sizeof(T) * 8); +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Definitions for 1-bit binary and 4-bit integer types +// + +/// 1-bit binary type +using bin1_t = bool; + +/// Defines the size of an element in bits - specialized for bin1_t +template <> +struct sizeof_bits { + static int const value = 1; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct index_sequence; + +template +struct index_sequence_helper : index_sequence_helper {}; + +template +struct index_sequence_helper<0, 0, Next...> { + using type = index_sequence<0, Next...>; +}; + +template +using make_index_sequence = typename index_sequence_helper::type; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + +#include "mctlass/integer_subbyte.h" + +#include "mctlass/half.h" +#include "mctlass/bfloat16.h" +#include "mctlass/tfloat32.h" +#include "mctlass/float8.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/pipeline/pipeline.hpp b/csrc/mctlass/include/mctlass/pipeline/pipeline.hpp new file mode 100644 index 0000000..12238d3 --- /dev/null +++ b/csrc/mctlass/include/mctlass/pipeline/pipeline.hpp @@ -0,0 +1,36 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#include "mctlass/pipeline/sm90_pipeline.hpp" +//////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/pipeline/sm90_pipeline.hpp b/csrc/mctlass/include/mctlass/pipeline/sm90_pipeline.hpp new file mode 100644 index 0000000..81bd40a --- /dev/null +++ b/csrc/mctlass/include/mctlass/pipeline/sm90_pipeline.hpp @@ -0,0 +1,995 @@ +/*************************************************************************************************** + * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/mctlass.h" +#include "cute/numeric/integral_constant.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/arch/copy_sm90_tma.hpp" +#include "mctlass/arch/barrier.h" + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +using namespace cute; + +enum class BarrierStatus : uint32_t { + WaitAgain = 0u, + WaitDone = 1u +}; + +class ArrivalToken { +public: + MCTLASS_HOST_DEVICE + ArrivalToken(BarrierStatus barrier_status) : barrier_status_(barrier_status) {} + + MCTLASS_HOST_DEVICE + ArrivalToken() = delete; + + MCTLASS_HOST_DEVICE + BarrierStatus get() const { + return barrier_status_;; + } + + MCTLASS_HOST_DEVICE + bool operator==(ArrivalToken const& other) const { + return barrier_status_ == other.get(); + } + +private: + BarrierStatus barrier_status_; + + MCTLASS_HOST_DEVICE + friend bool operator==(const ArrivalToken& left, const BarrierStatus& right) { + return left.get() == right; + } + + MCTLASS_HOST_DEVICE + friend bool operator==(const BarrierStatus& left, const ArrivalToken& right) { + return left == right.get(); + } +}; + +class ProducerToken : public ArrivalToken { + using ArrivalToken::ArrivalToken; +}; + +class ConsumerToken : public ArrivalToken { + using ArrivalToken::ArrivalToken; +}; + +// Circular Buffer Index + Associated Phase +// Assumes only one operation possible - i.e., ++ +template +struct PipelineState { + + static constexpr uint32_t Stages = Stages_; + +private: + int index_ = 0; + uint32_t phase_ = 0; + uint32_t phase_count_ = 0; + +public: + MCTLASS_DEVICE + PipelineState(): index_{}, phase_{}, phase_count_{} {} + + MCTLASS_DEVICE + PipelineState(int index, uint32_t phase, uint32_t phase_count) + : index_(index) + , phase_(phase) + , phase_count_(phase_count) {} + + MCTLASS_DEVICE + int index() const { + return index_; + } + + MCTLASS_DEVICE + uint32_t phase() const { + return phase_; + } + + MCTLASS_DEVICE + uint32_t phase_count() const { + return phase_count_; + } + + MCTLASS_DEVICE + void operator++() { + if constexpr (Stages > 0) { + ++index_; + if (index_ == Stages) { + index_ = 0; + phase_ ^= 1; + ++phase_count_; + } + } + } + + MCTLASS_DEVICE + PipelineState& operator=(const PipelineState& other) { + index_ = other.index(); + phase_ = other.phase(); + phase_count_ = other.phase_count(); + return *this; + } + + MCTLASS_DEVICE + PipelineState advance(uint32_t num_iterations) { + if constexpr (Stages > 0) { + // Number of iterations cross over the stage boundary => flipped phase + if ((num_iterations < Stages) && (index_ + num_iterations) >= Stages ) { + phase_ ^= 1; + } + // How many times number of iterations cross over the stage boundary and + // end up on a odd number => flipped phase + if ((num_iterations >= Stages) && (((index_ + num_iterations) / Stages) % 2) == 1) { + phase_ ^= 1; + } + phase_count_ += (index_ + num_iterations) / Stages; + index_ = (index_ + num_iterations) % Stages; + } + return *this; + } + + MCTLASS_DEVICE + static PipelineState make_pipeline_state(PipelineState start_state, uint32_t num_iterations) { + return start_state.advance(num_iterations); + } +}; + +template +MCTLASS_DEVICE +PipelineState make_producer_start_state() { + // Producer starts with an opposite phase as the buffers are initially empty + constexpr int InitialProducerStage = 0; + constexpr uint32_t InitialProducerPhase = 1; + constexpr uint32_t InitialProducerPhaseCount = 0; + return {InitialProducerStage, InitialProducerPhase, InitialProducerPhaseCount}; +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// TMA load (producer) Async Pipeline class +// +/////////////////////////////////////////////////////////////////////////////////////////////////// +// Assumptions : Constructor is visible Cluster-wide (as it needs a Cluster-Sync) +// We have exactly one thread elected in the Producer as the "leader" +// Currently, it is optional to elect a leader for the Consumers +template < + int Stages_, + class ClusterShape_ +> +class PipelineTmaAsync { +public : + using ClusterShape = ClusterShape_; + using FullBarrier = mctlass::arch::ClusterTransactionBarrier; + using EmptyBarrier = mctlass::arch::ClusterBarrier; + using ProducerBarrierType = FullBarrier::ValueType; + using ConsumerBarrierType = EmptyBarrier::ValueType; + static constexpr uint32_t Stages = Stages_; + + struct SharedStorage { + FullBarrier full_barrier_[Stages]; + EmptyBarrier empty_barrier_[Stages]; + }; + + enum class ThreadCategory { + NonParticipant, + Producer, + Consumer, + ProducerConsumer + }; + + struct Params { + uint32_t transaction_bytes = 0; + ThreadCategory role = ThreadCategory::NonParticipant; + uint32_t is_leader = 0; + uint32_t num_consumers = 0; + }; + + // Constructor + MCTLASS_DEVICE + PipelineTmaAsync(SharedStorage& storage, Params params) + : params_(params) + , full_barrier_ptr_(&storage.full_barrier_[0]) + , empty_barrier_ptr_(&storage.empty_barrier_[0]) { + + int warp_idx = canonical_warp_idx(); + int lane_predicate = cute::elect_one_sync(); + auto cluster_shape = ClusterShape{}; + + if (warp_idx == 0 && lane_predicate == 1) { + // Barrier FULL init + for (int i = 0; i < Stages; ++i) { + full_barrier_ptr_[i].init(1); + } + // Barrier EMPTY init + uint32_t const num_consumer_warpgroups_per_cluster = params_.num_consumers / NumThreadsPerWarpGroup; + uint32_t const multicast_consumer_arrival_count = (cute::size<0>(cluster_shape) + cute::size<1>(cluster_shape) - 1) * + num_consumer_warpgroups_per_cluster; + for (int i = 0; i < Stages; ++i) { + empty_barrier_ptr_[i].init(multicast_consumer_arrival_count); + } + } + + // Logic to optimally schedule Empty Arrives + // Goal : To divide SYNCS Empty Arrival duty equally amongst the Warp-Group (128 threads) + dim3 block_id = cute::block_id_in_cluster(); + auto cluster_size = cute::size(cluster_shape); + static constexpr int MaxClusterSize = 16; + static_assert(cluster_size <= MaxClusterSize, "ERROR : Cluster size too large !" ); + + // STEP 1 : Use Cute Layout function to generate an optimal dst block-id (0-15) + if (params_.num_consumers % NumThreadsPerWarpGroup == 0) { + int thread_idx = threadIdx.x % NumThreadsPerWarpGroup; + is_signalling_thread_ = (thread_idx % (NumThreadsPerWarpGroup / MaxClusterSize)) == 0; + auto layout = cute::composition(Swizzle<2,0,-2>{}, + Layout,Stride<_4,_1>>{}); + uint32_t thread_row = warp_idx % 4; + uint32_t thread_col = (thread_idx / 8) % 4; + dst_blockid_ = layout(thread_row, thread_col); + } + else if (params_.num_consumers == 32) { + int thread_idx = threadIdx.x % 32; + is_signalling_thread_ = (thread_idx % (32 / MaxClusterSize)) == 0; + auto layout = Layout,Stride<_4, _1>>{}; + uint32_t thread_row = thread_idx / 8; + uint32_t thread_col = (thread_idx % 8) / 2; + dst_blockid_ = layout(thread_row, thread_col); + } + else { + is_signalling_thread_ = 0; + #ifndef NDEBUG + asm volatile ("brkpt;\n" ::); + #endif + } + + // STEP 2: Find if this dst block-id needs an arrival for this problem + is_signalling_thread_ &= dst_blockid_ < cluster_size; + is_signalling_thread_ &= is_same_row_or_col(dst_blockid_, block_id, cluster_shape); + + mctlass::arch::fence_barrier_init(); + } + + MCTLASS_DEVICE + bool is_same_row_or_col(int dst_block_id, dim3 block_id, ClusterShape cluster_shape) { + return ((dst_block_id % cute::size<0>(cluster_shape)) == block_id.x || + (dst_block_id / cute::size<0>(cluster_shape)) == block_id.y); + } + + //////////////////// + // Producer APIs + //////////////////// + // Four member functions are always used in pairs: + // + // * producer_try_acquire and producer_acquire, and + // * consumer_try_wait and consumer_wait. + // + // The two functions with "try" in their names are called "try" functions, + // and the other two are conceptually "finalize" functions. + // The "try" function in each pair starts the process of waiting on the barrier to flip. + // It opportunistically waits for an implementation-dependent timeout. + // Whether or not the barrier has flipped yet, the try function will return a token. + // If the token indicates that the barrier has not flipped, + // then the token must be passed into the corresponding "finalize" function. + // The finalize function will then block until the barrier has flipped. + // If the token indicates that the barrier _has_ flipped, + // then it is still correct to pass it into the finalize function. + // The finalize function will return immediately in that case. + + MCTLASS_DEVICE + ProducerToken producer_try_acquire(PipelineState state, uint32_t skip_wait = false) { + return producer_try_acquire(state.index(), state.phase(), skip_wait); + } + + MCTLASS_DEVICE + void producer_acquire(PipelineState state, ProducerToken barrier_token = {BarrierStatus::WaitAgain}) { + producer_acquire(state.index(), state.phase(), barrier_token); + } + + MCTLASS_DEVICE + void producer_commit(PipelineState state, uint32_t bytes) { + producer_commit(state.index(), bytes); + } + + // Prevents early exit of producer blocks in Cluster. + // This should be called once before kernel exits. + MCTLASS_DEVICE + void producer_tail(PipelineState state) { + for (int count = 0; count < Stages; ++count) { + producer_acquire(state); + ++state; + } + } + + MCTLASS_DEVICE + ProducerBarrierType* producer_get_barrier(PipelineState state) { + return producer_get_barrier(state.index()); + } + + //////////////////// + // Consumer APIs + //////////////////// + MCTLASS_DEVICE + ConsumerToken consumer_try_wait(PipelineState state, uint32_t skip_wait = false) { + return consumer_try_wait(state.index(), state.phase(), skip_wait); + } + + MCTLASS_DEVICE + void consumer_wait(PipelineState state) { + consumer_wait(state.index(), state.phase()); + } + + MCTLASS_DEVICE + void consumer_wait(PipelineState state, ConsumerToken barrier_token) { + consumer_wait(state.index(), state.phase(), barrier_token); + } + + MCTLASS_DEVICE + void consumer_release(PipelineState state) { + consumer_release(state.index()); + } + +private : + uint32_t dst_blockid_ = 0; + uint32_t is_signalling_thread_ = 0; + FullBarrier *full_barrier_ptr_ = nullptr; + EmptyBarrier *empty_barrier_ptr_ = nullptr; + Params params_; + + MCTLASS_DEVICE + ProducerToken producer_try_acquire(uint32_t stage, uint32_t phase, uint32_t skip_wait) { + if (skip_wait) { + return {BarrierStatus::WaitDone}; + } + uint32_t barrier_status = empty_barrier_ptr_[stage].try_wait(phase); + return {static_cast(barrier_status)}; + } + + MCTLASS_DEVICE + void producer_acquire(uint32_t stage, uint32_t phase, ProducerToken barrier_token) { + if (barrier_token == BarrierStatus::WaitAgain) { + empty_barrier_ptr_[stage].wait(phase); + } + + if (params_.is_leader) { + full_barrier_ptr_[stage].arrive_and_reset_bytes(params_.transaction_bytes); + } + #ifndef NDEBUG + if (params_.role == ThreadCategory::Consumer || params_.role == ThreadCategory::NonParticipant) { + asm volatile ("brkpt;\n" ::); + } + + // Most likely you have elected more than one leader + if (params_.is_leader && (threadIdx.x % 32 != 0)) { + asm volatile ("brkpt;\n" ::); + } + #endif + } + + // NOP for TMA based mainloop + MCTLASS_DEVICE + void producer_commit(uint32_t stage, uint32_t bytes) { + // Below code is used only for unit-testing (in the absence of TMA commit) + #if MCTLASS_UNIT_TEST_PIPELINE + if (params_.is_leader) { + // STEP 1 : Commit to self + full_barrier_ptr_[stage].commit(bytes); + + // STEP 2 : Commit to other blocks in our cluster + auto cluster_shape = ClusterShape{}; + Layout block_layout_in_cluster = make_layout(cluster_shape); + dim3 local_block_id = cute::block_id_in_cluster(); + + MCTLASS_PRAGMA_UNROLL + for(int n = 0; n < size<1>(block_layout_in_cluster); ++n) { + uint32_t dst_block_id = block_layout_in_cluster(local_block_id.x,n,Int<0>{}); + full_barrier_ptr_[stage].commit(dst_block_id, bytes, n!=local_block_id.y); + } + + MCTLASS_PRAGMA_UNROLL + for(int m = 0; m < size<0>(block_layout_in_cluster); ++m) { + uint32_t dst_block_id = block_layout_in_cluster(m,local_block_id.y,Int<0>{}); + full_barrier_ptr_[stage].commit(dst_block_id, bytes, m!=local_block_id.x); + } + } + #endif + } + + MCTLASS_DEVICE + ConsumerToken consumer_try_wait(uint32_t stage, uint32_t phase, uint32_t skip_wait) { + if (skip_wait) { + return {BarrierStatus::WaitDone}; + } + uint32_t barrier_status = full_barrier_ptr_[stage].try_wait(phase); + return {static_cast(barrier_status)}; + } + + // Wait for producer to commit transactions (done by TMA) + MCTLASS_DEVICE + void consumer_wait(uint32_t stage, uint32_t phase) { + uint32_t done = full_barrier_ptr_[stage].test_wait(phase); + if (not done) { + full_barrier_ptr_[stage].wait(phase); + } + } + + // Wait for producer to commit transactions (done by TMA) + MCTLASS_DEVICE + void consumer_wait(uint32_t stage, uint32_t phase, ConsumerToken barrier_token) { + if (barrier_token == BarrierStatus::WaitAgain) { + full_barrier_ptr_[stage].wait(phase); + } + } + + // Consumer signalling Producer of completion + // Ensures all blocks in the Same Row and Column get notifed. + MCTLASS_DEVICE + void consumer_release(uint32_t stage, uint32_t skip = false) { + empty_barrier_ptr_[stage].arrive(dst_blockid_, is_signalling_thread_ & (!skip)); + #ifndef NDEBUG + if (params_.role == ThreadCategory::Producer || params_.role == ThreadCategory::NonParticipant) { + asm volatile ("brkpt;\n" ::); + } + #endif + } + + MCTLASS_DEVICE + ProducerBarrierType* producer_get_barrier(uint32_t stage) { + return reinterpret_cast(&full_barrier_ptr_[stage]); + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// TMA store (consumer) pipeline class +// producer-only class, no async barriers between threads because consumer is TMA unit +// +/////////////////////////////////////////////////////////////////////////////////////////////////// +template < + int Stages_ +> +class PipelineTmaStore { +public: + static constexpr uint32_t Stages = Stages_; + + struct Params { + bool always_wait = false; + }; + + MCTLASS_DEVICE + PipelineTmaStore(Params params = {}) : params_(params) {} + + //////////////////// + // Producer APIs + //////////////////// + // Wait for the least recently committed batch of TMA stores to complete + MCTLASS_DEVICE + void producer_acquire(PipelineState state) { + producer_acquire(state.index(), state.phase_count()); + } + + // Commit the most recently issued batch of TMA stores + MCTLASS_DEVICE + void producer_commit(PipelineState state) { + producer_commit(state.index(), state.phase_count()); + } + + // Wait for all TMA stores to complete + MCTLASS_DEVICE + void producer_tail([[maybe_unused]] PipelineState state) { + tma_store_wait<0>(); + } + +private: + Params params_; + + // Wait for the least recently committed batch of TMA stores to complete + MCTLASS_DEVICE + void producer_acquire([[maybe_unused]] uint32_t stage, uint32_t phase_count) { + if (params_.always_wait || phase_count > 0) { + tma_store_wait(); + } + } + + // Commit the most recently issued batch of TMA stores + MCTLASS_DEVICE + void producer_commit([[maybe_unused]] uint32_t stage, [[maybe_unused]] uint32_t phase_count) { + tma_store_arrive(); + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Simple producer-consumer async Pipeline class using producer transaction barriers +// +/////////////////////////////////////////////////////////////////////////////////////////////////// +template +class PipelineTransactionAsync { +public : + using FullBarrier = mctlass::arch::ClusterTransactionBarrier; + using EmptyBarrier = mctlass::arch::ClusterBarrier; + using ProducerBarrierType = FullBarrier::ValueType; + using ConsumerBarrierType = EmptyBarrier::ValueType; + static constexpr uint32_t Stages = Stages_; + + struct SharedStorage { + FullBarrier full_barrier_[Stages]; + EmptyBarrier empty_barrier_[Stages]; + }; + + enum class ThreadCategory { + NonParticipant, + Producer, + Consumer, + ProducerConsumer + }; + + struct Params { + ThreadCategory role = ThreadCategory::NonParticipant; + uint32_t transaction_bytes = 0; + uint32_t producer_arv_count = 1; + uint32_t consumer_arv_count = 1; + uint32_t dst_blockid = cute::block_rank_in_cluster(); + }; + + // Constructor + MCTLASS_DEVICE + PipelineTransactionAsync(SharedStorage& storage, Params const& params) + : params_(params) + , full_barrier_ptr_(&storage.full_barrier_[0]) + , empty_barrier_ptr_(&storage.empty_barrier_[0]) { + + int warp_idx = canonical_warp_idx(); + int lane_predicate = cute::elect_one_sync(); + + // Barrier FULL, EMPTY init + // Init is done only by thread 0 of the block + if (warp_idx == 0 && lane_predicate == 1) { + for (int i = 0; i < Stages; ++i) { + full_barrier_ptr_[i].init(params.producer_arv_count); + empty_barrier_ptr_[i].init(params.consumer_arv_count); + } + } + + mctlass::arch::fence_barrier_init(); + } + + //////////////////// + // Producer APIs + //////////////////// + // Four member functions are always used in pairs: + // + // * producer_try_acquire and producer_acquire, and + // * consumer_try_wait and consumer_wait. + // + // The two functions with "try" in their names are called "try" functions, + // and the other two are conceptually "finalize" functions. + // The "try" function in each pair starts the process of waiting on the barrier to flip. + // It opportunistically waits for an implementation-dependent timeout. + // Whether or not the barrier has flipped yet, the try function will return a token. + // If the token indicates that the barrier has not flipped, + // then the token must be passed into the corresponding "finalize" function. + // The finalize function will then block until the barrier has flipped. + // If the token indicates that the barrier _has_ flipped, + // then it is still correct to pass it into the finalize function. + // The finalize function will return immediately in that case. + MCTLASS_DEVICE + ProducerToken producer_try_acquire(PipelineState state, uint32_t skip_wait = false) { + return producer_try_acquire(state.index(), state.phase(), skip_wait); + } + + MCTLASS_DEVICE + void producer_acquire(PipelineState state, ProducerToken barrier_token = {BarrierStatus::WaitAgain}) { + producer_acquire(state.index(), state.phase(), barrier_token); + } + + MCTLASS_DEVICE + void producer_commit(PipelineState state) { + producer_commit(state.index()); + } + + // Prevents early exit of producer blocks in Cluster. + // This should be called once before kernel exits. + MCTLASS_DEVICE + void producer_tail(PipelineState state) { + for (int count = 0; count < Stages; ++count) { + producer_acquire(state); + ++state; + } + } + + MCTLASS_DEVICE + ProducerBarrierType* producer_get_barrier(PipelineState state) { + return producer_get_barrier(state.index()); + } + + //////////////////// + // Consumer APIs + //////////////////// + MCTLASS_DEVICE + ConsumerToken consumer_try_wait(PipelineState state, uint32_t skip_wait = false) { + return consumer_try_wait(state.index(), state.phase(), skip_wait); + } + + MCTLASS_DEVICE + void consumer_wait(PipelineState state, ConsumerToken barrier_token = {BarrierStatus::WaitAgain}) { + consumer_wait(state.index(), state.phase(), barrier_token); + } + + MCTLASS_DEVICE + void consumer_release(PipelineState state) { + consumer_release(state.index()); + } + +private: + FullBarrier *full_barrier_ptr_ = nullptr; + EmptyBarrier *empty_barrier_ptr_ = nullptr; + Params params_; + + MCTLASS_DEVICE + ProducerToken producer_try_acquire(uint32_t stage, uint32_t phase, uint32_t skip_wait) { + if (skip_wait) { + return {BarrierStatus::WaitDone}; + } + uint32_t barrier_status = empty_barrier_ptr_[stage].try_wait(phase); + return {static_cast(barrier_status)}; + } + + MCTLASS_DEVICE + void producer_acquire(uint32_t stage, uint32_t phase, ProducerToken barrier_token) { + if (barrier_token == BarrierStatus::WaitAgain) { + empty_barrier_ptr_[stage].wait(phase); + } + + full_barrier_ptr_[stage].arrive_and_reset_bytes(params_.transaction_bytes, params_.dst_blockid); + } + + MCTLASS_DEVICE + void producer_commit([[maybe_unused]] uint32_t stage) { + } + + MCTLASS_DEVICE + ProducerBarrierType* producer_get_barrier(uint32_t stage) { + return reinterpret_cast(&full_barrier_ptr_[stage]); + } + + MCTLASS_DEVICE + ConsumerToken consumer_try_wait(uint32_t stage, uint32_t phase, uint32_t skip_wait) { + if (skip_wait) { + return {BarrierStatus::WaitDone}; + } + uint32_t barrier_status = full_barrier_ptr_[stage].try_wait(phase); + return {static_cast(barrier_status)}; + } + + MCTLASS_DEVICE + void consumer_wait(uint32_t stage, uint32_t phase, ConsumerToken barrier_token) { + if (barrier_token == BarrierStatus::WaitAgain) { + full_barrier_ptr_[stage].wait(phase); + } + } + + MCTLASS_DEVICE + void consumer_release(uint32_t stage, uint32_t skip = false) { + empty_barrier_ptr_[stage].arrive(params_.dst_blockid, (not skip)); + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Simple producer-consumer async Pipeline class +// +/////////////////////////////////////////////////////////////////////////////////////////////////// +template +class PipelineAsync { +public : + using FullBarrier = mctlass::arch::ClusterBarrier; + using EmptyBarrier = mctlass::arch::ClusterBarrier; + using ProducerBarrierType = FullBarrier::ValueType; + using ConsumerBarrierType = EmptyBarrier::ValueType; + static constexpr uint32_t Stages = Stages_; + + struct SharedStorage { + FullBarrier full_barrier_[Stages]; + EmptyBarrier empty_barrier_[Stages]; + }; + + enum class ThreadCategory { + NonParticipant, + Producer, + Consumer, + ProducerConsumer + }; + + struct Params { + ThreadCategory role = ThreadCategory::NonParticipant; + uint32_t producer_arv_count = 1; + uint32_t consumer_arv_count = 1; + uint32_t dst_blockid = cute::block_rank_in_cluster(); + }; + + // Default assumption when only storage is passed is : + // => single producer, single consumer & they are in the same block (within the Cluster) + MCTLASS_DEVICE + PipelineAsync(SharedStorage& storage) + : PipelineAsync(storage, {}) {} + + MCTLASS_DEVICE + PipelineAsync( + SharedStorage& storage, + Params const& params) : + params_(params), + full_barrier_ptr_(&storage.full_barrier_[0]), + empty_barrier_ptr_(&storage.empty_barrier_[0]) { + + int warp_idx = canonical_warp_idx(); + int lane_predicate = cute::elect_one_sync(); + + // Barrier FULL, EMPTY init + // Init is done only by thread 0 of the block + if (warp_idx == 0 && lane_predicate == 1) { + for (int i = 0; i < Stages; ++i) { + full_barrier_ptr_[i].init(params.producer_arv_count); + empty_barrier_ptr_[i].init(params.consumer_arv_count); + } + } + + mctlass::arch::fence_barrier_init(); + } + + //////////////////// + // Producer APIs + //////////////////// + // Four member functions are always used in pairs: + // + // * producer_try_acquire and producer_acquire, and + // * consumer_try_wait and consumer_wait. + // + // The two functions with "try" in their names are called "try" functions, + // and the other two are conceptually "finalize" functions. + // The "try" function in each pair starts the process of waiting on the barrier to flip. + // It opportunistically waits for an implementation-dependent timeout. + // Whether or not the barrier has flipped yet, the try function will return a token. + // If the token indicates that the barrier has not flipped, + // then the token must be passed into the corresponding "finalize" function. + // The finalize function will then block until the barrier has flipped. + // If the token indicates that the barrier _has_ flipped, + // then it is still correct to pass it into the finalize function. + // The finalize function will return immediately in that case. + MCTLASS_DEVICE + ProducerToken producer_try_acquire(PipelineState state, uint32_t skip_wait = false) { + return producer_try_acquire(state.index(), state.phase(), skip_wait); + } + + MCTLASS_DEVICE + void producer_acquire(PipelineState state, ProducerToken barrier_token = {BarrierStatus::WaitAgain}) { + producer_acquire(state.index(), state.phase(), barrier_token); + } + + MCTLASS_DEVICE + void producer_commit(PipelineState state) { + producer_commit(state.index()); + } + + // Prevents early exit of producer blocks in Cluster. + // This should be called once before kernel exits. + MCTLASS_DEVICE + void producer_tail(PipelineState state) { + for (int count = 0; count < Stages; ++count) { + producer_acquire(state); + ++state; + } + } + + MCTLASS_DEVICE + ProducerBarrierType* producer_get_barrier(PipelineState state) { + return producer_get_barrier(state.index()); + } + + //////////////////// + // Consumer APIs + //////////////////// + MCTLASS_DEVICE + ConsumerToken consumer_try_wait(PipelineState state, uint32_t skip_wait = false) { + return consumer_try_wait(state.index(), state.phase(), skip_wait); + } + + MCTLASS_DEVICE + void consumer_wait(PipelineState state, ConsumerToken barrier_token = {BarrierStatus::WaitAgain}) { + consumer_wait(state.index(), state.phase(), barrier_token); + } + + MCTLASS_DEVICE + void consumer_release(PipelineState state) { + consumer_release(state.index()); + } + +private: + Params params_; + FullBarrier *full_barrier_ptr_; + EmptyBarrier *empty_barrier_ptr_; + + MCTLASS_DEVICE + ProducerToken producer_try_acquire(uint32_t stage, uint32_t phase, uint32_t skip_wait) { + if (skip_wait) { + return {BarrierStatus::WaitDone}; + } + uint32_t barrier_status = empty_barrier_ptr_[stage].try_wait(phase); + return {static_cast(barrier_status)}; + } + + MCTLASS_DEVICE + void producer_acquire(uint32_t stage, uint32_t phase, ProducerToken barrier_token) { + if (barrier_token == BarrierStatus::WaitAgain) { + empty_barrier_ptr_[stage].wait(phase); + } + } + + MCTLASS_DEVICE + void producer_commit(uint32_t stage) { + full_barrier_ptr_[stage].arrive(); + } + + MCTLASS_DEVICE + ProducerBarrierType* producer_get_barrier(uint32_t stage) { + return reinterpret_cast(&full_barrier_ptr_[stage]); + } + + MCTLASS_DEVICE + ConsumerToken consumer_try_wait(uint32_t stage, uint32_t phase, uint32_t skip_wait) { + if (skip_wait) { + return {BarrierStatus::WaitDone}; + } + uint32_t barrier_status = full_barrier_ptr_[stage].try_wait(phase); + return {static_cast(barrier_status)}; + } + + MCTLASS_DEVICE + void consumer_wait(uint32_t stage, uint32_t phase) { + uint32_t done = full_barrier_ptr_[stage].test_wait(phase); + if (!done) { + full_barrier_ptr_[stage].wait(phase); + } + } + + MCTLASS_DEVICE + void consumer_wait(uint32_t stage, uint32_t phase, ConsumerToken barrier_token) { + if (barrier_token == BarrierStatus::WaitAgain) { + full_barrier_ptr_[stage].wait(phase); + } + } + + MCTLASS_DEVICE + void consumer_release(uint32_t stage) { + empty_barrier_ptr_[stage].arrive(params_.dst_blockid); + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Barrier to ensure an Ordered Sequence between +// SequenceLength number of groups (each with group_size participants) executing SequenceDepth Stages +// i.e., for all i < j - only after id "i" arrives at a particular stage "m" +// will the wait() for id "j" succeed for the same stage +// +/////////////////////////////////////////////////////////////////////////////////////////////////// + +template +class OrderedSequenceBarrier { +public : + using Barrier = mctlass::arch::ClusterBarrier; + + struct SharedStorage { + Barrier barrier_[SequenceDepth][SequenceLength]; + }; + + struct Params { + uint32_t group_id; + uint32_t group_size; + }; + +private : + // In future this Params object can be replaced easily with a CG object + Params params_; + Barrier *barrier_ptr_; + PipelineState stage_; + + static constexpr int Depth = SequenceDepth; + static constexpr int Length = SequenceLength; + +public: + OrderedSequenceBarrier() = delete; + OrderedSequenceBarrier(const OrderedSequenceBarrier&) = delete; + OrderedSequenceBarrier(OrderedSequenceBarrier&&) = delete; + OrderedSequenceBarrier& operator=(const OrderedSequenceBarrier&) = delete; + OrderedSequenceBarrier& operator=(OrderedSequenceBarrier&&) = delete; + ~OrderedSequenceBarrier() = default; + + MCTLASS_DEVICE + OrderedSequenceBarrier(SharedStorage& storage, Params const& params) : + params_(params), + barrier_ptr_(&storage.barrier_[0][0]), + // Group 0 - starts with an opposite phase + stage_({0, params.group_id == 0, 0}) { + + int warp_idx = canonical_warp_idx(); + int lane_predicate = cute::elect_one_sync(); + + // Barrier FULL, EMPTY init + // Init is done only by the one elected thread of the block + if (warp_idx == 0 && lane_predicate == 1) { + for (int d = 0; d < Depth; ++d) { + for (int l = 0; l < Length; ++l) { + barrier_ptr_[d * Length + l].init(params.group_size); + } + } + } + + mctlass::arch::fence_barrier_init(); + } + + // Wait on a stage to be unlocked + MCTLASS_DEVICE + void wait() { + get_barrier_for_current_stage(params_.group_id).wait(stage_.phase()); + } + + // Signal completion of Stage and move to the next stage + // (group_id) signals to (group_id+1) + MCTLASS_DEVICE + void arrive() { + int signalling_id = (params_.group_id + 1) % Length; + get_barrier_for_current_stage(signalling_id).arrive(); + ++stage_; + } + + MCTLASS_DEVICE + void advance() { + ++stage_; + } + +private: + + MCTLASS_DEVICE + Barrier& get_barrier_for_current_stage(int group_id) { + return barrier_ptr_[stage_.index() * Length + group_id]; + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // end namespace mctlass diff --git a/csrc/mctlass/include/mctlass/pitch_linear_coord.h b/csrc/mctlass/include/mctlass/pitch_linear_coord.h new file mode 100644 index 0000000..c57b228 --- /dev/null +++ b/csrc/mctlass/include/mctlass/pitch_linear_coord.h @@ -0,0 +1,181 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines layout functions used by TensorRef and derived classes for pitch-linear memory. +*/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/coord.h" + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Template defining a shape used by pitch-linear operators +template < + int Contiguous, + int Strided +> +struct PitchLinearShape { + static int const kContiguous = Contiguous; + static int const kStrided = Strided; + static int const kCount = Contiguous * Strided; +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Coordinate in pitch-linear space +struct PitchLinearCoord : public Coord<2, int> { +public: + + /// Integer-valued index + using Index = int; + + /// Base type is a Coord of rank=2 + using Base = Coord<2, Index>; + + /// Long integer type + using LongIndex = typename Base::LongIndex; + +private: + + /// Rows dimension + static int const kContiguous = 0; + + /// Columns dimension + static int const kStrided = 1; + +public: + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + PitchLinearCoord() { } + + /// Constructs from Coord<2> + MCTLASS_HOST_DEVICE + PitchLinearCoord(Coord<2, Index> const &coord): Base(coord) { } + + /// Helper to construct from a row and column + MCTLASS_HOST_DEVICE + PitchLinearCoord(Index contiguous_, Index strided_): Base(make_Coord(contiguous_, strided_)) { } + + /// Helper to construct from a row and column based on LongIndex + MCTLASS_HOST_DEVICE + PitchLinearCoord(LongIndex contiguous_, LongIndex strided_) + : Base(make_Coord(Index(contiguous_), Index(strided_))) { } + + /// Returns the contiguous dimension + MCTLASS_HOST_DEVICE + Index const & contiguous() const { return this->at(kContiguous); } + + /// Returns the contiguous dimension + MCTLASS_HOST_DEVICE + Index & contiguous() { return this->at(kContiguous); } + + /// Returns the column of the coordinate + MCTLASS_HOST_DEVICE + Index const & strided() const { return this->at(kStrided); } + + /// Returns the column of the coordinate + MCTLASS_HOST_DEVICE + Index & strided() { return this->at(kStrided); } + + // + // Coord operators + // + + /// Element-wise addition + MCTLASS_HOST_DEVICE + PitchLinearCoord operator+(Base const& b) const { + return PitchLinearCoord(Base::operator+(b)); + } + + /// Element-wise subtraction + MCTLASS_HOST_DEVICE + PitchLinearCoord operator-(Base const& b) const { + return PitchLinearCoord(Base::operator-(b)); + } + + MCTLASS_HOST_DEVICE + PitchLinearCoord operator-() const { + return PitchLinearCoord(-at(0), -at(1)); + } + + /// Element-wise multiplication + MCTLASS_HOST_DEVICE + PitchLinearCoord operator*(Base const& b) const { + return PitchLinearCoord(Base::operator*(b)); + } + + /// Element-wise division + MCTLASS_HOST_DEVICE + PitchLinearCoord operator/(Base const& b) const { + return PitchLinearCoord(Base::operator/(b)); + } + + /// In-place addition + MCTLASS_HOST_DEVICE + PitchLinearCoord& operator+=(Base const& b) { + Base::operator+=(b); + return *this; + } + + /// In-place subtraction + MCTLASS_HOST_DEVICE + PitchLinearCoord& operator-=(Base const& b) { + Base::operator-=(b); + return *this; + } + + /// In-place multiplication + MCTLASS_HOST_DEVICE + PitchLinearCoord& operator*=(Base const& b) { + Base::operator*=(b); + return *this; + } + + /// In-place division + MCTLASS_HOST_DEVICE + PitchLinearCoord& operator/=(Base const& b) { + Base::operator/=(b); + return *this; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/platform/platform.h b/csrc/mctlass/include/mctlass/platform/platform.h new file mode 100644 index 0000000..01b8a24 --- /dev/null +++ b/csrc/mctlass/include/mctlass/platform/platform.h @@ -0,0 +1,883 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +/** + * \file + * \brief C++ features that may be otherwise unimplemented for CUDA device functions. + * + * This file has three components: + * + * (1) Macros: + * - Empty macro defines for C++ keywords not supported by the current + * version of C++. These simply allow compilation to proceed (but do + * not provide the added semantics). + * - \p noexcept + * - \p constexpr + * - \p nullptr + * - \p static_assert + * + * - Macro functions that we need in constant expressions because the + * C++ equivalents require constexpr compiler support. These are + * prefixed with \p __NV_STD_* + * - \p __NV_STD_MAX + * - \p __NV_STD_MIN + * + * (2) Re-implementations of STL functions and types: + * - C++ features that need the \p __device__ annotation. These are + * placed into the \p platform namespace. + * - \p abs + * - \p plus + * - \p less + * - \p greater + * - \p min + * - \p max + * - \p methods on std::pair (==, !=, <, <=, >, >=, and make_pair()) + * + * (3) Stop-gap implementations of unsupported STL functions and types: + * - STL functions and types defined by C++ 11/14/17/etc. that are not + * provided by the current version of C++. These are placed into the + * \p platform namespace + * - \p integral_constant + * - \p nullptr_t + * - \p true_type + * - \p false_type + * - \p bool_constant + * - \p enable_if + * - \p conditional + * - \p is_same + * - \p is_base_of + * - \p remove_const + * - \p remove_volatile + * - \p remove_cv + * - \p is_volatile + * - \p is_pointer + * - \p is_void + * - \p is_integral + * - \p is_floating_point + * - \p is_arithmetic + * - \p is_fundamental + +//----------------------------------------------------------------------------- +// Dependencies +//----------------------------------------------------------------------------- + +#if defined(__MACACC_RTC__) +#include +#else +#include +#endif + +#if !defined(__MACACC_RTC__) +//----------------------------------------------------------------------------- +// Include STL files that platform provides functionality for +//----------------------------------------------------------------------------- + +#include // Minimum/maximum operations +#include // nullptr_t +#include // Arithmetic operations +#include // For methods on std::pair +#if (!defined(_MSC_VER) && (__cplusplus >= 201103L)) || (defined(_MSC_VER) && (_MS_VER >= 1500)) +#include // For integral constants, conditional metaprogramming, and type traits +#endif + +#include "mctlass/mctlass.h" + +#endif + +//----------------------------------------------------------------------------- +// OS +//----------------------------------------------------------------------------- +#if defined(WIN32) || defined(_WIN32) || defined(__WIN32) && !defined(__CYGWIN__) +#define MCTLASS_OS_WINDOWS +#endif + +/****************************************************************************** + * Macros + ******************************************************************************/ +//----------------------------------------------------------------------------- +// Keywords +//----------------------------------------------------------------------------- + +/// noexcept, constexpr +#if (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1900)) +#ifndef noexcept +#define noexcept +#endif +#ifndef constexpr +#define constexpr +#endif +#endif + +/// nullptr +#if (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1310)) +#ifndef nullptr +#define nullptr 0 +#endif +#endif + +/// static_assert +#if (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1600)) +#ifndef static_assert +#define __platform_cat_(a, b) a##b +#define __platform_cat(a, b) __platform_cat_(a, b) +#define static_assert(__e, __m) typedef int __platform_cat(AsSeRt, __LINE__)[(__e) ? 1 : -1] +#endif +#endif + +//----------------------------------------------------------------------------- +// Functions +//----------------------------------------------------------------------------- + +/// Select maximum(a, b) +#ifndef __NV_STD_MAX +#define __NV_STD_MAX(a, b) (((b) > (a)) ? (b) : (a)) +#endif + +/// Select minimum(a, b) +#ifndef __NV_STD_MIN +#define __NV_STD_MIN(a, b) (((b) < (a)) ? (b) : (a)) +#endif + +/****************************************************************************** + * Re-implementations + ******************************************************************************/ +namespace mctlass { +namespace platform { + +//----------------------------------------------------------------------------- +// Abs operations +//----------------------------------------------------------------------------- + +#if defined(__MACACC_RTC__) +/// std::abs +MCTLASS_HOST_DEVICE constexpr int abs(int a) { + return (a < 0) ? -a : a; +} +MCTLASS_HOST_DEVICE constexpr long long abs(long long a) { + return (a < 0) ? -a : a; +} +#else +using std::abs; +#endif + +//----------------------------------------------------------------------------- +// Minimum/maximum operations +//----------------------------------------------------------------------------- + +/// std::min +template +MCTLASS_HOST_DEVICE constexpr const T& min(const T& a, const T& b) { + return (b < a) ? b : a; +} + +/// std::max +template +MCTLASS_HOST_DEVICE constexpr const T& max(const T& a, const T& b) { + return (a < b) ? b : a; +} + +#if !defined(__MACACC_RTC__) +//----------------------------------------------------------------------------- +// Methods on std::pair +//----------------------------------------------------------------------------- + +using std::pair; + +template +MCTLASS_HOST_DEVICE constexpr bool operator==(const pair& lhs, const pair& rhs) { + return (lhs.first == rhs.first) && (lhs.second == rhs.second); +} + +template +MCTLASS_HOST_DEVICE constexpr bool operator!=(const pair& lhs, const pair& rhs) { + return (lhs.first != rhs.first) && (lhs.second != rhs.second); +} + +template +MCTLASS_HOST_DEVICE constexpr bool operator<(const pair& lhs, const pair& rhs) { + return (lhs.first < rhs.first) ? true : (rhs.first < lhs.first) ? false + : (lhs.second < rhs.second); +} + +template +MCTLASS_HOST_DEVICE constexpr bool operator<=(const pair& lhs, const pair& rhs) { + return !(rhs < lhs); +} + +template +MCTLASS_HOST_DEVICE constexpr bool operator>(const pair& lhs, const pair& rhs) { + return (rhs < lhs); +} + +template +MCTLASS_HOST_DEVICE constexpr bool operator>=(const pair& lhs, const pair& rhs) { + return !(lhs < rhs); +} + +template +MCTLASS_HOST_DEVICE std::pair make_pair(T1 t, T2 u) { + std::pair retval; + retval.first = t; + retval.second = u; + return retval; +} +#endif + +} // namespace platform + +/****************************************************************************** + * Implementations of C++ 11/14/17/... STL features + ******************************************************************************/ + +namespace platform { + +//----------------------------------------------------------------------------- +// Integral constant helper types +//----------------------------------------------------------------------------- + +#if defined(__MACACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1500)) + +/// std::integral_constant +template +struct integral_constant; + +/// std::integral_constant +template +struct integral_constant { + static const value_t value = V; + + typedef value_t value_type; + typedef integral_constant type; + + MCTLASS_HOST_DEVICE operator value_type() const { return value; } + + MCTLASS_HOST_DEVICE const value_type operator()() const { return value; } +}; + +#else + +using std::integral_constant; +using std::pair; + +#endif + +/// The type used as a compile-time boolean with true value. +typedef integral_constant true_type; + +/// The type used as a compile-time boolean with false value. +typedef integral_constant false_type; + +#if defined(__MACACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus <= 201402L)) || (defined(_MSC_VER) && (_MSC_VER < 1900)) + +/// std::bool_constant +template +struct bool_constant : platform::integral_constant {}; + +#else + +using std::bool_constant; + +#endif + +#if defined(__MACACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1700)) + +/// std::nullptr_t +struct nullptr_t {}; + +#else + +using std::nullptr_t; + +#endif + +//----------------------------------------------------------------------------- +// Conditional metaprogramming +//----------------------------------------------------------------------------- + +#if defined(__MACACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1600)) + +/// std::enable_if (true specialization) +template +struct enable_if { + typedef T type; +}; + +/// std::enable_if (false specialization) +template +struct enable_if {}; + +/// std::conditional (true specialization) +template +struct conditional { + typedef T type; +}; + +/// std::conditional (false specialization) +template +struct conditional { + typedef F type; +}; + +#else + +using std::enable_if; +using std::conditional; + +#endif + +//----------------------------------------------------------------------------- +// Const/volatility specifiers +//----------------------------------------------------------------------------- + +#if defined(__MACACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1500)) + +/// std::remove_const (non-const specialization) +template +struct remove_const { + typedef T type; +}; + +/// std::remove_const (const specialization) +template +struct remove_const { + typedef T type; +}; + +/// std::remove_volatile (non-volatile specialization) +template +struct remove_volatile { + typedef T type; +}; + +/// std::remove_volatile (volatile specialization) +template +struct remove_volatile { + typedef T type; +}; + +/// std::remove_cv +template +struct remove_cv { + typedef typename remove_volatile::type>::type type; +}; + +#else + +using std::remove_const; +using std::remove_volatile; +using std::remove_cv; + +#endif + +//----------------------------------------------------------------------------- +// Type relationships +//----------------------------------------------------------------------------- + +#if defined(__MACACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1500)) + +/// std::is_same (false specialization) +template +struct is_same : false_type {}; + +/// std::is_same (true specialization) +template +struct is_same : true_type {}; + +/// Helper for std::is_base_of +template +struct is_base_of_helper { + typedef char (&yes)[1]; + typedef char (&no)[2]; + + template + struct dummy { + MCTLASS_HOST_DEVICE operator B*() const; + MCTLASS_HOST_DEVICE operator D*(); + }; + + template + MCTLASS_HOST_DEVICE static yes check(DerivedT*, T); + + MCTLASS_HOST_DEVICE static no check(BaseT*, int); + + static const bool value = sizeof(check(dummy(), int())) == sizeof(yes); +}; + +/// std::is_base_of +template +struct is_base_of + : integral_constant::type, + typename remove_cv::type>::value) || + (is_same::type, + typename remove_cv::type>::value)> {}; + +#else + +using std::is_same; +using std::is_base_of; + +#endif + +//----------------------------------------------------------------------------- +// Type properties +//----------------------------------------------------------------------------- + +#if defined(__MACACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1500)) + +/// std::is_volatile +template +struct is_volatile : false_type {}; +template +struct is_volatile : true_type {}; + +/// Helper for std::is_pointer (false specialization) +template +struct is_pointer_helper : false_type {}; + +/// Helper for std::is_pointer (true specialization) +template +struct is_pointer_helper : true_type {}; + +/// std::is_pointer +template +struct is_pointer : is_pointer_helper::type> {}; + +/// std::is_void +template +struct is_void : is_same::type> {}; + +/// std::is_integral +template +struct is_integral : false_type {}; +template <> +struct is_integral : true_type {}; +template <> +struct is_integral : true_type {}; +template <> +struct is_integral : true_type {}; +template <> +struct is_integral : true_type {}; +template <> +struct is_integral : true_type {}; +template <> +struct is_integral : true_type {}; +template <> +struct is_integral : true_type {}; +template <> +struct is_integral : true_type {}; +template <> +struct is_integral : true_type {}; +template <> +struct is_integral : true_type {}; +template <> +struct is_integral : true_type {}; +template +struct is_integral : is_integral {}; +template +struct is_integral : is_integral {}; +template +struct is_integral : is_integral {}; + +/// std::is_floating_point +template +struct is_floating_point + : integral_constant::type>::value || + is_same::type>::value)> {}; + +/// std::is_arithmetic +template +struct is_arithmetic + : integral_constant::value || is_floating_point::value)> {}; + +/// std::is_fundamental +template +struct is_fundamental + : integral_constant::value || is_void::value || + is_same::type>::value)> {}; + +#else + +using std::is_volatile; +using std::is_pointer; +using std::is_void; +using std::is_integral; +using std::is_floating_point; +using std::is_arithmetic; +using std::is_fundamental; + +#endif + +#if defined(__MACACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1800)) || \ + (defined(__GNUG__) && (__GNUC__ < 5)) + +/** + * std::is_trivially_copyable + * + * This implementation only evaluates true if T is fundamental or pointer + * + * Without help from partial template specializations provided by the user for + * a specific class or struct, this trait will never report that the specified + * class or struct is trivially-copyable ; this is always safe, + * if possibly sub-optimal. + */ +template +struct is_trivially_copyable + : integral_constant::value || is_pointer::value)> {}; + +#else + +using std::is_trivially_copyable; + +#endif + +//----------------------------------------------------------------------------- +// bit_cast +//----------------------------------------------------------------------------- + +template< class To, class From > +constexpr To MCTLASS_HOST_DEVICE bit_cast(const From& from ) noexcept; + +template +constexpr To MCTLASS_HOST_DEVICE bit_cast(const From& src) noexcept +{ + static_assert(sizeof(To) == sizeof(From), "sizes must match"); + return reinterpret_cast(src); +} + + +//----------------------------------------------------------------------------- +// Alignment and layout utilities +//----------------------------------------------------------------------------- + +#if defined(__MACACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1500)) + +/// std::alignment_of +template +struct alignment_of { + struct pad { + value_t val; + char byte; + }; + + enum { value = sizeof(pad) - sizeof(value_t) }; +}; + +#else + +template +struct alignment_of : std::alignment_of {}; + +#endif + +/* 16B specializations where 32-bit Win32 host compiler disagrees with device compiler */ +template <> +struct alignment_of { + enum { value = 16 }; +}; +template <> +struct alignment_of { + enum { value = 16 }; +}; +template <> +struct alignment_of { + enum { value = 16 }; +}; +template <> +struct alignment_of { + enum { value = 16 }; +}; +template <> +struct alignment_of { + enum { value = 16 }; +}; +template <> +struct alignment_of { + enum { value = 16 }; +}; +template <> +struct alignment_of { + enum { value = 16 }; +}; +template <> +struct alignment_of { + enum { value = 16 }; +}; +template <> +struct alignment_of { + enum { value = 16 }; +}; +template <> +struct alignment_of { + enum { value = 16 }; +}; +template <> +struct alignment_of { + enum { value = 16 }; +}; + +// Specializations for volatile/const qualified types +template +struct alignment_of : alignment_of {}; +template +struct alignment_of : alignment_of {}; +template +struct alignment_of : alignment_of {}; + +#if defined(__MACACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1800)) + +template +struct aligned_chunk; +template <> +struct __align__(1) aligned_chunk<1> { + uint8_t buff; +}; +template <> +struct __align__(2) aligned_chunk<2> { + uint16_t buff; +}; +template <> +struct __align__(4) aligned_chunk<4> { + uint32_t buff; +}; +template <> +struct __align__(8) aligned_chunk<8> { + uint32_t buff[2]; +}; +template <> +struct __align__(16) aligned_chunk<16> { + uint32_t buff[4]; +}; +template <> +struct __align__(32) aligned_chunk<32> { + uint32_t buff[8]; +}; +template <> +struct __align__(64) aligned_chunk<64> { + uint32_t buff[16]; +}; +template <> +struct __align__(128) aligned_chunk<128> { + uint32_t buff[32]; +}; +template <> +struct __align__(256) aligned_chunk<256> { + uint32_t buff[64]; +}; +template <> +struct __align__(512) aligned_chunk<512> { + uint32_t buff[128]; +}; +template <> +struct __align__(1024) aligned_chunk<1024> { + uint32_t buff[256]; +}; +template <> +struct __align__(2048) aligned_chunk<2048> { + uint32_t buff[512]; +}; +template <> +struct __align__(4096) aligned_chunk<4096> { + uint32_t buff[1024]; +}; + +/// std::aligned_storage +template +struct aligned_storage { + typedef aligned_chunk type[Len / sizeof(aligned_chunk)]; +}; + +#else + +using std::aligned_storage; + +#endif + +#if !defined(__MACACC_RTC__) +/// Default deleter +template +struct default_delete { + void operator()(T* ptr) const { delete ptr; } +}; + +/// Partial specialization for deleting array types +template +struct default_delete { + void operator()(T* ptr) const { delete[] ptr; } +}; + +/// std::unique_ptr +template > +class unique_ptr { + public: + typedef T* pointer; + typedef T element_type; + typedef Deleter deleter_type; + + private: + /// Pointer to memory + pointer _ptr; + + /// Deleter + deleter_type _deleter; + + public: + unique_ptr() : _ptr(nullptr) {} + unique_ptr(pointer p) : _ptr(p) {} + + ~unique_ptr() { + if (_ptr) { + _deleter(_ptr); + } + } + /// Returns a pointer to the managed object or nullptr if no object is owned. + pointer get() const noexcept { return _ptr; } + + /// Releases ownership of the managed object, if any + pointer release() noexcept { + pointer p(_ptr); + _ptr = nullptr; + return p; + } + + /// Replaces the managed object, deleting the old object. + void reset(pointer p = pointer()) noexcept { + pointer old_ptr = _ptr; + _ptr = p; + if (old_ptr != nullptr) { + get_deleter()(old_ptr); + } + } + + /// Swaps the managed objects with *this and another unique_ptr + void swap(unique_ptr& other) noexcept { std::swap(_ptr, other._ptr); } + + /// Returns the deleter object + Deleter& get_deleter() noexcept { return _deleter; } + + /// Returns the deleter object + Deleter const& get_deleter() const noexcept { return _deleter; } + + /// Checks whether an object is owned + operator bool() const noexcept { return _ptr != nullptr; } + + /// Dereferences the unique_ptr + T& operator*() const { return *_ptr; } + + /// Returns a pointer to the managed object + pointer operator->() const noexcept { return _ptr; } + + /// Array access to managed object + T& operator[](size_t i) const { return _ptr[i]; } +}; + +/// Specializes the swap algorithm +template +void swap(unique_ptr& lhs, unique_ptr& rhs) noexcept { + lhs.swap(rhs); +} +#endif + +/// std::numeric_limits +template +struct numeric_limits; + +template <> +struct numeric_limits { + MCTLASS_HOST_DEVICE + static constexpr int32_t lowest() noexcept { return -2147483647 - 1;} + MCTLASS_HOST_DEVICE + static constexpr int32_t max() noexcept { return 2147483647;} + static constexpr bool is_integer = true; +}; + +template <> +struct numeric_limits { + MCTLASS_HOST_DEVICE + static constexpr int16_t lowest() noexcept { return -32768;} + MCTLASS_HOST_DEVICE + static constexpr int16_t max() noexcept { return 32767;} + static constexpr bool is_integer = true; +}; + +template <> +struct numeric_limits { + MCTLASS_HOST_DEVICE + static constexpr int8_t lowest() noexcept { return -128;} + MCTLASS_HOST_DEVICE + static constexpr int8_t max() noexcept { return 127;} + static constexpr bool is_integer = true; +}; + + +template <> +struct numeric_limits { + MCTLASS_HOST_DEVICE + static constexpr uint32_t lowest() noexcept { return 0;} + MCTLASS_HOST_DEVICE + static constexpr uint32_t max() noexcept { return 4294967295U;} + static constexpr bool is_integer = true; +}; + +template <> +struct numeric_limits { + MCTLASS_HOST_DEVICE + static constexpr uint16_t lowest() noexcept { return 0;} + MCTLASS_HOST_DEVICE + static constexpr uint16_t max() noexcept { return 65535U;} + static constexpr bool is_integer = true; +}; + +template <> +struct numeric_limits { + MCTLASS_HOST_DEVICE + static constexpr uint8_t lowest() noexcept { return 0;} + MCTLASS_HOST_DEVICE + static constexpr uint8_t max() noexcept { return 255U;} + static constexpr bool is_integer = true; +}; + +#if !defined(__MACACC_RTC__) +template <> +struct numeric_limits { + MCTLASS_HOST_DEVICE + static constexpr float infinity() noexcept { return bit_cast(0x7f800000);} + static constexpr bool is_integer = false; + static constexpr bool has_infinity = true; +}; +#endif + +} // namespace platform +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/predicate_vector.h b/csrc/mctlass/include/mctlass/predicate_vector.h new file mode 100644 index 0000000..82e3398 --- /dev/null +++ b/csrc/mctlass/include/mctlass/predicate_vector.h @@ -0,0 +1,524 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines container classes and iterators for managing a statically sized vector + of boolean predicates. +*/ +#pragma once + +#if defined(__MACACC_RTC__) +#include +#include +#else +#include +#include +#endif + +#include "mctlass/mctlass.h" + +#include "mctlass/platform/platform.h" + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/*!@defgroup predicate_vector_concept Predicate Vector Concept +@{ + +Implementations of \ref predicate_vector_concept contain an ordered set of boolean predicates which +may be used as conditionals in other device-side operations. Both random access and iterators +offering sequential access are provided. + +@par Predicate Vector + A \ref predicate_vector_concept satisfies the following expressions + - at(int idx) - returns the value of the indexed predicate + - set(int idx, bool value) - sets the value of the indexed predicate + - begin() - returns a \ref predicate_iterator_concept pointing to the first predicate + +@} +*/ + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/*!@defgroup predicate_iterator_concept Predicate Iterator Concept +@{ + +Implementations of \ref predicate_iterator_concept enables accessing and traversing elements of a +bit vector. + +@par Const Predicate Iterator + A const \ref predicate_iterator_concept satisfies the following expressions + - ++it increments the iterator to the next predicate + - *it returns the value of the currently pointed-to predicate + +@par Mutable Predicate Iterator + A \ref predicate_iterator_concept that is non-const also satisfies the following expressions + - it.set(bool value) sets the value of the currently pointed-to predicate + +@} +*/ + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/*!@defgroup predicate_tile_adapter Predicate Tile Adapter Concept +@{ + +Implementations of \ref predicate_tile_adapter provide a mapping between a the elements of a \ref +tile_traits_concept and a \ref predicate_vector_concept. + +@par Predicate Tile Adapter + A \ref predicate_tile_adapter satisfies the following expressions + - at(int d, int h, int w, int c) - returns the value of a predicate corresponding to the + access (d, h, w, c) within the tile. + +@} +*/ + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Statically sized array of bits implementing @concept{predicate_vector_concept}. +template < + /// Number of predicates conatined in predicate vector + int kPredicates_, + /// Number of predicates contained in each byte of internal storage + int kPredicatesPerByte_ = 4, + /// Location of first predicate within byte of internal storage + int kPredicateStart_ = 0> +struct PredicateVector { + /// Number of bits stored by the PredicateVector + static int const kPredicates = kPredicates_; + + /// Number of bits stored within each byte of the predicate bit vector + static int const kPredicatesPerByte = kPredicatesPerByte_; + + /// First bit withing each byte containing predicates + static int const kPredicateStart = kPredicateStart_; + + // Make sure no one tries to put more than 8 bits in a byte :) + static_assert(kPredicatesPerByte <= 8, "kPredicatesPerByte must fit within an actual byte"); + // Make sure the "offsetted" bits fit in one byte. + static_assert(kPredicateStart + kPredicatesPerByte <= 8, + "The offsetted predicates must fit within an actual byte."); + + /// Storage type of individual elements + typedef uint32_t Storage; + + /// Number of bytes needed + static int const kBytes = (kPredicates + kPredicatesPerByte - 1) / kPredicatesPerByte; + + /// Number of storage elements needed + static int const kWordCount = (kBytes + int(sizeof(Storage)) - 1) / int(sizeof(Storage)); + + private: + // + // Data members + // + + /// Words of bit vector + Storage storageData[kWordCount]; + + // + // Methods + // + + /// Computes the word and bit corresponding to a logical predicate index + MCTLASS_HOST_DEVICE void computeStorageOffset(int &word, int &bit, int idx) const { + MCTLASS_ASSERT(idx < kPredicates); + + int byte = (idx / kPredicatesPerByte); + int bit_offset = (idx % kPredicatesPerByte); + + word = byte / sizeof(Storage); + int byte_offset = (byte % sizeof(Storage)); + + bit = byte_offset * 8 + bit_offset + kPredicateStart; + } + + /// Accesses a given word with optional assertions + MCTLASS_HOST_DEVICE Storage &storage(int word) { + MCTLASS_ASSERT(word < kWordCount); + return storageData[word]; + } + + /// Accesses a given word with optional assertions + MCTLASS_HOST_DEVICE Storage const &storage(int word) const { + MCTLASS_ASSERT(word < kWordCount); + return storageData[word]; + } + + public: + // + // Iterator + // + + /** + * @brief An iterator implementing \ref predicate_iterator_concept enabling sequential + * read and write access to predicates. + * @concept{predicate_iterator_concept} + */ + class Iterator { + /// Reference to PredicateVector instance + PredicateVector &vec_; + + /// Index into PredicateVector + int bit_; + + public: + /// Copy constructor + MCTLASS_HOST_DEVICE + Iterator(Iterator const &it) : vec_(it.vec_), bit_(it.bit_) {} + + /// Constructs an iterator from a PredicateVector + MCTLASS_HOST_DEVICE + Iterator(PredicateVector &vec, int _start = 0) : vec_(vec), bit_(_start) {} + + /// Pre-increment + MCTLASS_HOST_DEVICE + Iterator &operator++() { + ++bit_; + return *this; + } + + /// Increment + MCTLASS_HOST_DEVICE + Iterator &operator+=(int offset) { + bit_ += offset; + return *this; + } + + /// Pre-decrement + MCTLASS_HOST_DEVICE + Iterator &operator--() { + --bit_; + return *this; + } + + /// Decrement + MCTLASS_HOST_DEVICE + Iterator &operator-=(int offset) { + bit_ -= offset; + return *this; + } + + /// Post-increment + MCTLASS_HOST_DEVICE + Iterator operator++(int) { + Iterator ret(*this); + ret.bit_++; + return ret; + } + + /// Post-decrement + MCTLASS_HOST_DEVICE + Iterator operator--(int) { + Iterator ret(*this); + ret.bit_--; + return ret; + } + + /// Iterator advances by some amount + MCTLASS_HOST_DEVICE + Iterator operator+(int offset) { + Iterator ret(*this); + ret.bit_ += offset; + return ret; + } + + /// Iterator recedes by some amount + MCTLASS_HOST_DEVICE + Iterator operator-(int offset) { + ConstIterator ret(*this); + ret.bit_ -= offset; + return ret; + } + + /// Returns true if iterators point to the same bit + MCTLASS_HOST_DEVICE + bool operator==(Iterator const &it) const { return bit_ == it.bit_; } + + /// Returns false if iterators point to the same bit + MCTLASS_HOST_DEVICE + bool operator!=(Iterator const &it) const { return bit_ != it.bit_; } + + /// Gets the bit at the pointed to location + MCTLASS_HOST_DEVICE + bool get() { return vec_.at(bit_); } + + /// Gets the bit at the pointed to location + MCTLASS_HOST_DEVICE + bool at() const { return vec_.at(bit_); } + + /// Dereferences iterator + MCTLASS_HOST_DEVICE + bool operator*() const { return at(); } + + /// Sets the bit at the pointed to location + MCTLASS_HOST_DEVICE + void set(bool value = true) { vec_.set(bit_, value); } + }; + + /** + * @brief An iterator implementing \ref predicate_iterator_concept enabling sequential + * read and write access to predicates. + * @concept{predicate_iterator_concept} + */ + class ConstIterator { + /// Reference to PredicateVector instance + PredicateVector const &vec_; + + /// Index into PredicateVector + int bit_; + + public: + /// Copy constructor + MCTLASS_HOST_DEVICE + ConstIterator(ConstIterator const &it) : vec_(it.vec_), bit_(it.bit_) {} + + /// Constructs an iterator from a PredicateVector + MCTLASS_HOST_DEVICE + ConstIterator(PredicateVector const &vec, int _start = 0) : vec_(vec), bit_(_start) {} + + /// Pre-increment + MCTLASS_HOST_DEVICE + ConstIterator &operator++() { + ++bit_; + return *this; + } + + /// Increment + MCTLASS_HOST_DEVICE + ConstIterator &operator+=(int offset) { + bit_ += offset; + return *this; + } + + /// Pre-decrement + MCTLASS_HOST_DEVICE + ConstIterator &operator--() { + --bit_; + return *this; + } + + /// Decrement + MCTLASS_HOST_DEVICE + ConstIterator &operator-=(int offset) { + bit_ -= offset; + return *this; + } + + /// Post-increment + MCTLASS_HOST_DEVICE + ConstIterator operator++(int) { + ConstIterator ret(*this); + ret.bit_++; + return ret; + } + + /// Post-decrement + MCTLASS_HOST_DEVICE + ConstIterator operator--(int) { + ConstIterator ret(*this); + ret.bit_--; + return ret; + } + + /// Iterator advances by some amount + MCTLASS_HOST_DEVICE + ConstIterator operator+(int offset) { + ConstIterator ret(*this); + ret.bit_ += offset; + return ret; + } + + /// Iterator recedes by some amount + MCTLASS_HOST_DEVICE + ConstIterator operator-(int offset) { + ConstIterator ret(*this); + ret.bit_ -= offset; + return ret; + } + + /// Returns true if iterators point to the same bit + MCTLASS_HOST_DEVICE + bool operator==(ConstIterator const &it) const { return bit_ == it.bit_; } + + /// Returns false if iterators point to the same bit + MCTLASS_HOST_DEVICE + bool operator!=(ConstIterator const &it) const { return bit_ != it.bit_; } + + /// Gets the bit at the pointed to location + MCTLASS_HOST_DEVICE + bool get() { return vec_.at(bit_); } + + /// Gets the bit at the pointed to location + MCTLASS_HOST_DEVICE + bool at() const { return vec_.at(bit_); } + + /// Dereferences iterator + MCTLASS_HOST_DEVICE + bool operator*() const { return at(); } + }; + + /// Iterator that always returns true + struct TrivialIterator { + /// Constructor + MCTLASS_HOST_DEVICE + TrivialIterator() {} + + /// Copy constructor + MCTLASS_HOST_DEVICE + TrivialIterator(Iterator const &it) {} + + /// Constructs an iterator from a PredicateVector + MCTLASS_HOST_DEVICE + TrivialIterator(PredicateVector const &_vec) {} + + /// Pre-increment + MCTLASS_HOST_DEVICE + TrivialIterator &operator++() { return *this; } + + /// Post-increment + MCTLASS_HOST_DEVICE + TrivialIterator operator++(int) { return *this; } + + /// Dereferences iterator + MCTLASS_HOST_DEVICE + bool operator*() const { return true; } + }; + + public: + // + // Methods + // + + /// Initialize the predicate vector + MCTLASS_HOST_DEVICE PredicateVector(bool value = true) { fill(value); } + + /// Fills all predicates with a given value + MCTLASS_HOST_DEVICE void fill(bool value = true) { + Storage item = (value ? ~Storage(0) : Storage(0)); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kWordCount; ++i) { + storage(i) = item; + } + } + + /// Clears all predicates + MCTLASS_HOST_DEVICE void clear() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kWordCount; ++i) { + storage(i) = 0; + } + } + + /// Sets all predicates to true + MCTLASS_HOST_DEVICE void enable() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kWordCount; ++i) { + storage(i) = ~Storage(0); + } + } + + /// Accesses a bit within the predicate vector. + MCTLASS_HOST_DEVICE bool operator[](int idx) const { return at(idx); } + + /// Accesses a bit within the predicate vector. + MCTLASS_HOST_DEVICE bool at(int idx) const { + int bit, word; + computeStorageOffset(word, bit, idx); + + return ((storage(word) >> bit) & 1); + } + + /// Set a bit within the predicate vector. + MCTLASS_HOST_DEVICE void set(int idx, bool value = true) { + int bit, word; + computeStorageOffset(word, bit, idx); + + Storage disable_mask = (~(Storage(1) << bit)); + Storage enable_mask = (Storage(value) << bit); + + storage(word) = ((storage(word) & disable_mask) | enable_mask); + } + + /// Computes the intersection of two identical predicate vectors. + MCTLASS_HOST_DEVICE PredicateVector &operator&=(PredicateVector const &predicates) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kWordCount; ++i) { + storage(i) = (storage(i) & predicates.storage(i)); + } + return *this; + } + + /// Computes the union of two identical predicate vectors. + MCTLASS_HOST_DEVICE PredicateVector &operator|=(PredicateVector const &predicates) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kWordCount; ++i) { + storage(i) = (storage(i) | predicates.storage(i)); + } + return *this; + } + + /// Returns true if entire predicate array is zero. + MCTLASS_HOST_DEVICE bool is_zero() const { + Storage mask(0); + for (int byte = 0; byte < sizeof(Storage); ++byte) { + Storage byte_mask = (((1 << kPredicatesPerByte) - 1) << kPredicateStart); + mask |= (byte_mask << (byte * 8)); + } + uint32_t result = 0; + for (int word = 0; word < kWordCount; ++word) { + result |= storage(word); + } + return result == 0; + } + + /// Returns an iterator to the start of the bit vector + MCTLASS_DEVICE + Iterator begin() { return Iterator(*this); } + + /// Returns an iterator + MCTLASS_DEVICE + Iterator end() { return Iterator(*this, kPredicates); } + + /// Returns a ConstIterator + MCTLASS_DEVICE + ConstIterator const_begin() const { return ConstIterator(*this); } + + /// Returns a ConstIterator + MCTLASS_DEVICE + ConstIterator const_end() const { return ConstIterator(*this, kPredicates); } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/quaternion.h b/csrc/mctlass/include/mctlass/quaternion.h new file mode 100644 index 0000000..4e4f0fe --- /dev/null +++ b/csrc/mctlass/include/mctlass/quaternion.h @@ -0,0 +1,752 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines a densely packed quaternion object intended for storing data in registers and + executing quaternion operations within a CUDA or host thread. +*/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/functional.h" +#include "mctlass/array.h" +#include "mctlass/real.h" +#include "mctlass/coord.h" +#include "mctlass/matrix.h" +#include "mctlass/fast_math.h" +#include "mctlass/layout/vector.h" + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Quaternion: xi + yj + zk + w +template < + typename Element_ = float ///< element type +> +class Quaternion : public Array { +public: + + /// Logical rank of tensor index space + static int const kRank = 1; + + /// Number of elements + static int const kExtent = 4; + + /// Base class is a four-element array + using Base = Array; + + /// Element type + using Element = typename Base::Element; + + /// Reference type to an element + using Reference = typename Base::reference; + + /// Index type + using Index = int; + + /// Quaternion storage - imaginary part + static int const kX = 0; + + /// Quaternion storage - imaginary part + static int const kY = 1; + + /// Quaternion storage - imaginary part + static int const kZ = 2; + + /// Quaternion storage - real part + static int const kW = 3; + +public: + + // + // Methods + // + + /// Constructs a quaternion q = 0 + MCTLASS_HOST_DEVICE + Quaternion() { + Base::at(kX) = Element(); + Base::at(kY) = Element(); + Base::at(kZ) = Element(); + Base::at(kW) = Element(); + } + + /// Constructs a quaternion q = w + 0*i + 0*j + 0*k + MCTLASS_HOST_DEVICE + Quaternion( + Element w_ + ) { + Base::at(kX) = Element(); + Base::at(kY) = Element(); + Base::at(kZ) = Element(); + Base::at(kW) = w_; + } + + /// Constructs a quaternion q = w + x*i + y*j + z*k + MCTLASS_HOST_DEVICE + Quaternion( + Element x_, + Element y_, + Element z_, + Element w_ + ) { + Base::at(kX) = x_; + Base::at(kY) = y_; + Base::at(kZ) = z_; + Base::at(kW) = w_; + } + + /// Constructs a quaternion from a vector representing the imaginary part and a real number + MCTLASS_HOST_DEVICE + Quaternion( + Matrix3x1 const &imag_, + Element w_ = Element() + ) { + Base::at(kX) = imag_[0]; + Base::at(kY) = imag_[1]; + Base::at(kZ) = imag_[2]; + Base::at(kW) = w_; + } + + /// Returns a reference to the element at a given Coord + MCTLASS_HOST_DEVICE + Reference at(Index idx) const { + return Base::at(idx); + } + + /// Returns a reference to the element at a given Coord + MCTLASS_HOST_DEVICE + Reference at(Index idx) { + return Base::at(idx); + } + + /// Accesses the x element of the imaginary part of the quaternion + MCTLASS_HOST_DEVICE + Element x() const { + return Base::at(kX); + } + + /// Accesses the x element of the imaginary part of the quaternion + MCTLASS_HOST_DEVICE + Reference x() { + return Base::at(kX); + } + + /// Accesses the y element of the imaginary part of the quaternion + MCTLASS_HOST_DEVICE + Element y() const { + return Base::at(kY); + } + + /// Accesses the y element of the imaginary part of the quaternion + MCTLASS_HOST_DEVICE + Reference y() { + return Base::at(kY); + } + + /// Accesses the z element of the imaginary part of the quaternion + MCTLASS_HOST_DEVICE + Element z() const { + return Base::at(kZ); + } + + /// Accesses the z element of the imaginary part of the quaternion + MCTLASS_HOST_DEVICE + Reference z() { + return Base::at(kZ); + } + + /// Accesses the real part of the quaternion + MCTLASS_HOST_DEVICE + Element w() const { + return Base::at(kW); + } + + /// Accesses the real part of the quaternion + MCTLASS_HOST_DEVICE + Reference w() { + return Base::at(kW); + } + + /// Returns the pure imaginary part of the quaternion as a 3-vector + MCTLASS_HOST_DEVICE + Matrix3x1 pure() const { + return Matrix3x1(x(), y(), z()); + } + + /// Returns a quaternion representation of a spatial rotation given a unit-length axis and + /// a rotation in radians. + MCTLASS_HOST_DEVICE + static Quaternion rotation( + Matrix3x1 const &axis_unit, ///< axis of rotation (assumed to be unit length) + Element theta) { ///< angular rotation in radians + + Element s = fast_sin(theta / Element(2)); + + return Quaternion( + s * axis_unit[0], + s * axis_unit[1], + s * axis_unit[2], + fast_cos(theta / Element(2)) + ); + } + + /// Returns a quaternion representation of a spatial rotation represented as a + /// unit-length rotation axis (r_x, r_y, r_z) and an angular rotation in radians + MCTLASS_HOST_DEVICE + static Quaternion rotation( + Element r_x, + Element r_y, + Element r_z, + Element theta) { ///< angular rotation in radians + + return rotation({r_x, r_y, r_z}, theta); + } + + /// Geometric rotation of a 3-element vector + MCTLASS_HOST_DEVICE + Matrix3x1 rotate(Matrix3x1 const &rhs) const { + return (*this * Quaternion(rhs, 0) * reciprocal(*this)).pure(); + } + + /// Inverse rotation operation + MCTLASS_HOST_DEVICE + Matrix3x1 rotate_inv(Matrix3x1 const &rhs) const { + return (reciprocal(*this) * Quaternion(rhs, 0) * *this).pure(); + } + + /// Rotates a 3-vector assuming this is a unit quaternion (a spinor) + MCTLASS_HOST_DEVICE + Matrix3x1 spinor(Matrix3x1 const &rhs) const { + return (*this * Quaternion(rhs, 0) * conj(*this)).pure(); + } + + /// Inverse rotation of 3-vector assuming this is a unit quaternion (a spinor) + MCTLASS_HOST_DEVICE + Matrix3x1 spinor_inv(Matrix3x1 const &rhs) const { + return (conj(*this) * Quaternion(rhs, 0) * *this).pure(); + } + + /// In-place addition + template + MCTLASS_HOST_DEVICE + Quaternion &operator+=(Quaternion const &rhs) { + *this = (*this + rhs); + return *this; + } + + /// In-place subtraction + template + MCTLASS_HOST_DEVICE + Quaternion &operator-=(Quaternion const &rhs) { + *this = (*this - rhs); + return *this; + } + + /// In-place multiplication + template + MCTLASS_HOST_DEVICE + Quaternion &operator*=(Quaternion const &rhs) { + *this = (*this * rhs); + return *this; + } + + /// Scalar multiplication + template + MCTLASS_HOST_DEVICE + Quaternion &operator*=(Element s) { + *this = (*this * s); + return *this; + } + + /// In-place Division + template + MCTLASS_HOST_DEVICE + Quaternion &operator/=(Quaternion const &rhs) { + *this = (*this / rhs); + return *this; + } + + /// In-place Division + template + MCTLASS_HOST_DEVICE + Quaternion &operator/=(Element s) { + *this = (*this / s); + return *this; + } + + /// Computes a 3x3 rotation matrix (row-major representation) + MCTLASS_HOST_DEVICE + Matrix3x3 as_rotation_matrix_3x3() const { + Matrix3x3 m( + w() * w() + x() * x() - y() * y() - z() * z(), + 2 * x() * y() - 2 * w() * z(), + 2 * x() * z() + 2 * w() * y(), + + 2 * x() * y() + 2 * w() * z(), + w() * w() - x() * x() + y() * y() - z() * z(), + 2 * y() * z() - 2 * w() * x(), + + 2 * x() * z() - 2 * w() * y(), + 2 * y() * z() + 2 * w() * x(), + w() * w() - x() * x() - y() * y() + z() * z() + ); + return m; + } + + /// Computes a 4x4 rotation matrix (row-major representation) + MCTLASS_HOST_DEVICE + Matrix4x4 as_rotation_matrix_4x4() const { + Matrix4x4 m = Matrix4x4::identity(); + m.set_slice_3x3(as_rotation_matrix_3x3()); + return m; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Constructs a quaternion that is non-zero only in its real element. +template +MCTLASS_HOST_DEVICE +Quaternion make_Quaternion( + Element w) { ///< real part + + return Quaternion(w); +} + +/// Constructs a quaternion from a vector and real +template +MCTLASS_HOST_DEVICE +Quaternion make_Quaternion( + Matrix3x1 const &imag, ///< imaginary party as a vector + Element w) { ///< real part + + return Quaternion(imag, w); +} + +/// Constructs a quaternion from a unit-length rotation axis and a rotation +/// angle in radians +template +MCTLASS_HOST_DEVICE +Quaternion make_QuaternionRotation( + Matrix3x1 const &axis_unit, ///< rotation axis (unit-length) + Element w) { ///< rotation angle in radians + + return Quaternion::rotation(axis_unit, w); +} + +/// Constructs a quaternion q = xi + yj + zk + w +template +MCTLASS_HOST_DEVICE +Quaternion make_Quaternion(Element x, Element y, Element z, Element w) { + return Quaternion(x, y, z, w); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Returns the real part of the quaternion number +template +MCTLASS_HOST_DEVICE +Element const &real(Quaternion const &q) { + return q.w(); +} + +/// Returns the real part of the quaternion number +template +MCTLASS_HOST_DEVICE +Element &real(Quaternion &q) { + return q.w(); +} + +/// Returns the magnitude of the quaternion number +template +MCTLASS_HOST_DEVICE +Element abs(Quaternion const &q) { + return fast_sqrt(norm(q)); +} + +/// Quaternion conjugate +template +MCTLASS_HOST_DEVICE +Quaternion conj(Quaternion const &q) { + return make_Quaternion( + -q.x(), + -q.y(), + -q.z(), + q.w() + ); +} + +/// Computes the squared magnitude of the quaternion +template +MCTLASS_HOST_DEVICE +Element norm(Quaternion const &q) { + return q.x() * q.x() + q.y() * q.y() + q.z() * q.z() + q.w() * q.w(); +} + +/// Quaternion reciprocal +template +MCTLASS_HOST_DEVICE +Quaternion reciprocal(Quaternion const &q) { + + Element nsq = norm(q); + + return make_Quaternion( + -q.x() / nsq, + -q.y() / nsq, + -q.z() / nsq, + q.w() / nsq + ); +} + +/// Returns a unit-length quaternion +template +MCTLASS_HOST_DEVICE +Quaternion unit(Quaternion const &q) { + + Element rcp_mag = Element(1) / abs(q); + + return make_Quaternion( + q.x() * rcp_mag, + q.y() * rcp_mag, + q.z() * rcp_mag, + q.w() * rcp_mag + ); +} + +/// Quaternion exponential +template +MCTLASS_HOST_DEVICE +Quaternion exp(Quaternion const &q) { + + Element exp_ = fast_exp(q.w()); + Element imag_norm = fast_sqrt(q.x() * q.x() + q.y() * q.y() + q.z() * q.z()); + Element sin_norm = fast_sin(imag_norm); + + return make_Quaternion( + exp_ * q.x() * sin_norm / imag_norm, + exp_ * q.y() * sin_norm / imag_norm, + exp_ * q.z() * sin_norm / imag_norm, + exp_ * fast_cos(imag_norm) + ); +} + +/// Quaternion natural logarithm +template +MCTLASS_HOST_DEVICE +Quaternion log(Quaternion const &q) { + + Element v = fast_sqrt(q.x() * q.x() + q.y() * q.y() + q.z() * q.z()); + Element s = fast_acos(q.w() / abs(q)) / v; + + return make_Quaternion( + q.x() * s, + q.y() * s, + q.z() * s, + fast_log(q.w()) + ); +} + +/// Gets the rotation angle from a unit-length quaternion +template +MCTLASS_HOST_DEVICE +Element get_rotation_angle(Quaternion const &q_unit) { + return fast_acos(q_unit.w()) * Element(2); +} + +/// Gets the rotation axis from a unit-length quaternion +template +MCTLASS_HOST_DEVICE +Matrix3x1 get_rotation_axis(Quaternion const &q_unit) { + return q_unit.pure().unit(); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Equality operator +template +MCTLASS_HOST_DEVICE +bool operator==(Quaternion const &lhs, Quaternion const &rhs) { + return lhs.x() == rhs.x() && + lhs.y() == rhs.y() && + lhs.z() == rhs.z() && + lhs.w() == rhs.w(); +} + +/// Inequality operator +template +MCTLASS_HOST_DEVICE +bool operator!=(Quaternion const &lhs, Quaternion const &rhs) { + return !(lhs == rhs); +} + +/// Quaternion scalar multiplication +template +MCTLASS_HOST_DEVICE +Quaternion operator*(Quaternion q, Element s) { + return make_Quaternion( + q.x() * s, + q.y() * s, + q.z() * s, + q.w() * s + ); +} + +/// Quaternion scalar multiplication +template +MCTLASS_HOST_DEVICE +Quaternion operator*(Element s, Quaternion const &q) { + return make_Quaternion( + s * q.x(), + s * q.y(), + s * q.z(), + s * q.w() + ); +} + +/// Quaternion scalar division +template +MCTLASS_HOST_DEVICE +Quaternion operator/(Quaternion const &q, Element s) { + return make_Quaternion( + q.x() / s, + q.y() / s, + q.z() / s, + q.w() / s + ); +} + +/// Quaternion unary negation +template +MCTLASS_HOST_DEVICE +Quaternion operator-(Quaternion const &q) { + return make_Quaternion( + -q.x(), + -q.y(), + -q.z(), + -q.w() + ); +} + +/// Quaternion addition +template +MCTLASS_HOST_DEVICE +Quaternion operator+(Quaternion const &lhs, Quaternion const &rhs) { + return make_Quaternion( + lhs.x() + rhs.x(), + lhs.y() + rhs.y(), + lhs.z() + rhs.z(), + lhs.w() + rhs.w() + ); +} + +/// Quaternion subtraction +template +MCTLASS_HOST_DEVICE +Quaternion operator-(Quaternion const &lhs, Quaternion const &rhs) { + return make_Quaternion( + lhs.x() - rhs.x(), + lhs.y() - rhs.y(), + lhs.z() - rhs.z(), + lhs.w() - rhs.w() + ); +} + +/// Quaternion product +template +MCTLASS_HOST_DEVICE +Quaternion operator*(Quaternion const &lhs, Quaternion const &rhs) { + return make_Quaternion( + lhs.w() * rhs.x() + rhs.w() * lhs.x() + lhs.y() * rhs.z() - lhs.z() * rhs.y(), + lhs.w() * rhs.y() + rhs.w() * lhs.y() + lhs.z() * rhs.x() - lhs.x() * rhs.z(), + lhs.w() * rhs.z() + rhs.w() * lhs.z() + lhs.x() * rhs.y() - lhs.y() * rhs.x(), + lhs.w() * rhs.w() - lhs.x() * rhs.x() - lhs.y() * rhs.y() - lhs.z() * rhs.z() + ); +} + +/// Quaternion division +template +MCTLASS_HOST_DEVICE +Quaternion operator/(Quaternion const &lhs, Quaternion const &rhs) { + return lhs * reciprocal(rhs); +} + +/// Quaternion scalar division +template +MCTLASS_HOST_DEVICE +Quaternion operator/(Element s, Quaternion const &q) { + return s * reciprocal(q); +} + +/// Comparison +template +MCTLASS_HOST_DEVICE +bool operator<(Quaternion const &lhs, Quaternion const &rhs) { + return true; +} + +/// Rotates a 3-vector assuming this is a unit quaternion (a spinor). This avoids computing +/// a reciprocal. +template +MCTLASS_HOST_DEVICE +Matrix3x1 spinor_rotation( + Quaternion const &spinor, /// unit-length quaternion + Matrix3x1 const &rhs) { /// arbitrary 3-vector + + return (spinor * Quaternion(rhs, 0) * conj(spinor)).pure(); +} + +/// Inverse rotation of 3-vector assuming this is a unit quaternion (a spinor). This avoids computing +/// a reciprocal. +template +MCTLASS_HOST_DEVICE +Matrix3x1 spinor_rotation_inv( + Quaternion const &spinor, /// unit-length quaternion + Matrix3x1 const &rhs) { /// arbitrary 3-vector + + return (conj(spinor) * Quaternion(rhs, 0) * spinor).pure(); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization for Quaternion-valued type. +template +struct RealType< Quaternion > { + using Type = T; + + /// Number of elements + static int const kExtent = Quaternion::kExtent; + +MCTLASS_HOST_DEVICE + static Quaternion from_real(double x) { + return Quaternion(static_cast(x)); + } +}; + + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Factories +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +MCTLASS_HOST_DEVICE +mctlass::Quaternion from_real >(double r) { + return mctlass::Quaternion(half_t(r)); +} + +template <> +MCTLASS_HOST_DEVICE +mctlass::Quaternion from_real >(double r) { + return mctlass::Quaternion(float(r)); +} + +template <> +MCTLASS_HOST_DEVICE +mctlass::Quaternion from_real >(double r) { + return mctlass::Quaternion(r); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + + +///////////////////////////////////////////////////////////////////////////////////////////////// +// functional.h numeric specializations +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct multiplies> { + MCTLASS_HOST_DEVICE + Quaternion operator()(Quaternion lhs, Quaternion const &rhs) const { + lhs = lhs * rhs; + return lhs; + } +}; + +/// Squares with optional conversion +template +struct magnitude_squared, Output> { + MCTLASS_HOST_DEVICE + Output operator()(Quaternion lhs) const { + multiplies mul_op; + + Output y_w = Output(lhs.w()); + Output y_x = Output(lhs.x()); + Output y_y = Output(lhs.y()); + Output y_z = Output(lhs.z()); + + return mul_op(y_w, y_w) + mul_op(y_x, y_x) + mul_op(y_y, y_y) + \ + mul_op(y_z, y_z); + } +}; + +template +struct multiply_add, Quaternion, Quaternion> { + MCTLASS_HOST_DEVICE + Quaternion operator()( + Quaternion const &a, + Quaternion const &b, + Quaternion const &c) const { + + T x = c.x(); + T y = c.y(); + T z = c.z(); + T w = c.w(); + + x += a.w() * b.x(); + x += b.w() * a.x(); + x += a.y() * b.z(); + x += -a.z() * b.y(), + + y += a.w() * b.y(); + y += b.w() * a.y(); + y += a.z() * b.x(); + y += -a.x() * b.z(); + + z += a.w() * b.z(); + z += b.w() * a.z(); + z += a.x() * b.y(); + z += -a.y() * b.x(); + + w += a.w() * b.w(); + w += -a.x() * b.x(); + w += -a.y() * b.y(); + w += -a.z() * b.z(); + + return mctlass::make_Quaternion(x, y, z, w); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/real.h b/csrc/mctlass/include/mctlass/real.h new file mode 100644 index 0000000..41e392d --- /dev/null +++ b/csrc/mctlass/include/mctlass/real.h @@ -0,0 +1,61 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/** + \file + \brief This class provides helpers to support real<> and complex<> types in generic code. +*/ + +#pragma once + +namespace mctlass { + +/// Used to determine the real-valued underlying type of a numeric type T. +template +struct RealType { + using Type = T; + + /// Number of elements + static int const kExtent = 1; + +MCTLASS_HOST_DEVICE + static T from_real(double x) { + return static_cast(x); + } +}; + +template +MCTLASS_HOST_DEVICE +static T from_real(double r) { + return T(r); +} + + +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/reduction/device/reduce_split_k.h b/csrc/mctlass/include/mctlass/reduction/device/reduce_split_k.h new file mode 100644 index 0000000..e870cfc --- /dev/null +++ b/csrc/mctlass/include/mctlass/reduction/device/reduce_split_k.h @@ -0,0 +1,223 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Kernel performing a reduction over densely packed tensors in global memory +*/ + +#pragma once + +#include "mctlass/device_kernel.h" +#include "mctlass/reduction/kernel/reduce_split_k.h" +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace reduction { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename ReductionKernel_ +> +class ReduceSplitK { +public: + using ReductionKernel = ReductionKernel_; + + using Shape = typename ReductionKernel::Shape; + using ReductionOp = typename ReductionKernel::ReductionOp; + using OutputOp = typename ReductionKernel::OutputOp; + + using ElementWorkspace = typename ReductionKernel::ElementWorkspace; + using ElementAccumulator = typename ReductionKernel::ElementAccumulator; + using ElementOutput = typename ReductionKernel::ElementOutput; + + using WorkspaceTensorRef = typename ReductionKernel::WorkspaceTensorRef; + using OutputTensorRef = typename ReductionKernel::OutputTensorRef; + + using StrideIndex = typename ReductionKernel::StrideIndex; + + /// Argument structure + struct Arguments { + + // + // Data members + // + + MatrixCoord problem_size; + int partitions; + size_t partition_stride; + WorkspaceTensorRef workspace; + OutputTensorRef destination; + OutputTensorRef source; + typename OutputOp::Params output; + typename ReductionOp::Params reduction; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Arguments() : + problem_size(0, 0), + partitions(1), + partition_stride(0) { } + + MCTLASS_HOST_DEVICE + Arguments( + MatrixCoord const & problem_size + ): + problem_size(problem_size) { } + + MCTLASS_HOST_DEVICE + Arguments( + MatrixCoord problem_size_, + int partitions_, + size_t partition_stride_, + WorkspaceTensorRef workspace_, + OutputTensorRef destination_, + OutputTensorRef source_, + typename OutputOp::Params output_ = typename OutputOp::Params(), + typename ReductionOp::Params reduction_ = typename ReductionOp::Params() + ): + problem_size(problem_size_), + partitions(partitions_), + partition_stride(partition_stride_), + workspace(workspace_), + destination(destination_), + source(source_), + output(output_), + reduction(reduction_) + { + + } + + }; + +private: + /// Kernel parameters object + typename ReductionKernel::Params params_; + +public: + /// Constructs Reduction SplitK + ReduceSplitK() { } + + /// Determines whether the ReduceSplitK can execute the given problem. + static Status can_implement(Arguments const &args) { + + return Status::kSuccess; + } + + /// Gets the workspace size + static size_t get_workspace_size(Arguments const &args) { + // needs no additional workspace + return 0; + } + + /// Initializes Reduction state from arguments. + Status initialize( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + // initialize the params structure from the arguments + params_ = typename ReductionKernel::Params( + args.problem_size, + args.partitions, + args.partition_stride, + args.workspace, + args.destination, + args.source, + args.output, + args.reduction + ); + + return Status::kSuccess; + + } + + /// Initializes Reduction kernel state from arguments. + Status update(Arguments const &args, void *workspace = nullptr) { + + // update the params structure from the arguments + params_.workspace.reset(args.workspace.non_const_ref().data()); + params_.destination.reset(args.destination.non_const_ref().data()); + params_.source.reset(args.source.non_const_ref().data()); + params_.output = args.output; + params_.reduction = args.reduction; + + return Status::kSuccess; + } + + /// Runs the kernel using initialized state. + Status run(mcStream_t stream = nullptr) { + + // + // Launch reduction kernel + // + dim3 block = ReductionKernel::block_shape(); + dim3 grid = ReductionKernel::grid_shape(params_.problem_size); + + Kernel<<< grid, block, 0, stream >>>(params_); + + mcError_t result = mcGetLastError(); + + return result == mcSuccess ? Status::kSuccess : Status::kErrorInternal; + } + + + /// Runs the kernel using initialized state. + Status operator()(mcStream_t stream = nullptr) { + return run(stream); + } + + /// Runs the kernel using initialized state. + Status operator()( + Arguments const &args, + void *workspace = nullptr, + mcStream_t stream = nullptr) { + + Status status = initialize(args, workspace, stream); + + if (status == Status::kSuccess) { + status = run(stream); + } + + return status; + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace reduction +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/reduction/device/tensor_reduce.h b/csrc/mctlass/include/mctlass/reduction/device/tensor_reduce.h new file mode 100644 index 0000000..1ff9854 --- /dev/null +++ b/csrc/mctlass/include/mctlass/reduction/device/tensor_reduce.h @@ -0,0 +1,264 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Kernel performing a reduction over one or more ranks of an affine tensor +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/fast_math.h" +#include "mctlass/numeric_types.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/reduction/device/tensor_reduce_affine_strided.h" +#include "mctlass/reduction/device/tensor_reduce_affine_contiguous.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace reduction { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tensor reduction operator on specific MCTLASS layouts over exactly one index +template < + typename ElementOutput_, + typename ElementSource_, + typename Layout_, + typename ReductionOp_, + int VectorLength_ = 1, + typename ElementCompute_ = ElementOutput_ +> +struct TensorReduction { + + using ElementOutput = ElementOutput_; + using ElementSource = ElementSource_; + using Layout = Layout_; + using ReductionOp = ReductionOp_; + static int const kVectorLength = VectorLength_; + using ElementCompute = ElementCompute_; + + using TensorCoord = typename Layout::TensorCoord; + + /// Reduction operator + using ReductionDeviceStridedOperator = TensorReductionAffineStrided< + 4, 3, ElementOutput, ElementSource, ReductionOp, kVectorLength, ElementCompute + >; + + using ReductionDeviceContiguousOperator = TensorReductionAffineContiguous< + 4, 3, ElementOutput, ElementSource, ReductionOp, kVectorLength, ElementCompute + >; + + // + // Data members + // + + ReductionDeviceStridedOperator reduction_strided; + ReductionDeviceContiguousOperator reduction_contiguous; + int reduction_index; + + // + // Methods + // + + /// + TensorReduction( + TensorCoord extent, + int reduction_index_ + ): + reduction_index(reduction_index_) { + + Coord<4> extent_affine; + + switch (reduction_index) { + case 0: + extent_affine[0] = extent[1]; + extent_affine[1] = extent[2]; + extent_affine[2] = extent[0]; + extent_affine[3] = extent[3]; + break; + case 1: + extent_affine[0] = extent[0]; + extent_affine[1] = extent[2]; + extent_affine[2] = extent[1]; + extent_affine[3] = extent[3]; + break; + case 2: + extent_affine[0] = extent[0]; + extent_affine[1] = extent[1]; + extent_affine[2] = extent[2]; + extent_affine[3] = extent[3]; + break; + case 3: + extent_affine[0] = extent[0]; + extent_affine[1] = extent[1]; + extent_affine[2] = extent[2]; + extent_affine[3] = extent[3]; + break; + default: break; + } + + if (reduction_index == 3) { + reduction_contiguous = ReductionDeviceContiguousOperator(extent_affine); + } + else { + reduction_strided = ReductionDeviceStridedOperator(extent_affine); + } + } + + /// Simple check to verify the object is initialized correctly + bool good() const { + if (reduction_index == 3) { + return reduction_contiguous.good(); + } + return reduction_strided.good(); + } + + /// Size of one workspace + int64_t workspace_stride() const { + if (reduction_index == 3) { + return reduction_contiguous.workspace_stride(); + } + else { + return reduction_strided.workspace_stride(); + } + } + + /// Returns the size (in bytes) of a temporary workspace needed for reduction across CTAs + int64_t workspace_size() const { + if (reduction_index == 3) { + return reduction_contiguous.workspace_size(); + } + else { + return reduction_strided.workspace_size(); + } + } + + /// Helper to use overloaded function call operator + Status reduce( + TensorRef dst_ref, + TensorRef src_ref, + void *device_workspace_ptr = nullptr, + ElementCompute reduction_identity = ElementCompute(), + ReductionOp reduction_op = ReductionOp(), + mcStream_t stream = nullptr) { + + int64_t src_stride[3]; + int64_t dst_stride[3]; + + switch (reduction_index) { + case 0: + src_stride[0] = src_ref.stride()[1]; + src_stride[1] = src_ref.stride()[0]; + src_stride[2] = src_ref.stride()[2]; + dst_stride[0] = dst_ref.stride()[1]; + dst_stride[1] = dst_ref.stride()[0]; + break; + case 1: + src_stride[0] = src_ref.stride()[2]; + src_stride[1] = src_ref.stride()[0]; + src_stride[2] = src_ref.stride()[1]; + dst_stride[0] = dst_ref.stride()[2]; + dst_stride[1] = dst_ref.stride()[0]; + break; + case 2: + src_stride[0] = src_ref.stride()[2]; + src_stride[1] = src_ref.stride()[1]; + src_stride[2] = src_ref.stride()[0]; + dst_stride[0] = dst_ref.stride()[2]; + dst_stride[1] = dst_ref.stride()[1]; + break; + case 3: + src_stride[0] = src_ref.stride()[2]; + src_stride[1] = src_ref.stride()[1]; + src_stride[2] = src_ref.stride()[0]; + + dst_stride[0] = dst_ref.stride()[2]; + dst_stride[1] = dst_ref.stride()[1]; + dst_stride[2] = dst_ref.stride()[0]; + + default: break; + } + + if (reduction_index == 3) { + return reduction_contiguous( + dst_ref.data(), + dst_stride, + src_ref.data(), + src_stride, + device_workspace_ptr, + reduction_identity, + reduction_op, + stream); + } + else { + return reduction_strided( + dst_ref.data(), + dst_stride, + src_ref.data(), + src_stride, + device_workspace_ptr, + reduction_identity, + reduction_op, + stream); + } + } + + Status operator()( + TensorRef dst_ref, + TensorRef src_ref, + void *device_workspace_ptr = nullptr, + ElementCompute reduction_identity = ElementCompute(), + ReductionOp reduction_op = ReductionOp(), + mcStream_t stream = nullptr) { + + return reduce( + dst_ref, + src_ref, + device_workspace_ptr, + reduction_identity, + reduction_op, + stream); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace reduction +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/reduction/device/tensor_reduce_affine_contiguous.h b/csrc/mctlass/include/mctlass/reduction/device/tensor_reduce_affine_contiguous.h new file mode 100644 index 0000000..8d2668c --- /dev/null +++ b/csrc/mctlass/include/mctlass/reduction/device/tensor_reduce_affine_contiguous.h @@ -0,0 +1,373 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Kernel performing a reduction over one or more ranks of an affine tensor +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/fast_math.h" +#include "mctlass/numeric_types.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/reduction/kernel/tensor_reduce_affine_contiguous.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace reduction { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tensor reduction operator on layouts which are affine +template < + int Rank, ///< Rank of source tensor (e.g. NDHWC => 5) + int ReducedRank, ///< Rank of reduced tensor (e.g. ND => 2) + typename ElementOutput_, + typename ElementSource_, + typename ReductionOp_, + int VectorLength = 1, + typename ElementCompute_ = ElementOutput_, + int Threads = 256, ///< Number of participating threads + int BatchSize = 4 ///< Number of elements to load per batch +> +struct TensorReductionAffineContiguous { + + static int const kRank = Rank; + static int const kReducedRank = ReducedRank; + static int const kVectorLength = VectorLength; + static int const kInnerRank = kRank - kReducedRank; + static int const kThreads = Threads; + static int const kBatchSize = BatchSize; + + using ElementOutput = ElementOutput_; + using ElementSource = ElementSource_; + using ReductionOp = ReductionOp_; + using ElementCompute = ElementCompute_; + + // + // Data members + // + + /// Internal status field + Status status; + + /// Extent of tensor in source layout + Coord extent; + + /// Number of points in the outer index space + int64_t outer_count; + + /// Number of elements in the inner index space + int64_t inner_count; + + /// Number of workspaces needed + int workspace_count; + + /// CUDA Grid shape (.x => contiguous, .y => outer, .z => inner) + dim3 grid_shape; + + /// CUDA Threadblock shape (.x => contiguous, .y => outer, .z => inner) + dim3 threadblock_shape; + + /// CUDA grid shape for the final reduction step if needed + dim3 grid_final; + + /// CUDA threadblock shape for the final reduction step if needed + dim3 threadblock_final; + +private: + // + // Methods + // + + /// Helper to reshape 'count' such that it is less than 2 x 'ext' + static int reshape_pow2(int ext, int count) { + if (ext > count) { + return 1; + } + int x = 1; + for (; count >= ext * 2; ) { + count >>= 1; + x <<= 1; + } + return x; + } + +public: + + /// Default ctor + TensorReductionAffineContiguous(): + status(Status::kErrorInvalidProblem), + extent(), + outer_count(0), + inner_count(0), + workspace_count(0), + grid_shape(0, 0, 0), + threadblock_shape(0, 0, 0) { } + + /// Constructor + TensorReductionAffineContiguous( + Coord extent_, + int target_threadblock_count = 128 + ): + status(Status::kSuccess), + extent(extent_), + outer_count(0), + inner_count(0), + workspace_count(0) { + + // + // Plan the parallel mapping strategy. + // + + outer_count = 1; + inner_count = 1; + + // Compute number of elements in strided ranks + for (int p = 0; p < kReducedRank; ++p) { + outer_count *= extent[p]; + } + + for (int p = 0; p < kInnerRank; ++p) { + inner_count *= extent[kReducedRank + p]; + } + + int cta_count_x = 1; + int cta_count_y = 1; + int cta_count_z = 1; + + int cta_threads_x = kThreads; + int cta_threads_y = 1; + int cta_threads_z = 1; + + // Determine CTA shape + int64_t inner_vector_count = inner_count / kVectorLength; + + // Priority 1. Assign threadblocks to outer indices if possible + if (outer_count > target_threadblock_count) { + cta_count_x = 1; + cta_count_y = target_threadblock_count; + cta_count_z = 1; + } + else { + + cta_count_y = int(outer_count); + int remaining_ctas = target_threadblock_count / cta_count_y; + + // Priority 2. Assign inner dimensions to one CTA + if (inner_vector_count > cta_threads_x) { + int64_t cta_z_bound = inner_vector_count / cta_threads_x; + if (cta_z_bound > remaining_ctas) { + cta_count_z = remaining_ctas; + } + else { + cta_count_z = int(cta_z_bound); + } + } + else { + cta_threads_x = reshape_pow2(int(inner_vector_count), cta_threads_x); + cta_count_z = 1; + } + } + + grid_shape = dim3(cta_count_x, cta_count_y, cta_count_z); + threadblock_shape = dim3(cta_threads_x, cta_threads_y, cta_threads_z); + + workspace_count = (cta_count_z > 1 ? cta_count_z : 0); + + // Determine shape of final reduction kernel if needed + if (workspace_count) { + + int final_threads = kThreads; + int final_ctas = 1; + + if (outer_count > kThreads) { + final_ctas = int(outer_count + kThreads - 1) / kThreads; + } + else { + final_threads = int(outer_count); + } + + grid_final = dim3(final_ctas, 1, 1); + threadblock_final = dim3(final_threads, 1, 1); + } + else { + grid_final = dim3(0, 0, 0); + threadblock_final = dim3(0, 0, 0); + } + } + + /// Simple check to verify the object is initialized correctly + bool good() const { + return status == Status::kSuccess; + } + + /// Size (in bytes) of workspace elements which are densely packed together + int64_t workspace_stride() const { + + // Error condition + if (!good()) { + return 0; + } + + return outer_count * sizeof_bits::value / 8; + } + + /// Returns the size (in bytes) of a temporary workspace needed for reduction across CTAs + int64_t workspace_size() const { + + // Error condition + if (!good()) { + return 0; + } + + // No reduction across CTAs + if (grid_shape.z == 1) { + return 0; + } + + return workspace_stride() * grid_shape.z; + } + + /// Performs a reduction + Status reduce( + ElementOutput *dst_ptr, ///< Pointer to destination tensor + int64_t dst_stride[], ///< Stride vector (of length kReducedRank - 1) + ElementSource const *src_ptr, ///< Pointer to source tensor + int64_t src_stride[], ///< Stride vector (of length kRank - 1) + void *device_workspace_ptr = nullptr, ///< Device workspace + ElementCompute reduction_identity = ElementCompute(), ///< Reduction identity element + ReductionOp reduction_op = ReductionOp(), ///< Reduction operator + mcStream_t stream = nullptr) { ///< MACA Stream into which all kernels are launched + + // Initial status check + if (!good()) { + return status; + } + + // Guard against null workspace + if (workspace_count > 1 && device_workspace_ptr == nullptr) { + return Status::kErrorWorkspaceNull; + } + + // Define reduction kernel + using ReductionKernel = kernel::TensorReductionAffineContiguous< + kRank, + kReducedRank, + ElementOutput, + ElementSource, + ReductionOp, + kVectorLength, + ElementCompute, + kThreads>; + + using FinalReductionKernel = kernel::TensorReductionAffineContiguousFinal< + kRank, + kReducedRank, + ElementOutput, + ElementSource, + ReductionOp, + kVectorLength, + ElementCompute, + kThreads>; + + using Params = typename ReductionKernel::Params; + + // Construct the parameters + Params params( + extent, + dst_ptr, + dst_stride, + src_ptr, + src_stride, + static_cast(device_workspace_ptr), + workspace_stride(), + workspace_count, + reduction_op, + reduction_identity); + + // Shared memory size + int shared_mem_bytes = sizeof(typename ReductionKernel::SharedStorage); + + // Launch the kernel + Kernel<<< grid_shape, threadblock_shape, shared_mem_bytes, stream >>>(params); + + // Check error condition + if (mcPeekAtLastError() == mcSuccess) { + status = Status::kSuccess; + } + else { + status = Status::kErrorInternal; + } + + // Final reduction kernel + if (workspace_count) { + Kernel<<< grid_final, threadblock_final, 0, stream >>>(params); + } + + // Check error condition + if (mcPeekAtLastError() == mcSuccess) { + status = Status::kSuccess; + } + else { + status = Status::kErrorInternal; + } + + return status; + } + + /// Helper to use overloaded function call operator + Status operator()( + ElementOutput *dst_ptr, ///< Pointer to destination tensor + int64_t dst_stride[], ///< Stride vector (of length kReducedRank - 1) + ElementSource const *src_ptr, ///< Pointer to source tensor + int64_t src_stride[], ///< Stride vector (of length kRank - 1) + void *device_workspace_ptr = nullptr, ///< Pointer to device workspace + ElementCompute reduction_identity = ElementCompute(), ///< Reduction identity element + ReductionOp reduction_op = ReductionOp(), ///< Reduction operator + mcStream_t stream = nullptr) { ///< MACA Stream into which all kernels are launched + + return reduce(dst_ptr, dst_stride, src_ptr, src_stride, device_workspace_ptr, reduction_identity, reduction_op, stream); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace reduction +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/reduction/device/tensor_reduce_affine_strided.h b/csrc/mctlass/include/mctlass/reduction/device/tensor_reduce_affine_strided.h new file mode 100644 index 0000000..71177b2 --- /dev/null +++ b/csrc/mctlass/include/mctlass/reduction/device/tensor_reduce_affine_strided.h @@ -0,0 +1,361 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Kernel performing a reduction over one or more ranks of an affine tensor +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/fast_math.h" +#include "mctlass/numeric_types.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/reduction/kernel/tensor_reduce_affine_strided.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace reduction { +namespace device { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tensor reduction operator on layouts which are affine +template < + int Rank, ///< Rank of source tensor (e.g. NDHWC => 5) + int ReducedRank, ///< Rank of reduced tensor (includes contiguous, e.g. NC => 2) + typename ElementOutput_, + typename ElementSource_, + typename ReductionOp_, + int VectorLength = 1, + typename ElementCompute_ = ElementOutput_, + int Threads = 256, ///< Number of participating threads + int BatchSize = 4 ///< Number of elements to load per batch +> +struct TensorReductionAffineStrided { + + static int const kRank = Rank; + static int const kReducedRank = ReducedRank; + static int const kVectorLength = VectorLength; + static int const kInnerRank = kRank - kReducedRank; + static int const kThreads = Threads; + static int const kBatchSize = BatchSize; + + using ElementOutput = ElementOutput_; + using ElementSource = ElementSource_; + using ReductionOp = ReductionOp_; + using ElementCompute = ElementCompute_; + + // + // Data members + // + + /// Internal status field + Status status; + + /// Extent of tensor in source layout + Coord extent; + + /// Number of points in the outer index space + int64_t outer_count; + + /// Number of elements in the inner index space + int64_t inner_count; + + /// Number of workspaces needed + int workspace_count; + + /// CUDA Grid shape (.x => contiguous, .y => outer, .z => inner) + dim3 grid_shape; + + /// CUDA Threadblock shape (.x => contiguous, .y => outer, .z => inner) + dim3 threadblock_shape; + + /// CUDA grid shape for the final reduction step if needed + dim3 grid_final; + + /// CUDA threadblock shape for the final reduction step if needed + dim3 threadblock_final; + +private: + // + // Methods + // + + /// Helper to reshape 'count' such that it is less than 2 x 'ext' + static int reshape_pow2(int ext, int count) { + if (ext > count) { + return 1; + } + int x = 1; + for (; count >= ext * 2; ) { + count >>= 1; + x <<= 1; + } + return x; + } + +public: + + /// Default ctor + TensorReductionAffineStrided(): + status(Status::kErrorInvalidProblem), + extent(), + outer_count(0), + inner_count(0), + workspace_count(0), + grid_shape(0, 0, 0), + threadblock_shape(0, 0, 0) { } + + /// Constructor + TensorReductionAffineStrided( + Coord extent_, + int target_threadblock_count = 128 + ): + status(Status::kSuccess), + extent(extent_), + outer_count(0), + inner_count(0), + workspace_count(0) { + + // + // Plan the parallel mapping strategy. + // + + outer_count = 1; + inner_count = 1; + + // Compute number of elements in strided ranks + for (int p = 0; p < kReducedRank - 1; ++p) { + outer_count *= extent[p]; + } + + for (int p = 0; p < kInnerRank; ++p) { + inner_count *= extent[kReducedRank + p - 1]; + } + + // Compute plan for the reduction + int extent_c = extent[kRank - 1]; + int vectors_c = (extent_c -1 + kVectorLength) / kVectorLength; + + // Determine CTA shape + int cta_width = kThreads * kVectorLength; + int cta_ways = reshape_pow2(extent_c, cta_width); + int cta_threads_x = kThreads / cta_ways; + + threadblock_shape = dim3(cta_threads_x, 1, std::min(cta_ways, 64)); + + // This leads to an error. + if (threadblock_shape.z > 1) { + if (threadblock_shape.y != 1) { + status = Status::kErrorInternal; + return; + } + } + + // Determine grid shape + int cta_count_x = (vectors_c + cta_threads_x - 1) / cta_threads_x; + int cta_count_y = std::max(1, target_threadblock_count / cta_count_x); + + // Limit the number of CTAs assigned to outer dimension + if (int64_t(cta_count_y * threadblock_shape.y) > outer_count) { + cta_count_y = int(outer_count + threadblock_shape.y - 1) / threadblock_shape.y; + } + + // Limit the number of CTAs assigned to inner dimension + int cta_count_z = std::max(1, target_threadblock_count / cta_count_y); + if (int64_t(cta_count_z * threadblock_shape.z) > inner_count) { + cta_count_z = int(inner_count + threadblock_shape.z - 1) / threadblock_shape.z; + } + + grid_shape = dim3(cta_count_x, cta_count_y, cta_count_z); + workspace_count = (cta_count_z > 1 ? cta_count_z : 0); + + // Determine shape of final reduction kernel if needed + grid_final = dim3(cta_count_x, int(outer_count)); + threadblock_final = dim3(cta_threads_x, 1, 1); + } + + /// Simple check to verify the object is initialized correctly + bool good() const { + return status == Status::kSuccess; + } + + /// Size of one CTA's workspace + int64_t workspace_stride() const { + + // Error condition + if (!good()) { + return 0; + } + + int vector_size_bytes = kVectorLength * sizeof_bits::value / 8; + + return extent[kRank - 1] * vector_size_bytes; + } + + /// Returns the size (in bytes) of a temporary workspace needed for reduction across CTAs + int64_t workspace_size() const { + + // Error condition + if (!good()) { + return 0; + } + + // No reduction across CTAs + if (grid_shape.z == 1) { + return 0; + } + + return workspace_stride() * outer_count * grid_shape.z; + } + + /// Performs a reduction + Status reduce( + ElementOutput *dst_ptr, ///< Pointer to destination tensor + int64_t dst_stride[], ///< Stride vector (of length kReducedRank - 1) + ElementSource const *src_ptr, ///< Pointer to source tensor + int64_t src_stride[], ///< Stride vector (of length kRank - 1) + void *device_workspace_ptr = nullptr, ///< Device workspace + ElementCompute reduction_identity = ElementCompute(), ///< Reduciton identity + ReductionOp reduction_op = ReductionOp(), ///< Reduction operator + mcStream_t stream = nullptr) { ///< MACA Stream into which all kernels are launched + + // Initial status check + if (!good()) { + return status; + } + + // Guard against null workspace + if (workspace_count > 1 && device_workspace_ptr == nullptr) { + return Status::kErrorWorkspaceNull; + } + + // Define reduction kernel + using ReductionKernel = kernel::TensorReductionAffineStrided< + kRank, + kReducedRank, + ElementOutput, + ElementSource, + ReductionOp, + kVectorLength, + ElementCompute, + kThreads>; + + using FinalReductionKernel = kernel::TensorReductionAffineStridedFinal< + kRank, + kReducedRank, + ElementOutput, + ElementSource, + ReductionOp, + kVectorLength, + ElementCompute, + kThreads>; + + using Params = typename ReductionKernel::Params; + + // Construct the parameters + Params params( + extent, + dst_ptr, + dst_stride, + src_ptr, + src_stride, + static_cast(device_workspace_ptr), + workspace_stride(), + workspace_count, + reduction_op, + reduction_identity); + + // Shared memory size + int shared_mem_bytes = sizeof(typename ReductionKernel::SharedStorage); + + // Launch the kernel + Kernel<<< grid_shape, threadblock_shape, shared_mem_bytes, stream >>>(params); + + // Check error condition + if (mcPeekAtLastError() == mcSuccess) { + status = Status::kSuccess; + } + else { + status = Status::kErrorInternal; + } + + // Final reduction kernel + if (workspace_count) { + + Kernel<<< grid_final, threadblock_final, 0, stream >>>(params); + + // Check error condition + if (mcPeekAtLastError() == mcSuccess) { + status = Status::kSuccess; + } + else { + status = Status::kErrorInternal; + } + } + + return status; + } + + /// Helper to use overloaded function call operator + Status operator()( + ElementOutput *dst_ptr, ///< Pointer to destination tensor + int64_t dst_stride[], ///< Stride vector (of length kReducedRank - 1) + ElementSource const *src_ptr, ///< Pointer to source tensor + int64_t src_stride[], ///< Stride vector (of length kRank - 1) + void *device_workspace_ptr = nullptr, ///< Pointer to device workspace + ElementCompute reduction_identity = ElementCompute(), ///< Reduciton identity + ReductionOp reduction_op = ReductionOp(), ///< Reduction operator + mcStream_t stream = nullptr) { ///< MACA Stream into which all kernels are launched + + return reduce( + dst_ptr, + dst_stride, + src_ptr, + src_stride, + device_workspace_ptr, + reduction_identity, + reduction_op, + stream); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace device +} // namespace reduction +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/reduction/kernel/reduce_softmax_final.h b/csrc/mctlass/include/mctlass/reduction/kernel/reduce_softmax_final.h new file mode 100644 index 0000000..1c7fbe6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/reduction/kernel/reduce_softmax_final.h @@ -0,0 +1,267 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Kernel performing a final reduction for softmax +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/arch/memory.h" +#include "mctlass/arch/memory_sm75.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace reduction { +namespace kernel { + +template < + typename ElementNorm_, + typename ElementSum_, + typename ElementSoftmaxCompute_, + typename ThreadblockShape_, + bool GroupedProblem = false +> +class ApplySoftmaxFinalReduction { +public: + + using ElementNorm = ElementNorm_; + using ElementSum = ElementSum_; + using ElementSoftmaxCompute = ElementSoftmaxCompute_; + using ThreadblockShape = ThreadblockShape_; + static const bool isGroupedProblem = GroupedProblem; + + // + // Arguments + // + + struct Arguments { + + mctlass::gemm::GemmCoord* problem_sizes; + mctlass::gemm::GemmCoord problem_size; + ElementNorm* block_Norm; + ElementSum* block_Sum; + int64_t* offset_Norm_Device; + int64_t* offset_Sum_Device; + int64_t batch_stride_Max; + int64_t batch_stride_Sum; + + // + // Methods + // + Arguments() { } + + // Non-grouped constructor without batching + Arguments( + mctlass::gemm::GemmCoord problem_size, + ElementNorm* block_Norm, + ElementSum* block_Sum + ): + problem_size(problem_size), + block_Norm(block_Norm), + block_Sum(block_Sum), + problem_sizes(nullptr), + offset_Norm_Device(nullptr), + offset_Sum_Device(nullptr), + batch_stride_Max(0), + batch_stride_Sum(0) + { + + } + + // Non-grouped constructor with batching + Arguments( + mctlass::gemm::GemmCoord problem_size, + ElementNorm* block_Norm, + ElementSum* block_Sum, + int64_t batch_stride_Max, + int64_t batch_stride_Sum + ): + problem_size(problem_size), + block_Norm(block_Norm), + block_Sum(block_Sum), + batch_stride_Max(batch_stride_Max), + batch_stride_Sum(batch_stride_Sum), + problem_sizes(nullptr), + offset_Norm_Device(nullptr), + offset_Sum_Device(nullptr) + { + + } + + + // Grouped constructor + Arguments( + mctlass::gemm::GemmCoord *problem_sizes, + ElementNorm* block_Norm, + ElementSum* block_Sum, + int64_t* offset_Norm_Device, + int64_t* offset_Sum_Device + ): + problem_sizes(problem_sizes), + problem_size(mctlass::gemm::GemmCoord(0, 0, 0)), + block_Norm(block_Norm), + block_Sum(block_Sum), + offset_Norm_Device(offset_Norm_Device), + offset_Sum_Device(offset_Sum_Device) + { + + } + }; + + struct SharedStorage { + + + }; + + // + // Params struct + // + + struct Params { + Arguments args; + + // + // Methods + // + Params() { } + + Params(Arguments const &args_): args(args_) { } + }; + +private: + +public: + + MCTLASS_DEVICE + ApplySoftmaxFinalReduction() { } + + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + apply(params, shared_storage); + } + +private: + + /// Full reduction + MCTLASS_DEVICE + void apply(Params const ¶ms, SharedStorage &shared_storage) { + + int tid = threadIdx.x; + int bid = blockIdx.x; + int bdim = blockDim.x; + + int block_batch = blockIdx.z; + + // defining three vars for a general reduction module + mctlass::gemm::GemmCoord problem_size = isGroupedProblem ? params.args.problem_sizes[bid] : params.args.problem_size; + int m_dim_in_loop = isGroupedProblem ? problem_size.m() : tid + bdim; + int access_offset = isGroupedProblem ? 0 : bid * bdim; + + if (!isGroupedProblem && access_offset + tid >= problem_size.m()) return; + + ElementNorm *curr_ptr_Max = isGroupedProblem ? \ + params.args.block_Norm + params.args.offset_Norm_Device[bid] : \ + params.args.block_Norm + block_batch * params.args.batch_stride_Max; + ElementSum *curr_ptr_Sum = isGroupedProblem ? \ + params.args.block_Sum + params.args.offset_Sum_Device[bid] : \ + params.args.block_Sum + block_batch * params.args.batch_stride_Sum; + + int threadblock_num = (problem_size.n() + ThreadblockShape::kN - 1) / ThreadblockShape::kN; + + using ConvertSumOutput = mctlass::NumericConverter; + using ConvertNormOutput = mctlass::NumericConverter; + + using ConvertSum = mctlass::NumericConverter; + using ConvertNorm = mctlass::NumericConverter; + + ConvertSum convert_sum; + ConvertNorm convert_norm; + + ConvertSumOutput convert_sum_output; + ConvertNormOutput convert_norm_output; + + uint32_t float_max_bits = 0xff7fffff; + float min_float = reinterpret_cast(float_max_bits); + + MCTLASS_PRAGMA_UNROLL + for (int idx_m = tid; idx_m < m_dim_in_loop; idx_m += bdim) { + ElementNorm *access_n = curr_ptr_Max + idx_m + access_offset; + ElementSum *access_s = curr_ptr_Sum + idx_m + access_offset; + ElementNorm *access_n_bak = access_n; + ElementSum *access_s_bak = access_s; + ElementSoftmaxCompute max_val = ElementSoftmaxCompute(min_float); + ElementSoftmaxCompute sum_val = ElementSoftmaxCompute(0); + ElementNorm fetch_n; + ElementSum fetch_s; + + MCTLASS_PRAGMA_UNROLL + for (int idx_n = 0; idx_n < threadblock_num; idx_n++) { + mctlass::arch::global_load(fetch_n, access_n, true); + max_val = mctlass::fast_max(max_val, convert_norm(fetch_n)); + access_n += problem_size.m(); + } + + access_n = access_n_bak; + + MCTLASS_PRAGMA_UNROLL + for (int idx_n = 0; idx_n < threadblock_num; idx_n++) { + mctlass::arch::global_load(fetch_n, access_n, true); + mctlass::arch::global_load(fetch_s, access_s, true); + sum_val += convert_sum(fetch_s) * mctlass::fast_exp(convert_norm(fetch_n) - max_val); + access_n += problem_size.m(); + access_s += problem_size.m(); + } + + ElementSoftmaxCompute inv_sum = mctlass::constants::one() / sum_val; + + access_n = access_n_bak; + access_s = access_s_bak; + + access_n[0] = convert_norm_output(max_val); + access_s[0] = convert_sum_output(inv_sum); + } + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace reduction +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/reduction/kernel/reduce_split_k.h b/csrc/mctlass/include/mctlass/reduction/kernel/reduce_split_k.h new file mode 100644 index 0000000..886506d --- /dev/null +++ b/csrc/mctlass/include/mctlass/reduction/kernel/reduce_split_k.h @@ -0,0 +1,248 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Kernel performing a reduction over densely packed tensors in global memory +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/numeric_conversion.h" + +#include "mctlass/layout/matrix.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace reduction { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape_, ///< shape of CTA (concept: MatrixShape) + typename OutputOp_ , ///< output operator (concept: epilogue::thread operator) + typename ReductionOp_, ///< reduction operator (concept: ReductionOperator) + int PartitionsPerStage = 4 ///< number of partitions to issue +> +class ReduceSplitK { +public: + + using Shape = Shape_; + using ReductionOp = ReductionOp_; + using OutputOp = OutputOp_; + static int const kElementsPerAccess = OutputOp::kCount; + static int const kPartitionsPerStage = PartitionsPerStage; + + using ElementWorkspace = typename ReductionOp::Element; + using ElementAccumulator = typename ReductionOp::ElementAccumulator; + using ElementOutput = typename OutputOp::ElementOutput; + + using WorkspaceTensorRef = TensorRef; + using OutputTensorRef = TensorRef; + using StrideIndex = typename WorkspaceTensorRef::Layout::Stride::Index; + + using FragmentWorkspace = AlignedArray; + using FragmentAccumulator = Array; + using FragmentOutput = AlignedArray; + + // + // Types + // + + /// Params structure + struct Params { + + MatrixCoord problem_size; + int partitions; + size_t partition_stride; + WorkspaceTensorRef workspace; + OutputTensorRef destination; + OutputTensorRef source; + typename OutputOp::Params output; + typename ReductionOp::Params reduction; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params( + MatrixCoord problem_size_, + int partitions_, + size_t partition_stride_, + WorkspaceTensorRef workspace_, + OutputTensorRef destination_, + OutputTensorRef source_, + typename OutputOp::Params output_ = typename OutputOp::Params(), + typename ReductionOp::Params reduction_ = typename ReductionOp::Params() + ): + problem_size(problem_size_), + partitions(partitions_), + partition_stride(sizeof(FragmentWorkspace) * partition_stride_ / kElementsPerAccess), + workspace(workspace_), + destination(destination_), + source(source_), + output(output_), + reduction(reduction_) { + + } + }; + + struct SharedStorage { }; + + +public: + + /// Computes the grid size given a chosen threadblock shape + MCTLASS_HOST_DEVICE + static dim3 grid_shape( + mctlass::MatrixCoord problem_size) { + + return dim3( + (problem_size.row() + Shape::kRow - 1) / Shape::kRow, + (problem_size.column() + Shape::kColumn - 1) / Shape::kColumn); + } + + /// Determines the threadblock shape + MCTLASS_HOST_DEVICE + static dim3 block_shape() { + return dim3(Shape::kColumn / kElementsPerAccess, Shape::kRow); + } + + /// Perform a reduction + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &storage) { + + // Determine CTA position + MatrixCoord thread_offset( + MatrixCoord::Index(int(blockIdx.x) * Shape::kRow + threadIdx.y), + MatrixCoord::Index(int(blockIdx.y) * Shape::kColumn + threadIdx.x * kElementsPerAccess) + ); + + // One guard conditional + if (!(thread_offset.row() < params.problem_size.row() && + thread_offset.column() < params.problem_size.column())) { + + return; + } + + + ReductionOp reduction_op(params.reduction); + + FragmentAccumulator accumulator; + + accumulator.clear(); + + // + // Load the first slice + // + + char const *workspace_ptr = + reinterpret_cast( + params.workspace.data() + params.workspace.offset(thread_offset)); + + FragmentWorkspace workspace_frag[kPartitionsPerStage]; + + // + // Construct the output operator + // + + OutputOp output_op(params.output); + + // + // Load and accumulate with a simple batched loading sequence. + // + + MCTLASS_PRAGMA_NO_UNROLL + for (int k = 0; k < params.partitions; k += kPartitionsPerStage) { + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPartitionsPerStage; ++i) { + if (k + i < params.partitions) { + workspace_frag[i] = *reinterpret_cast(workspace_ptr); + workspace_ptr += params.partition_stride; + } + } + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPartitionsPerStage; ++i) { + if (k + i < params.partitions) { + accumulator = reduction_op(accumulator, workspace_frag[i]); + } + } + } + + // + // Conditionally load the source + // + + FragmentOutput source_frag; + + source_frag.clear(); + + FragmentOutput const *source_ptr = reinterpret_cast( + params.source.data() + params.source.offset(thread_offset)); + + if (output_op.is_source_needed()) { + reinterpret_cast(source_frag) = *source_ptr; + } + + // + // Compute the output + // + + typename OutputOp::FragmentOutput output_frag = output_op(accumulator, source_frag); + + // + // Store + // + + FragmentOutput *dest_ptr = reinterpret_cast( + params.destination.data() + params.destination.offset(thread_offset)); + + *dest_ptr = reinterpret_cast(output_frag); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace reduction +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/reduction/kernel/tensor_reduce_affine_contiguous.h b/csrc/mctlass/include/mctlass/reduction/kernel/tensor_reduce_affine_contiguous.h new file mode 100644 index 0000000..54ac92d --- /dev/null +++ b/csrc/mctlass/include/mctlass/reduction/kernel/tensor_reduce_affine_contiguous.h @@ -0,0 +1,606 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Kernel performing a reduction over one or more ranks of an affine tensor +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/fast_math.h" +#include "mctlass/numeric_types.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/reduction/thread/reduction_operators.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace reduction { +namespace kernel { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Parameters structure +template < + int Rank, ///< Rank of source tensor (e.g. NDHWC => 5) + int ReducedRank, ///< Rank of reduced tensor (i.e. number of outer ranks) + typename ElementOutput, ///< Data type of output tensor + typename ElementSource, ///< Data type of source tensor + typename ReductionOp, ///< Reduction operator + int VectorLength = 1, ///< Vector length for memory + typename ElementCompute = ElementOutput, ///< Internal compute type - input type of reduction operation + int Threads = 256, ///< Number of participating threads + int BatchSize = 4 ///< Number of elements to load per batch +> +struct TensorReductionAffineContiguousParams { + + static int const kRank = Rank; + static int const kReducedRank = ReducedRank; + static int const kVectorLength = VectorLength; + static int const kInnerRank = kRank - kReducedRank; + static int const kThreads = Threads; + static int const kBatchSize = BatchSize; + + Coord extent; /// Extent of source tensor + FastDivmodU64 divmod[kRank - 1]; /// FastDivmod by each strided rank + int64_t dst_stride[kReducedRank]; /// stride (units of bytes) - I, J + int64_t src_stride[kRank - 1]; /// stride (units of bytes) - I, J, K + int64_t workspace_stride; /// stride (units of bytes) between workspace + int workspace_count; /// number of workspaces + + uint64_t inner_count; /// Number of elements in reduced index space + uint64_t outer_count; /// Number of elements in outer index space + + ElementOutput * destination; /// Pointer to output tensor of rank kReducedRank + ElementSource const * source; /// Poitner to source pointer of rank kRank + ReductionOp reduction_op; /// Reduction operator + ElementCompute reduction_identity; /// Identity element used by reduction operator + ElementCompute *device_workspace; /// Pointer to device workspace for inter-CTA reductions + + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + TensorReductionAffineContiguousParams() { + + } + + /// Ctor + TensorReductionAffineContiguousParams( + Coord extent_, ///< Extent of source tensor + ElementOutput * dst_ptr_, ///< Output tensor data + int64_t dst_stride_[], ///< Stride (units of elements) + ElementSource const * src_ptr_, ///< Source tensor data + int64_t src_stride_[], ///< Stride (units of elements) + ElementCompute *device_workspace_, ///< Pointer to device workspace for inter-CTA reductions + int64_t workspace_stride_, ///< Stride between workspaces + int workspace_count_, ///< Number of workspaces + ReductionOp reduction_op_, ///< Reduction operator + ElementCompute reduction_identity_ = ElementCompute() ///< Identity element used by reduction operator + ): + extent(extent_), + inner_count(1), + outer_count(1), + destination(dst_ptr_), + source(src_ptr_), + device_workspace(device_workspace_), + workspace_stride(workspace_stride_), + workspace_count(workspace_count_), + reduction_op(reduction_op_), + reduction_identity(reduction_identity_) { + + // Initialize divisors for fast div-mod + for (int p = 1; p < kRank; ++p) { + divmod[p - 1] = FastDivmodU64(uint64_t(extent[p])); + } + + int input_size_bits = sizeof_bits::value; + int output_size_bits = sizeof_bits::value; + + // Compute strides in units of bytes + for (int p = 0; p < kReducedRank; ++p) { + dst_stride[p] = dst_stride_[p] * output_size_bits / 8; + } + + for (int p = 0; p < kRank - 1; ++p) { + src_stride[p] = src_stride_[p] * input_size_bits / 8; + } + + // Compute number of elements in strided ranks + for (int p = 0; p < kReducedRank; ++p) { + outer_count *= uint64_t(extent[p]); + } + + for (int p = 0; p < kInnerRank; ++p) { + inner_count *= uint64_t(extent[kRank - 1 - p]); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Kernel to reduce a tensor with affine layout over a set of ranks *INCLUDING* the contiguous +/// rank. This leads to favorable vectorized memory accesses over the contiguous rank. +template < + int Rank, ///< Rank of source tensor (e.g. NDHWC => 5) + int ReducedRank, ///< Rank of reduced tensor (includes contiguous, e.g. NC => 2) + typename ElementOutput, ///< Data type of output tensor + typename ElementSource, ///< Data type of source tensor + typename ReductionOp, ///< Reduction operator + int VectorLength = 1, ///< Vector length for memory + typename ElementCompute = ElementOutput, ///< Internal compute type - input type of reduction operation + int Threads = 256, ///< Number of participating threads + int BatchSize = 4 ///< Number of elements to load per batch +> +class TensorReductionAffineContiguous { +public: + + static int const kRank = Rank; + static int const kReducedRank = ReducedRank; + static int const kVectorLength = VectorLength; + static int const kInnerRank = kRank - kReducedRank; + static int const kThreads = Threads; + static int const kBatchSize = BatchSize; + using ComputeFragment = Array; + using SourceFragment = AlignedArray; + using OutputFragment = AlignedArray; + + /// Shared memory allocation used for reduction within the CTA + struct SharedStorage { + Array workspace; + }; + + /// Parameters structure + using Params = TensorReductionAffineContiguousParams< + Rank, + ReducedRank, + ElementOutput, + ElementSource, + ReductionOp, + VectorLength, + ElementCompute, + Threads, + BatchSize + >; + +private: + + /// Computes the coordinate and offset of a given linear index + MCTLASS_DEVICE + void compute_inner_coord_and_offset_( + Params const ¶ms, + Coord & coord, + int64_t &src_offset, + uint64_t linear_idx) const { + + // Decompose into a coordinate of rank + coord = CoordinateDecomposition(linear_idx, ¶ms.divmod[kRank - kInnerRank]); + + // Compute an offset using the souce stride + src_offset = 0; + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kInnerRank - 1; ++i) { + src_offset += coord[i] * params.src_stride[kReducedRank + i]; + } + src_offset += coord[kInnerRank - 1] * sizeof_bits::value / 8; + } + + /// Computes the coordinate and offset of a given linear index + MCTLASS_DEVICE + void compute_outer_coord_and_offset_( + Params const ¶ms, + Coord & coord, + int64_t &dst_offset, + int64_t &src_offset, + uint64_t linear_idx) const { + + // Decompose into coordinate of rank + coord = CoordinateDecomposition(linear_idx, params.divmod); + + // Compute offsets using destination and source strides + dst_offset = 0; + src_offset = 0; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kReducedRank; ++i) { + dst_offset += params.dst_stride[i] * coord[i]; + src_offset += params.src_stride[i] * coord[i]; + } + } + + /// Reduces over the reduction indices yielding a single element + MCTLASS_DEVICE + ElementCompute reduce_indices_( + Params const ¶ms, + ElementCompute *threadblock_workspace, + char const *src_byte_ptr, + int coord_c) { + + NumericArrayConverter convert_source; + ReductionOp reduction_op(params.reduction_op); + + // + // Early exit or initialize to identity element + // + if (!params.inner_count) { + return params.reduction_identity; + } + + ComputeFragment accumulator; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < accumulator.size(); ++i) { + accumulator[i] = params.reduction_identity; + } + + // Compute the coordinate of the first access + int64_t src_byte_offset = 0; + Coord coord; + + uint64_t linear_idx = (threadIdx.x + blockDim.x * threadIdx.z + blockDim.x * blockIdx.z * blockDim.z) * kVectorLength; + compute_inner_coord_and_offset_(params, coord, src_byte_offset, linear_idx); + + // Load the first vector + SourceFragment source_fragment[kBatchSize]; + + bool not_done = true; + + // Iterate over vectors in a linearized reduction index space + while (not_done) { + + bool guards[kBatchSize]; + + // Issue a batch of loads + MCTLASS_PRAGMA_UNROLL + for (int b = 0; b < kBatchSize; ++b) { + + if (linear_idx < params.inner_count) { + source_fragment[b] = *reinterpret_cast(src_byte_ptr + src_byte_offset); + guards[b] = true; + } + else { + guards[b] = false; + not_done = false; + } + + linear_idx += (blockDim.z * gridDim.z * blockDim.x) * kVectorLength; + compute_inner_coord_and_offset_(params, coord, src_byte_offset, linear_idx); + } + + // Perform a batch of reduction operations + MCTLASS_PRAGMA_UNROLL + for (int b = 0; b < kBatchSize; ++b) { + if (guards[b]) { + auto cvt = convert_source(source_fragment[b]); + + accumulator = mctlass::reduction::thread::detail::ApplyArrayOperator( + reduction_op, + accumulator, + cvt); + } + } + }; + + // + // Reduction of vectors to scalar + // + + ElementCompute reduced_accumulator = accumulator[0]; + + MCTLASS_PRAGMA_UNROLL + for (int i = 1; i < kVectorLength; ++i) { + reduced_accumulator = reduction_op(reduced_accumulator, accumulator[i]); + } + + // + // Reduction within CTA across threadIdx.xz => threadIdx{.x = 0, .z = 0} + // + // This re-arranges data so threadIdx.y is effectively a row index and threadIdx.xz is a column + // + + int thread_count = blockDim.x * blockDim.z; + int thread_j = threadIdx.x + blockDim.x * threadIdx.z; + int thread_i = threadIdx.y; + + ElementCompute *frag_ptr = reinterpret_cast(threadblock_workspace) + thread_i * thread_count; + + frag_ptr[thread_j] = reduced_accumulator; + + // + // Reduce + // + MCTLASS_PRAGMA_NO_UNROLL + while (thread_count > 1) { + thread_count /= 2; + + __syncthreads(); + + if (thread_j < thread_count) { + ElementCompute other = frag_ptr[thread_j + thread_count]; + + reduced_accumulator = reduction_op(reduced_accumulator, other); + + frag_ptr[thread_j] = reduced_accumulator; + } + + __syncthreads(); + } + + + return reduced_accumulator; + } + +public: + + /// Perform a reduction + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + int coord_c = (blockIdx.x * blockDim.x + threadIdx.x) * kVectorLength; + + char const * src_byte_ptr = reinterpret_cast(params.source); + char * dst_byte_ptr = nullptr; + + // If performing a reduction across CTAs, redirect output to device workspace + if (gridDim.z == 1) { + dst_byte_ptr = reinterpret_cast(params.destination); + } + else { + dst_byte_ptr = reinterpret_cast(params.device_workspace); + } + + uint64_t idx_linear = blockIdx.y * blockDim.y + threadIdx.y; + + // Use modulo division to compute location + Coord outer_coord; + int64_t dst_byte_offset; + int64_t src_byte_offset; + + compute_outer_coord_and_offset_( + params, + outer_coord, + dst_byte_offset, + src_byte_offset, + idx_linear); + + if (gridDim.z == 1) { + + /// Complete the reduction with no workspace + while (idx_linear < params.outer_count) { + + ElementCompute result = reduce_indices_( + params, + shared_storage.workspace.data(), + src_byte_ptr + src_byte_offset, + coord_c); + + // Store the result after possible final reduction within the CTA + if (threadIdx.z == 0 && threadIdx.x == 0) { + + // Convert to output type and store + NumericConverter convert_output; + ElementOutput cvt = convert_output(result); + + *reinterpret_cast(dst_byte_ptr + dst_byte_offset) = cvt; + } + + __syncthreads(); + + // Update indices and pointers + idx_linear += gridDim.y * blockDim.y; + + compute_outer_coord_and_offset_( + params, + outer_coord, + dst_byte_offset, + src_byte_offset, + idx_linear); + + } // while + } + else { + + /// Complete the reduction with workspace + while (idx_linear < params.outer_count) { + + ElementCompute result = reduce_indices_( + params, + shared_storage.workspace.data(), + src_byte_ptr + src_byte_offset, + coord_c); + + int64_t byte_offset = + blockIdx.z * params.workspace_stride + idx_linear * sizeof_bits::value / 8; + + // Store the result for final reduction + if (threadIdx.z == 0 && threadIdx.x == 0) { + *reinterpret_cast(dst_byte_ptr + byte_offset) = result; + } + + __syncthreads(); + + // Update indices and pointers + idx_linear += gridDim.y * blockDim.y; + + compute_outer_coord_and_offset_( + params, + outer_coord, + dst_byte_offset, + src_byte_offset, + idx_linear); + } // while + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Kernel to perform final reduction +template < + int Rank, ///< Rank of source tensor (e.g. NDHWC => 5) + int ReducedRank, ///< Rank of reduced tensor (includes contiguous, e.g. NC => 2) + typename ElementOutput, ///< Data type of output tensor + typename ElementSource, ///< Data type of source tensor + typename ReductionOp, ///< Reduction operator + int VectorLength = 1, ///< Vector length for memory + typename ElementCompute = ElementOutput, ///< Internal compute type - input type of reduction operation + int Threads = 256, ///< Number of participating threads + int BatchSize = 4 ///< Number of elements to load per batch +> +class TensorReductionAffineContiguousFinal { +public: + + static int const kRank = Rank; + static int const kReducedRank = ReducedRank; + static int const kVectorLength = VectorLength; + static int const kInnerRank = kRank - kReducedRank; + static int const kThreads = Threads; + static int const kBatchSize = BatchSize; + + /// Shared memory + struct SharedStorage { }; + + /// Parameters structure + using Params = TensorReductionAffineContiguousParams< + Rank, + ReducedRank, + ElementOutput, + ElementSource, + ReductionOp, + VectorLength, + ElementCompute, + Threads, + BatchSize + >; + +private: + + /// Computes the coordinate and offset of a given linear index + MCTLASS_DEVICE + void compute_outer_coord_and_offset_( + Params const ¶ms, + Coord & coord, + int64_t &dst_offset, + uint64_t linear_idx) const { + + // Decompose into coordinate of rank + coord = CoordinateDecomposition(linear_idx, params.divmod); + + // Compute offsets using destination and source strides + dst_offset = 0; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kReducedRank; ++i) { + dst_offset += params.dst_stride[i] * coord[i]; + } + } + + /// Reduces over the reduction indices + MCTLASS_DEVICE + ElementCompute reduce_indices_( + Params const ¶ms, + ElementCompute const *device_workspace) { + + ReductionOp reduction_op(params.reduction_op); + char const *src_byte_ptr = reinterpret_cast(device_workspace); + + // Accumulated output + ElementCompute accumulator = params.reduction_identity; + + for (int iter = 0; iter < params.workspace_count; ++iter) { + ElementCompute workspace_item = *reinterpret_cast(src_byte_ptr); + + accumulator = reduction_op(accumulator, workspace_item); + + src_byte_ptr += params.workspace_stride; + } + + return accumulator; + } + +public: + + // + // Methods + // + + /// Perform a reduction + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + uint64_t idx_linear = blockIdx.x * blockDim.x + threadIdx.x; + + char * dst_byte_ptr = reinterpret_cast(params.destination); + + // Use modulo division to compute location + Coord outer_coord; + int64_t dst_byte_offset; + + compute_outer_coord_and_offset_( + params, + outer_coord, + dst_byte_offset, + idx_linear); + + /// Complete the reduction + while (idx_linear < params.outer_count) { + + ElementCompute result = reduce_indices_(params, params.device_workspace + idx_linear); + + // Convert to output type and store + NumericConverter convert_output; + + *reinterpret_cast(dst_byte_ptr + dst_byte_offset) = convert_output(result); + + // Update indices and pointers + idx_linear += gridDim.x * blockDim.x; + + compute_outer_coord_and_offset_( + params, + outer_coord, + dst_byte_offset, + idx_linear); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace reduction +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/reduction/kernel/tensor_reduce_affine_strided.h b/csrc/mctlass/include/mctlass/reduction/kernel/tensor_reduce_affine_strided.h new file mode 100644 index 0000000..bcdb700 --- /dev/null +++ b/csrc/mctlass/include/mctlass/reduction/kernel/tensor_reduce_affine_strided.h @@ -0,0 +1,641 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Kernel performing a reduction over one or more ranks of an affine tensor +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/fast_math.h" +#include "mctlass/numeric_types.h" +#include "mctlass/numeric_conversion.h" +#include "mctlass/device_kernel.h" + +#include "mctlass/reduction/thread/reduction_operators.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace reduction { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace kernel { + +/// Parameters structure +template < + int Rank, ///< Rank of source tensor (e.g. NDHWC => 5) + int ReducedRank, ///< Rank of reduced tensor (includes contiguous, e.g. NC => 2) + typename ElementOutput, ///< Data type of output tensor + typename ElementSource, ///< Data type of source tensor + typename ReductionOp, ///< Reduction operator + int VectorLength = 1, ///< Vector length for memory + typename ElementCompute = ElementOutput, ///< Internal compute type - input type of reduction operation + int Threads = 256, ///< Number of participating threads + int BatchSize = 4 ///< Number of elements to load per batch +> +struct TensorReductionAffineStridedParams { + + static int const kRank = Rank; + static int const kReducedRank = ReducedRank; + static int const kVectorLength = VectorLength; + static int const kInnerRank = kRank - kReducedRank; + static int const kThreads = Threads; + static int const kBatchSize = BatchSize; + + Coord extent; /// Extent of source tensor + FastDivmodU64 divmod[kRank - 1]; /// FastDivmod by each strided rank + int64_t dst_stride[kReducedRank - 1]; /// stride (units of bytes) - I, J + int64_t src_stride[kRank - 1]; /// stride (units of bytes) - I, J, K + int64_t workspace_stride; /// stride (units of bytes) between workspace + int64_t workspace_outer_stride; /// stride (units of bytes) between 'rows' of the workspace + int workspace_count; /// number of workspaces + + uint64_t inner_count; /// Number of elements in reduced index space + uint64_t outer_count; /// Number of elements in outer index space + + ElementOutput * destination; /// Pointer to output tensor of rank kReducedRank + ElementSource const * source; /// Poitner to source pointer of rank kRank + ReductionOp reduction_op; /// Reduction operator + ElementCompute reduction_identity; /// Identity element for reduction operator + ElementCompute *device_workspace; /// Pointer to device workspace for inter-CTA reductions + + // + // Methods + // + + /// Ctor + MCTLASS_HOST_DEVICE + TensorReductionAffineStridedParams() { + + } + + /// Ctor + TensorReductionAffineStridedParams( + Coord extent_, ///< Extent of source tensor + ElementOutput * dst_ptr_, ///< Output tensor data + int64_t dst_stride_[], ///< Stride (units of elements) + ElementSource const * src_ptr_, ///< Source tensor data + int64_t src_stride_[], ///< Stride (units of elements) + ElementCompute *device_workspace_, ///< Pointer to device workspace for inter-CTA reductions + int64_t workspace_stride_, ///< Stride between workspaces + int workspace_count_, ///< Number of workspaces + ReductionOp reduction_op_, ///< Reduction operator + ElementCompute reduction_identity_ = ElementCompute() ///< Identity element for reduction operator + ): + extent(extent_), + inner_count(1), + outer_count(1), + destination(dst_ptr_), + source(src_ptr_), + device_workspace(device_workspace_), + workspace_outer_stride(0), + workspace_stride(workspace_stride_), + workspace_count(workspace_count_), + reduction_op(reduction_op_), + reduction_identity(reduction_identity_) { + + // Initialize divisors for fast div-mod + for (int p = 1; p < kRank; ++p) { + divmod[p - 1] = FastDivmodU64(uint64_t(extent[p])); + } + + int input_size_bits = sizeof_bits::value; + int output_size_bits = sizeof_bits::value; + + workspace_outer_stride = workspace_stride * workspace_count; + + // Compute strides in units of bytes + for (int p = 0; p < kReducedRank - 1; ++p) { + dst_stride[p] = dst_stride_[p] * output_size_bits / 8; + } + + for (int p = 0; p < kRank - 1; ++p) { + src_stride[p] = src_stride_[p] * input_size_bits / 8; + } + + // Compute number of elements in strided ranks + for (int p = 0; p < kReducedRank - 1; ++p) { + outer_count *= uint64_t(extent[p]); + } + + for (int p = 0; p < kInnerRank; ++p) { + inner_count *= uint64_t(extent[kReducedRank + p - 1]); + } + } +}; + +/// Kernel to reduce a tensor with affine layout over a set of ranks *EXCLUDING* the contiguous +/// rank. This leads to favorable vectorized memory accesses over the contiguous rank. +template < + int Rank, ///< Rank of source tensor (e.g. NDHWC => 5) + int ReducedRank, ///< Rank of reduced tensor (includes contiguous, e.g. NC => 2) + typename ElementOutput, ///< Data type of output tensor + typename ElementSource, ///< Data type of source tensor + typename ReductionOp, ///< Reduction operator + int VectorLength = 1, ///< Vector length for memory + typename ElementCompute = ElementOutput, ///< Internal compute type - input type of reduction operation + int Threads = 256, ///< Number of participating threads + int BatchSize = 4 ///< Number of elements to load per batch +> +class TensorReductionAffineStrided { +public: + + static int const kRank = Rank; + static int const kReducedRank = ReducedRank; + static int const kVectorLength = VectorLength; + static int const kInnerRank = kRank - kReducedRank; + static int const kThreads = Threads; + static int const kBatchSize = BatchSize; + using ComputeFragment = Array; + using SourceFragment = AlignedArray; + using OutputFragment = AlignedArray; + + /// Shared memory allocation used for reduction within the CTA + struct SharedStorage { + Array workspace; + }; + + /// Parameters structure + using Params = TensorReductionAffineStridedParams< + Rank, + ReducedRank, + ElementOutput, + ElementSource, + ReductionOp, + VectorLength, + ElementCompute, + Threads, + BatchSize + >; + +private: + + /// Computes the coordinate and offset of a given linear index + MCTLASS_DEVICE + void compute_inner_coord_and_offset_( + Params const ¶ms, + Coord & coord, + int64_t &src_offset, + uint64_t linear_idx) const { + + // Decompose into coordinate + coord = CoordinateDecomposition(linear_idx, ¶ms.divmod[kReducedRank - 1]); + + // Compute linear offset + src_offset = 0; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kInnerRank; ++i) { + src_offset += params.src_stride[kReducedRank + i - 1] * coord[i]; + } + } + + /// Computes the coordinate and offset of a given linear index + MCTLASS_DEVICE + void compute_outer_coord_and_offset_( + Params const ¶ms, + Coord & coord, + int64_t &dst_offset, + int64_t &src_offset, + uint64_t linear_idx) const { + + // Decompose linear coordinate + coord = CoordinateDecomposition(linear_idx, params.divmod); + + // Compute offset into tensors + dst_offset = 0; + src_offset = 0; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kReducedRank - 1; ++i) { + dst_offset += params.dst_stride[i] * coord[i]; + src_offset += params.src_stride[i] * coord[i]; + } + } + + /// Reduces over the reduction indices + MCTLASS_DEVICE + ComputeFragment reduce_indices_( + Params const ¶ms, + ElementCompute *threadblock_workspace, + char const *src_byte_ptr) { + + NumericArrayConverter convert_source; + ReductionOp reduction_op(params.reduction_op); + + // Accumulated output + ComputeFragment identity_frag; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < identity_frag.size(); ++i) { + identity_frag[i] = params.reduction_identity; + } + + if (!params.inner_count) { + return identity_frag; + } + + ComputeFragment accumulator = identity_frag; + + // Compute the coordinate of the first access + int64_t src_byte_offset = 0; + Coord coord; + + uint64_t linear_idx = threadIdx.z + blockIdx.z * blockDim.z; + compute_inner_coord_and_offset_(params, coord, src_byte_offset, linear_idx); + + // Load the first vector + SourceFragment source_fragment[kBatchSize]; + + bool not_done = true; + + // Iterate over vectors in a linearized reduction index space + while (not_done) { + + bool guards[kBatchSize]; + + // Issue a batch of loads + MCTLASS_PRAGMA_UNROLL + for (int b = 0; b < kBatchSize; ++b) { + + if (linear_idx < params.inner_count) { + source_fragment[b] = *reinterpret_cast(src_byte_ptr + src_byte_offset); + guards[b] = true; + } + else { + guards[b] = false; + not_done = false; + } + + linear_idx += blockDim.z * gridDim.z; + compute_inner_coord_and_offset_(params, coord, src_byte_offset, linear_idx); + } + + // Perform a batch of reduction operations + MCTLASS_PRAGMA_UNROLL + for (int b = 0; b < kBatchSize; ++b) { + if (guards[b]) { + + auto cvt = convert_source(source_fragment[b]); + + accumulator = mctlass::reduction::thread::detail::ApplyArrayOperator( + reduction_op, + accumulator, + cvt); + } + } + }; + + // Optional reduction within a CTA + if (blockDim.z > 1) { + + // Linearized thread ID + int thread_idx = threadIdx.x + blockDim.x * (threadIdx.y + blockDim.y * threadIdx.z); + + // all threads store to workspace + ComputeFragment *frag_ptr = reinterpret_cast(threadblock_workspace); + + frag_ptr[thread_idx] = accumulator; + + __syncthreads(); + + if (threadIdx.z == 0) { + // Load all additional block indices + for (int z = 1; z < blockDim.z; ++z) { + ComputeFragment frag = frag_ptr[thread_idx + z * blockDim.x * blockDim.y]; + + accumulator = mctlass::reduction::thread::detail::ApplyArrayOperator( + reduction_op, + accumulator, + frag); + } + } + + __syncthreads(); + } + + return accumulator; + } + +public: + + /// Perform a reduction + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + int coord_c = (blockIdx.x * blockDim.x + threadIdx.x) * kVectorLength; + + char const * src_byte_ptr = reinterpret_cast(params.source + coord_c); + char * dst_byte_ptr = nullptr; + + // If performing a reduction across CTAs, redirect output to device workspace + if (gridDim.z == 1) { + dst_byte_ptr = reinterpret_cast(params.destination + coord_c); + } + else { + dst_byte_ptr = reinterpret_cast(params.device_workspace + coord_c); + } + + // If the C index is out of bounds, exit + if (coord_c >= params.extent[kRank - 1]) { + return; + } + + int64_t idx_linear = blockIdx.y * blockDim.y + threadIdx.y; + + // Use modulo division to compute location + Coord outer_coord; + int64_t dst_byte_offset; + int64_t src_byte_offset; + + compute_outer_coord_and_offset_( + params, + outer_coord, + dst_byte_offset, + src_byte_offset, + idx_linear); + + if (gridDim.z == 1) { + + /// Complete the reduction with no workspace + while (idx_linear < params.outer_count) { + + ComputeFragment result; + + result = reduce_indices_( + params, + shared_storage.workspace.data(), + src_byte_ptr + src_byte_offset); + + // Store the result after possible final reduction within the CTA + if (threadIdx.z == 0) { + + // Convert to output type and store + NumericArrayConverter convert_output; + auto cvt = convert_output(result); + + *reinterpret_cast(dst_byte_ptr + dst_byte_offset) = + reinterpret_cast(cvt); + } + + // Update indices and pointers + idx_linear += gridDim.y * blockDim.y; + + compute_outer_coord_and_offset_( + params, + outer_coord, + dst_byte_offset, + src_byte_offset, + idx_linear); + + } // while + } + else { + + /// Complete the reduction with a device workspace + while (idx_linear < params.outer_count) { + + ComputeFragment result; + + result = reduce_indices_( + params, + shared_storage.workspace.data(), + src_byte_ptr + src_byte_offset); + + // Store the result after possible final reduction within the CTA + if (threadIdx.z == 0) { + + int64_t byte_offset = + blockIdx.z * params.workspace_stride + idx_linear * params.workspace_outer_stride; + + // No conversion - store in compute type + *reinterpret_cast(dst_byte_ptr + byte_offset) = + reinterpret_cast(result); + } + + // Update indices and pointers + idx_linear += gridDim.y * blockDim.y; + + compute_outer_coord_and_offset_( + params, + outer_coord, + dst_byte_offset, + src_byte_offset, + idx_linear); + + } // while (outer index) + } // if () + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Kernel to perform final reduction +template < + int Rank, ///< Rank of source tensor (e.g. NDHWC => 5) + int ReducedRank, ///< Rank of reduced tensor (includes contiguous, e.g. NC => 2) + typename ElementOutput, ///< Data type of output tensor + typename ElementSource, ///< Data type of source tensor + typename ReductionOp, ///< Reduction operator + int VectorLength = 1, ///< Vector length for memory + typename ElementCompute = ElementOutput, ///< Internal compute type - input type of reduction operation + int Threads = 256, ///< Number of participating threads + int BatchSize = 4 ///< Number of elements to load per batch +> +class TensorReductionAffineStridedFinal { +public: + + static int const kRank = Rank; + static int const kReducedRank = ReducedRank; + static int const kVectorLength = VectorLength; + static int const kInnerRank = kRank - kReducedRank; + static int const kThreads = Threads; + static int const kBatchSize = BatchSize; + using ComputeFragment = Array; + using SourceFragment = AlignedArray; + using OutputFragment = AlignedArray; + + /// Shared memory + struct SharedStorage { }; + + /// Parameters structure + using Params = TensorReductionAffineStridedParams< + Rank, + ReducedRank, + ElementOutput, + ElementSource, + ReductionOp, + VectorLength, + ElementCompute, + Threads, + BatchSize + >; + +private: + + /// Computes the coordinate and offset of a given linear index + MCTLASS_DEVICE + void compute_outer_coord_and_offset_( + Params const ¶ms, + Coord & coord, + int64_t &dst_offset, + uint64_t linear_idx) const { + + // Decompose linear index + coord = CoordinateDecomposition(linear_idx, params.divmod); + + // Compute tensor offset + dst_offset = 0; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kReducedRank - 1; ++i) { + dst_offset += params.dst_stride[i] * coord[i]; + } + } + + /// Reduces over the reduction indices + MCTLASS_DEVICE + ComputeFragment reduce_indices_( + Params const ¶ms, + char *src_byte_ptr) { + + ReductionOp reduction_op(params.reduction_op); + + // Accumulated output + ComputeFragment identity_frag; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < identity_frag.size(); ++i) { + identity_frag[i] = params.reduction_identity; + } + + ComputeFragment accumulator = identity_frag; + ComputeFragment workspace_fragments[kBatchSize]; + + // Partially unrolled loop + for (int idx = 0; idx < params.workspace_count; idx += kBatchSize) { + + // Issue a batch of loads + MCTLASS_PRAGMA_UNROLL + for (int b = 0; b < kBatchSize; ++b) { + if (idx + b < params.workspace_count) { + workspace_fragments[b] = + *reinterpret_cast(src_byte_ptr); + } + else { + workspace_fragments[b] = identity_frag; + } + src_byte_ptr += + params.workspace_stride; + } + + // Perform a reduction + MCTLASS_PRAGMA_UNROLL + for (int b = 0; b < kBatchSize; ++b) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kVectorLength; ++i) { + accumulator[i] = reduction_op(accumulator[i], workspace_fragments[b][i]); + } + } + } + + return accumulator; + } + +public: + + // + // Methods + // + + /// Perform a reduction + MCTLASS_DEVICE + void operator()(Params const ¶ms, SharedStorage &shared_storage) { + + int coord_c = (blockIdx.x * blockDim.x + threadIdx.x) * kVectorLength; + + char * src_byte_ptr = reinterpret_cast(params.device_workspace + coord_c); + char * dst_byte_ptr = reinterpret_cast(params.destination + coord_c); + + // If the C index is out of bounds, exit + if (coord_c >= params.extent[kRank - 1]) { + return; + } + + int64_t idx_linear = blockIdx.y * blockDim.y + threadIdx.y; + + // Use modulo division to compute location + Coord outer_coord; + int64_t dst_byte_offset; + + compute_outer_coord_and_offset_( + params, + outer_coord, + dst_byte_offset, + idx_linear); + + /// Complete the reduction + while (idx_linear < params.outer_count) { + + int64_t src_byte_offset = idx_linear * params.workspace_outer_stride; + + ComputeFragment result = reduce_indices_( + params, + src_byte_ptr + src_byte_offset); + + // Convert to output type and store + NumericArrayConverter convert_output; + auto cvt = convert_output(result); + + *reinterpret_cast(dst_byte_ptr + dst_byte_offset) = + reinterpret_cast(cvt); + + // Update indices and pointers + idx_linear += gridDim.y * blockDim.y; + + compute_outer_coord_and_offset_( + params, + outer_coord, + dst_byte_offset, + idx_linear); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace kernel +} // namespace reduction +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/reduction/thread/reduce.h b/csrc/mctlass/include/mctlass/reduction/thread/reduce.h new file mode 100644 index 0000000..afc31a5 --- /dev/null +++ b/csrc/mctlass/include/mctlass/reduction/thread/reduce.h @@ -0,0 +1,234 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines basic thread level reduction with specializations for Array. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/half.h" +#include "mctlass/functional.h" + +namespace mctlass { +namespace reduction { +namespace thread { + +/// Structure to compute the thread level reduction +template +struct Reduce; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial Specialization of Reduce for "plus" (a functional operator) +template +struct Reduce< plus, T > { + + MCTLASS_HOST_DEVICE + T operator()(T lhs, T const &rhs) const { + plus _op; + return _op(lhs, rhs); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specialization of Reduce for Array +template +struct Reduce < plus, Array> { + + MCTLASS_HOST_DEVICE + Array operator()(Array const &in) const { + + Array result; + Reduce< plus, T > scalar_reduce; + result.clear(); + + MCTLASS_PRAGMA_UNROLL + for (auto i = 0; i < N; ++i) { + result[0] = scalar_reduce(result[0], in[i]); + } + + return result; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specializations of Reduce for Array +template +struct Reduce < plus, Array > { + + MCTLASS_HOST_DEVICE + Array operator()(Array const &input) { + + Array result; + + // If there is only 1 element - there is nothing to reduce + if( N ==1 ){ + + result[0] = input.front(); + + } else { + + #if defined(__MACA_ARCH__) + + __half result_d; + Array const *in_ptr_half = reinterpret_cast const *>(&input); + Array const *in_ptr_half2 = reinterpret_cast const *>(&input); + __half2 const *x_in_half2 = reinterpret_cast<__half2 const *>(in_ptr_half2); + + // Set initial result = first half2, in case N==2 + __half2 tmp_result = x_in_half2[0]; + + MCTLASS_PRAGMA_UNROLL + for (int i = 1; i < N/2; ++i) { + + tmp_result = __hadd2(x_in_half2[i], tmp_result); + + } + + result_d = __hadd(__low2half(tmp_result), __high2half(tmp_result)); + + // One final step is needed for odd "N" (to add the (N-1)th element) + if( N%2 ){ + + __half last_element; + Array tmp_last; + Array *tmp_last_ptr = &tmp_last; + tmp_last_ptr[0] = in_ptr_half[N-1]; + last_element = reinterpret_cast<__half const &>(tmp_last); + + result_d = __hadd(result_d, last_element); + + } + + Array *result_ptr = &result; + *result_ptr = reinterpret_cast &>(result_d); + + #else + + Reduce< plus, half_t > scalar_reduce; + result.clear(); + + MCTLASS_PRAGMA_UNROLL + for (auto i = 0; i < N; ++i) { + + result[0] = scalar_reduce(result[0], input[i]); + + } + + #endif + } + + return result; + + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Partial specializations of Reduce for AlignedArray +template +struct Reduce < plus, AlignedArray > { + + MCTLASS_HOST_DEVICE + Array operator()(AlignedArray const &input) { + + Array result; + + // If there is only 1 element - there is nothing to reduce + if( N ==1 ){ + + result[0] = input.front(); + + } else { + + #if defined(__MACA_ARCH__) + + __half result_d; + AlignedArray const *in_ptr_half = reinterpret_cast const *>(&input); + AlignedArray const *in_ptr_half2 = reinterpret_cast const *>(&input); + __half2 const *x_in_half2 = reinterpret_cast<__half2 const *>(in_ptr_half2); + + // Set initial result = first half2, in case N==2 + __half2 tmp_result = x_in_half2[0]; + + MCTLASS_PRAGMA_UNROLL + for (int i = 1; i < N/2; ++i) { + + tmp_result = __hadd2(x_in_half2[i], tmp_result); + + } + + result_d = __hadd(__low2half(tmp_result), __high2half(tmp_result)); + + // One final step is needed for odd "N" (to add the (N-1)th element) + if( N%2 ){ + + __half last_element; + AlignedArray tmp_last; + AlignedArray *tmp_last_ptr = &tmp_last; + tmp_last_ptr[0] = in_ptr_half[N-1]; + last_element = reinterpret_cast<__half const &>(tmp_last); + + result_d = __hadd(result_d, last_element); + + } + + Array *result_ptr = &result; + *result_ptr = reinterpret_cast &>(result_d); + + #else + + Reduce< plus, half_t > scalar_reduce; + result.clear(); + + MCTLASS_PRAGMA_UNROLL + for (auto i = 0; i < N; ++i) { + + result[0] = scalar_reduce(result[0], input[i]); + + } + + #endif + } + + return result; + + } +}; +} +} +} diff --git a/csrc/mctlass/include/mctlass/reduction/thread/reduction_operators.h b/csrc/mctlass/include/mctlass/reduction/thread/reduction_operators.h new file mode 100644 index 0000000..94d1c6b --- /dev/null +++ b/csrc/mctlass/include/mctlass/reduction/thread/reduction_operators.h @@ -0,0 +1,235 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Kernel performing a reduction over densely packed tensors in global memory +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/numeric_types.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" +#include "mctlass/numeric_conversion.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace reduction { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Mixed-precision reduction +template < + typename ElementAccumulator_, + typename Element_, + int Count = 1 +> +struct ReduceAdd { + + // + // Type definitions + // + + using ElementAccumulator = ElementAccumulator_; + using Element = Element_; + static int const kCount = Count; + + using FragmentAccumulator = mctlass::Array; + using FragmentElement = mctlass::Array; + + struct Params { }; + + // + // Data members + // + + /// Parameters object + Params params; + + // + // Methods + // + + /// Constructor + MCTLASS_HOST_DEVICE + ReduceAdd(Params params_ = Params()): params(params_) { } + + /// Operator + MCTLASS_HOST_DEVICE + FragmentAccumulator operator()( + FragmentAccumulator accumulator, + FragmentElement element) const { + + plus op; + + NumericArrayConverter< + ElementAccumulator, + Element, + kCount, + PreferredRoundingMode::kRound> converter; + + return op(accumulator, converter(element)); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +/// Special handling for binary operators +template +struct VectorizeArrayOperation { + + using ValueType = Array; + + MCTLASS_HOST_DEVICE + ValueType operator()( + ReductionOp const &reduction_op, + ValueType const &lhs, + ValueType const &rhs) const { + + ValueType result; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + result[i] = reduction_op(lhs[i], rhs[i]); + } + + return result; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct ReduceArrayOperation { + + using ArrayType = Array; + + MCTLASS_HOST_DEVICE + Element operator()( + ReductionOp const &reduction_op, + ArrayType const &array) const { + + Element item = reduction_op(array[0], array[1]); + + MCTLASS_PRAGMA_UNROLL + for (int i = 2; i < N; ++i) { + item = reduction_op(item, array[i]); + } + + return item; + } +}; + +template +struct ReduceArrayOperation, uint1b_t, N> { + + using ArrayType = Array; + + MCTLASS_HOST_DEVICE + uint1b_t operator()( + logical_and const &reduction_op, + ArrayType const &array) const { + + uint8_t const *ptr = reinterpret_cast(&array); + bool item = false; + + MCTLASS_PRAGMA_UNROLL + for (int byte = 0; byte < (N + 7) / 8; ++byte) { + uint8_t bits = ptr[byte]; + item = (item || !bits); + } + + return uint1b_t(!item); + } +}; + +template +struct ReduceArrayOperation, uint1b_t, N> { + + using ArrayType = Array; + + MCTLASS_HOST_DEVICE + uint1b_t operator()( + logical_and const &reduction_op, + ArrayType const &array) const { + + uint8_t const *ptr = reinterpret_cast(&array); + bool item = true; + + MCTLASS_PRAGMA_UNROLL + for (int byte = 0; byte < (N + 7) / 8; ++byte) { + uint8_t bits = ptr[byte]; + item = (item || bits); + } + + return uint1b_t(item); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Helper function to infer template argument types +template +MCTLASS_HOST_DEVICE +Array ApplyArrayOperator( + ReductionOp const &reduction_op, + Array const &lhs, + Array const &rhs) { + + VectorizeArrayOperation vectorize_op; + + return vectorize_op(reduction_op, lhs, rhs); +} + +/// Helper to reduce an array +template +Element ReduceArray(ReductionOp const &reduction_op, Array const &array) { + ReduceArrayOperation reduce_array_op; + + return reduce_array_op(reduction_op, array); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace reduction +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/reduction/threadblock_swizzle.h b/csrc/mctlass/include/mctlass/reduction/threadblock_swizzle.h new file mode 100644 index 0000000..3bc9f1a --- /dev/null +++ b/csrc/mctlass/include/mctlass/reduction/threadblock_swizzle.h @@ -0,0 +1,67 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +* +**************************************************************************************************/ +/*! \file +\brief Defies functors for mapping blockIdx to partitions of the batched reduction computation. +*/ +#pragma once +#include "mctlass/coord.h" + +namespace mctlass { +namespace reduction { +struct DefaultBlockSwizzle { + /// Ctor + MCTLASS_HOST_DEVICE DefaultBlockSwizzle() {} + + /// Swizzle the block index. + MCTLASS_DEVICE dim3 swizzle() { return blockIdx; } + + /// + MCTLASS_HOST_DEVICE dim3 get_grid_layout(Coord<3> const &problem_size, + Coord<3> const &OutputTile) { + assert(OutputTile[0] == 1 && OutputTile[1] == 1); + assert((problem_size[0] * problem_size[1] * problem_size[2]) % OutputTile[2] == 0); + dim3 grid; + grid.x = problem_size[0] * problem_size[1] * problem_size[2] + / OutputTile[2] ; + return grid; + } + + /// + MCTLASS_DEVICE Coord<3> get_threadblock_offset(Coord<3> const &SubTile) { + assert(SubTile[0] == 1 && SubTile[1] == 1); + dim3 block = swizzle(); + Coord<3> threadblock_offset = + make_Coord(0, 0, block.x * SubTile[2]); + return threadblock_offset; + } +}; +} // namespace reduction +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/relatively_equal.h b/csrc/mctlass/include/mctlass/relatively_equal.h new file mode 100644 index 0000000..ff86ad0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/relatively_equal.h @@ -0,0 +1,219 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/* \file + \brief Performs comparison between two elements with support for floating-point comparisons. +*/ + +#pragma once + +#include "numeric_types.h" + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +MCTLASS_HOST_DEVICE +bool relatively_equal(T a, T b, T epsilon, T nonzero_floor); + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { + +// This floating-point comparison function implements the method described in +// +// https://floating-point-gui.de/errors/comparison/ +// +template +MCTLASS_HOST_DEVICE +bool relatively_equal_float(T a, T b, T epsilon, T nonzero_floor) { + +#if defined(__MACACC_RTC__) + using std::abs; // using maca::std::abs; +#else + using std::abs; +#endif + + T abs_A = abs(a); + T abs_B = abs(b); + T diff = abs(a - b); + T zero = T(0); + + if (a == b) { + return true; + } + else if (a == zero || b == zero || diff < nonzero_floor) { + return diff < epsilon * nonzero_floor; + } + + return diff < epsilon * (abs_A + abs_B); +} + +} // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(uint1b_t a, uint1b_t b, uint1b_t, uint1b_t) { + return (a == b); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(int2b_t a, int2b_t b, int2b_t, int2b_t) { + return (a == b); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(uint2b_t a, uint2b_t b, uint2b_t, uint2b_t) { + return (a == b); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(int4b_t a, int4b_t b, int4b_t, int4b_t) { + return (a == b); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(uint4b_t a, uint4b_t b, uint4b_t, uint4b_t) { + return (a == b); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(int8_t a, int8_t b, int8_t, int8_t) { + return (a == b); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(uint8_t a, uint8_t b, uint8_t, uint8_t) { + return (a == b); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(int16_t a, int16_t b, int16_t, int16_t) { + return (a == b); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(uint16_t a, uint16_t b, uint16_t, uint16_t) { + return (a == b); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(int32_t a, int32_t b, int32_t, int32_t) { + return (a == b); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(uint32_t a, uint32_t b, uint32_t, uint32_t) { + return (a == b); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(int64_t a, int64_t b, int64_t, int64_t) { + return (a == b); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(uint64_t a, uint64_t b, uint64_t, uint64_t) { + return (a == b); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(float_e4m3_t a, float_e4m3_t b, float_e4m3_t epsilon, float_e4m3_t nonzero_floor) { + return detail::relatively_equal_float(a, b, epsilon, nonzero_floor); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(float_e5m2_t a, float_e5m2_t b, float_e5m2_t epsilon, float_e5m2_t nonzero_floor) { + return detail::relatively_equal_float(a, b, epsilon, nonzero_floor); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(half_t a, half_t b, half_t epsilon, half_t nonzero_floor) { + return detail::relatively_equal_float(a, b, epsilon, nonzero_floor); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal( + bfloat16_t a, + bfloat16_t b, + bfloat16_t epsilon, + bfloat16_t nonzero_floor) { + + return detail::relatively_equal_float(a, b, epsilon, nonzero_floor); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal( + tfloat32_t a, + tfloat32_t b, + tfloat32_t epsilon, + tfloat32_t nonzero_floor) { + + return detail::relatively_equal_float(a, b, epsilon, nonzero_floor); +} + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(float a, float b, float epsilon, float nonzero_floor) { + return detail::relatively_equal_float(a, b, epsilon, nonzero_floor); +} + + +template <> +MCTLASS_HOST_DEVICE +bool relatively_equal(double a, double b, double epsilon, double nonzero_floor) { + return detail::relatively_equal_float(a, b, epsilon, nonzero_floor); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/semaphore.h b/csrc/mctlass/include/mctlass/semaphore.h new file mode 100644 index 0000000..06ac3e2 --- /dev/null +++ b/csrc/mctlass/include/mctlass/semaphore.h @@ -0,0 +1,125 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Implementation of a CTA-wide semaphore for inter-CTA synchronization. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" + +#include "mctlass/numeric_types.h" +#include "mctlass/matrix_shape.h" + +#include "mctlass/gemm/gemm.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// CTA-wide semaphore for inter-CTA synchronization. +class Semaphore { +public: + + int *lock; + bool wait_thread; + int state; + +public: + + /// Implements a semaphore to wait for a flag to reach a given value + MCTLASS_HOST_DEVICE + Semaphore(int *lock_, int thread_id): + lock(lock_), + wait_thread(thread_id < 0 || thread_id == 0), + state(-1) { + + } + + /// Permit fetching the synchronization mechanism early + MCTLASS_DEVICE + void fetch() { + if (wait_thread) { + // #if defined(__MACA_ARCH__) + // asm volatile ("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); + // #else + // asm volatile ("ld.global.cg.b32 %0, [%1];\n" : "=r"(state) : "l"(lock)); + // #endif +#if defined(__MACA_ARCH__) + state = __ldcg(lock); +#endif + } + } + + /// Gets the internal state + MCTLASS_DEVICE + int get_state() const { + return state; + } + + /// Waits until the semaphore is equal to the given value + MCTLASS_DEVICE + void wait(int status = 0) { + while( __syncthreads_and(state != status) ) { + fetch(); + } + + __syncthreads(); + } + + /// Updates the lock with the given result + MCTLASS_DEVICE + void release(int status = 0) { + __syncthreads(); + + if (wait_thread) { + #if 0 + #if defined(__MACA_ARCH__) + asm volatile ("st.global.release.gpu.b32 [%0], %1;\n" : : "l"(lock), "r"(status)); + #else + asm volatile ("st.global.cg.b32 [%0], %1;\n" : : "l"(lock), "r"(status)); + #endif + #else + __stcg(lock, status); + #endif + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/subbyte_reference.h b/csrc/mctlass/include/mctlass/subbyte_reference.h new file mode 100644 index 0000000..5f5f30a --- /dev/null +++ b/csrc/mctlass/include/mctlass/subbyte_reference.h @@ -0,0 +1,633 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Provides a mechanism for packing and unpacking elements smaller than one byte +*/ +#pragma once + +#include "mctlass/numeric_types.h" + +namespace mctlass { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// This class provides a mechanism for packing and unpacking elements smaller than one byte. It +/// assumes these sub-byte elements are packed in a traditional C++ numeric type. +/// +/// The intended application is to provide a mechanism to indirectly reference elements in +/// memory or Array<> objects whose addresses cannot otherwise be taken since they are smaller +/// than one byte. +/// +/// Supports basic pointer arithmetic: +/// +/// Example: +/// +/// int4b_t *ptr = ...; +/// +/// SubbyteReference ref = ptr; +/// ref += 15; +/// +/// int4b_t x = ref; // load an int4b_t +/// ref = x + 2_s4; // perform arithmetic on int4b_t and then store +/// +template < + typename Element_, /// MCTLASS numeric element type. + typename Storage_ = uint8_t /// Underlying storage type. Must be able to hold an integer + /// number of objects of type Element. +> +class ConstSubbyteReference { +public: + + using Element = Element_; + using Storage = Storage_; + using StoragePointer = Storage const *; + + static_assert(sizeof_bits::value <= sizeof_bits::value, + "Size of Element must not be greater than Storage."); + + static_assert(!(sizeof_bits::value % sizeof_bits::value), + "Storage must be divisible by Element"); + +private: + + ///! Number of elements per storage vector + int const kElementsPerVector = sizeof_bits::value / sizeof_bits::value; + + ///! Bit mask + Storage const kMask = + ((sizeof_bits::value < sizeof_bits::value) ? + (Storage(1) << sizeof_bits::value) - Storage(1) : + ~Storage(0)); + +private: + + /// Pointer to array containing element + StoragePointer ptr_; + + /// Offset (in units of elements) from pointer. + /// + /// Invariant: must always be in range [0, kElementsPerVector) + int offset_; + +public: + + MCTLASS_HOST_DEVICE + ConstSubbyteReference(): ptr_(nullptr), offset_(0) { } + + /// Constructor + MCTLASS_HOST_DEVICE + ConstSubbyteReference( + Element const *ptr, /// pointer to memory + int64_t offset /// logical offset in units of Element + ): + ptr_(reinterpret_cast(ptr)), + offset_(0) { + + int64_t offset_in_vectors = offset / kElementsPerVector; + int64_t offset_in_elements = offset % kElementsPerVector; + + ptr_ += offset_in_vectors; + offset_ = int(offset_in_elements); + } + + /// Constructor + MCTLASS_HOST_DEVICE + ConstSubbyteReference( + Element *ptr = nullptr + ): ConstSubbyteReference(ptr, 0) { } + + /// Gets storage pointer + MCTLASS_HOST_DEVICE + StoragePointer storage_pointer() const { + return ptr_; + } + + /// Gets element offset within storage vector + MCTLASS_HOST_DEVICE + int element_offset() const { + return offset_; + } + + /// Unpacks an element from memory + MCTLASS_HOST_DEVICE + Element get() const { + Storage item = Storage((*ptr_ >> (offset_ * sizeof_bits::value)) & kMask); + return reinterpret_cast(item); + } + + /// Unpacks an element from memory + MCTLASS_HOST_DEVICE + operator Element() const { + return get(); + } + + /// Adds an offset in units of elements to the reference + MCTLASS_HOST_DEVICE + ConstSubbyteReference &operator+=(int offset) { + + offset += offset_; + + int offset_in_vectors = offset / kElementsPerVector; + int offset_in_elements = offset % kElementsPerVector; + + ptr_ += offset_in_vectors; + offset_ = offset_in_elements; + + return *this; + } + + /// Adds an offset in units of elements to the reference + MCTLASS_HOST_DEVICE + ConstSubbyteReference &operator+=(long long offset) { + + offset += offset_; + + long long offset_in_vectors = offset / kElementsPerVector; + int offset_in_elements = int(offset % kElementsPerVector); + + ptr_ += offset_in_vectors; + offset_ = offset_in_elements; + + return *this; + } + + /// Adds an offset in units of elements to the reference + MCTLASS_HOST_DEVICE + ConstSubbyteReference &operator-=(int offset) { + + int offset_in_vectors = offset / kElementsPerVector; + int offset_in_elements = offset % kElementsPerVector; + + ptr_ -= offset_in_vectors; + offset_ -= offset_in_elements; + + if (offset_ < 0) { + offset_ += kElementsPerVector; + --ptr_; + } + + return *this; + } + + /// Adds an offset in units of elements to the reference + MCTLASS_HOST_DEVICE + ConstSubbyteReference &operator-=(long long offset) { + + long long offset_in_vectors = offset / kElementsPerVector; + int offset_in_elements = int(offset % kElementsPerVector); + + ptr_ -= offset_in_vectors; + offset_ -= offset_in_elements; + + if (offset_ < 0) { + offset_ += kElementsPerVector; + --ptr_; + } + + return *this; + } + + /// Returns a reference to an element with a given offset from the current reference + MCTLASS_HOST_DEVICE + ConstSubbyteReference operator+(int offset) const { + + ConstSubbyteReference ref(ptr_, offset_); + ref += offset; + + return ref; + } + + /// Returns a reference to an element with a given offset from the current reference + MCTLASS_HOST_DEVICE + ConstSubbyteReference operator+(long long offset) const { + + ConstSubbyteReference ref(ptr_, offset_); + ref += offset; + + return ref; + } + + /// Returns a reference to an element with a given offset from the current reference + MCTLASS_HOST_DEVICE + ConstSubbyteReference operator-(int offset) const { + + ConstSubbyteReference ref(ptr_, offset_); + ref -= offset; + + return ref; + } + + /// Returns a reference to an element with a given offset from the current reference + MCTLASS_HOST_DEVICE + ConstSubbyteReference operator-=(long long offset) const { + + ConstSubbyteReference ref(ptr_, offset_); + ref -= offset; + + return ref; + } + + /// Computes the difference in elements between references + MCTLASS_HOST_DEVICE + ptrdiff_t operator-(ConstSubbyteReference ref) const { + return (ptr_ - ref.ptr_) * kElementsPerVector + (offset_ - ref.offset_); + } + + /// Explicit cast to int + MCTLASS_HOST_DEVICE + explicit operator int() const { + return int(get()); + } + + /// Explicit cast to signed 64-bit integer + MCTLASS_HOST_DEVICE + explicit operator int64_t() const { + return int64_t(get()); + } + + /// Explicit cast to unsigned 64-bit integer + MCTLASS_HOST_DEVICE + explicit operator uint64_t() const { + return uint64_t(get()); + } + + /// Explicit cast to float + MCTLASS_HOST_DEVICE + explicit operator float() const { + return float(get()); + } + + /// Explicit cast to double + MCTLASS_HOST_DEVICE + explicit operator double() const { + return double(get()); + } +}; + +template < + typename Element_, /// MCTLASS numeric element type. + typename Storage_ = /// Underlying storage type. Must be able to hold an integer + /// number of objects of type Element. + +#if defined(__MACA_ARCH__) /// Default size depends on width of atomicCas() overloads. + uint16_t +#else + uint8_t +#endif +> +class SubbyteReference { +public: + + using Element = Element_; + using Storage = Storage_; + using StoragePointer = Storage *; + + static_assert(sizeof_bits::value <= sizeof_bits::value, + "Size of Element must not be greater than Storage."); + + static_assert(!(sizeof_bits::value % sizeof_bits::value), + "Storage must be divisible by Element"); + +private: + + ///! Number of elements per storage vector + int const kElementsPerVector = sizeof_bits::value / sizeof_bits::value; + + ///! Bit mask + Storage const kMask = + ((sizeof_bits::value < sizeof_bits::value) ? + (Storage(1) << sizeof_bits::value) - Storage(1) : + ~Storage(0)); + +private: + + /// Pointer to array containing element + StoragePointer ptr_; + + /// Offset (in units of elements) from pointer. + /// + /// Invariant: must always be in range [0, kElementsPerVector) + int offset_; + +public: + + MCTLASS_HOST_DEVICE + SubbyteReference(): ptr_(nullptr), offset_(0) { } + + /// Constructor + MCTLASS_HOST_DEVICE + SubbyteReference( + Element *ptr, /// pointer to memory + int64_t offset /// logical offset in units of Element + ): + ptr_(reinterpret_cast(ptr)), + offset_(0) { + + int64_t offset_in_vectors = offset / kElementsPerVector; + int64_t offset_in_elements = offset % kElementsPerVector; + + ptr_ += offset_in_vectors; + offset_ = int(offset_in_elements); + } + + /// Constructor + MCTLASS_HOST_DEVICE + SubbyteReference( + Element *ptr = nullptr + ): SubbyteReference(ptr, 0) { } + + /// Gets storage pointer + MCTLASS_HOST_DEVICE + StoragePointer storage_pointer() const { + return ptr_; + } + + /// Gets storage pointer + MCTLASS_HOST_DEVICE + Element * operator&() const { + return reinterpret_cast(ptr_); + } + + /// Gets element offset within storage vector + MCTLASS_HOST_DEVICE + int element_offset() const { + return offset_; + } + + /// Unpacks an element from memory + MCTLASS_HOST_DEVICE + Element get() const { + Storage item = Storage((*ptr_ >> (offset_ * sizeof_bits::value)) & kMask); + return reinterpret_cast(item); + } + + /// Stores an element to memory + MCTLASS_HOST_DEVICE + SubbyteReference & set(Element const &x) { + + Storage item = (reinterpret_cast(x) & kMask); + Storage kUpdateMask = Storage(~(kMask << (offset_ * mctlass::sizeof_bits::value))); + Storage new_bits = Storage(item << (offset_ * mctlass::sizeof_bits::value)); + +#if defined(__MACA_ARCH__) + + // + // Homebrew read-modify-write + // + Storage original; + Storage updated; + + do { + + original = (*ptr_); + + updated = Storage((original & kUpdateMask) | new_bits); + + original = atomicCAS(ptr_, original, updated); + + } while (updated != original); + +#else + + Storage original = (*ptr_); + Storage updated = Storage((original & kUpdateMask) | new_bits); + *ptr_ = updated; + +#endif + + return *this; + } + + //// + + /// Unpacks an element from memory + MCTLASS_HOST_DEVICE + operator Element() const { + return get(); + } + + /// Stores an element to memory + MCTLASS_HOST_DEVICE + SubbyteReference &operator=(Element const & x) { + return set(x); + } + + /// Stores an element to memory + MCTLASS_HOST_DEVICE + SubbyteReference &operator=(SubbyteReference const & x) { + return set(x.get()); + } + + /// Stores an element to memory + MCTLASS_HOST_DEVICE + SubbyteReference &operator=( + ConstSubbyteReference const &x) { + return set(x.get()); + } + + /// Adds an offset in units of elements to the reference + MCTLASS_HOST_DEVICE + SubbyteReference &operator+=(int offset) { + + offset += offset_; + + int offset_in_vectors = offset / kElementsPerVector; + int offset_in_elements = offset % kElementsPerVector; + + ptr_ += offset_in_vectors; + offset_ = offset_in_elements; + + return *this; + } + + /// Adds an offset in units of elements to the reference + MCTLASS_HOST_DEVICE + SubbyteReference &operator+=(long long offset) { + + offset += offset_; + + long long offset_in_vectors = offset / kElementsPerVector; + int offset_in_elements = int(offset % kElementsPerVector); + + ptr_ += offset_in_vectors; + offset_ = offset_in_elements; + + return *this; + } + + /// Adds an offset in units of elements to the reference + MCTLASS_HOST_DEVICE + SubbyteReference &operator-=(int offset) { + + int offset_in_vectors = offset / kElementsPerVector; + int offset_in_elements = offset % kElementsPerVector; + + ptr_ -= offset_in_vectors; + offset_ -= offset_in_elements; + + if (offset_ < 0) { + offset_ += kElementsPerVector; + --ptr_; + } + + return *this; + } + + /// Adds an offset in units of elements to the reference + MCTLASS_HOST_DEVICE + SubbyteReference &operator-=(long long offset) { + + long long offset_in_vectors = offset / kElementsPerVector; + int offset_in_elements = int(offset % kElementsPerVector); + + ptr_ -= offset_in_vectors; + offset_ -= offset_in_elements; + + if (offset_ < 0) { + offset_ += kElementsPerVector; + --ptr_; + } + + return *this; + } + + /// Returns a reference to an element with a given offset from the current reference + MCTLASS_HOST_DEVICE + SubbyteReference operator+(int offset) const { + + SubbyteReference ref(ptr_, offset_); + ref += offset; + + return ref; + } + + /// Returns a reference to an element with a given offset from the current reference + MCTLASS_HOST_DEVICE + SubbyteReference operator+(long long offset) const { + + SubbyteReference ref(ptr_, offset_); + ref += offset; + + return ref; + } + + /// Returns a reference to an element with a given offset from the current reference + MCTLASS_HOST_DEVICE + SubbyteReference operator-(int offset) const { + + SubbyteReference ref(ptr_, offset_); + ref -= offset; + + return ref; + } + + /// Returns a reference to an element with a given offset from the current reference + MCTLASS_HOST_DEVICE + SubbyteReference operator-=(long long offset) const { + + SubbyteReference ref(ptr_, offset_); + ref -= offset; + + return ref; + } + + /// Computes the difference in elements between references + MCTLASS_HOST_DEVICE + ptrdiff_t operator-(SubbyteReference ref) const { + return (ptr_ - ref.ptr_) * kElementsPerVector + (offset_ - ref.offset_); + } + + /// Explicit cast to int + MCTLASS_HOST_DEVICE + explicit operator int() const { + return int(get()); + } + + /// Explicit cast to signed 64-bit integer + MCTLASS_HOST_DEVICE + explicit operator int64_t() const { + return int64_t(get()); + } + + /// Explicit cast to unsigned 64-bit integer + MCTLASS_HOST_DEVICE + explicit operator uint64_t() const { + return uint64_t(get()); + } + + /// Explicit cast to float + MCTLASS_HOST_DEVICE + explicit operator float() const { + return float(get()); + } + + /// Explicit cast to double + MCTLASS_HOST_DEVICE + explicit operator double() const { + return double(get()); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template ::value < 8)> +struct ReferenceFactory; + +template +struct ReferenceFactory { + MCTLASS_HOST_DEVICE + static Element &get(Element *ptr, int64_t offset) { + return ptr[offset]; + } + + MCTLASS_HOST_DEVICE + static Element const &get(Element const *ptr, int64_t offset) { + return ptr[offset]; + } +}; + +template +struct ReferenceFactory { + MCTLASS_HOST_DEVICE + static SubbyteReference get(Element *ptr, int64_t offset) { + return SubbyteReference(ptr, offset); + } + + MCTLASS_HOST_DEVICE + static ConstSubbyteReference get(Element const *ptr, + int64_t offset) { + return ConstSubbyteReference(ptr, offset); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/tensor_coord.h b/csrc/mctlass/include/mctlass/tensor_coord.h new file mode 100644 index 0000000..2c1dbad --- /dev/null +++ b/csrc/mctlass/include/mctlass/tensor_coord.h @@ -0,0 +1,326 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines a canonical coordinate for rank=4 tensors offering named indices. +*/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/coord.h" + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a canonical 4D coordinate used by tensor operations. +struct Tensor4DCoord : public Coord<4> { + + /// Base class + using Base = Coord<4>; + + /// Index type + using Index = typename Base::Index; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// Batch dimension + static int const kN = 0; + + /// Height dimension + static int const kH = 1; + + /// Width dimension + static int const kW = 2; + + /// Channels dimension + static int const kC = 3; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Tensor4DCoord() { } + + /// Constructs from Coord<4> + MCTLASS_HOST_DEVICE + Tensor4DCoord(Coord<4> const &coord): Base(coord) { } + + /// Helper to construct from N, H, W, and C. + MCTLASS_HOST_DEVICE + Tensor4DCoord(Index n, Index h, Index w, Index c): Base(make_Coord(n, h, w, c)) { } + + /// Helper to construct from N, H, W, and C, which are LongIndex type + MCTLASS_HOST_DEVICE + Tensor4DCoord(LongIndex n, LongIndex h, LongIndex w, LongIndex c) + : Base(make_Coord(Index(n), Index(h), Index(w), Index(c))) { } + + /// Returns the batch of the coordinate + MCTLASS_HOST_DEVICE + Index const & n() const { return this->at(kN); } + + /// Returns the batch of the coordinate + MCTLASS_HOST_DEVICE + Index & n() { return this->at(kN); } + + /// Returns the row of the coordinate + MCTLASS_HOST_DEVICE + Index const & h() const { return this->at(kH); } + + /// Returns the row of the coordinate + MCTLASS_HOST_DEVICE + Index & h() { return this->at(kH); } + + /// Returns the column of the coordinate + MCTLASS_HOST_DEVICE + Index const & w() const { return this->at(kW); } + + /// Returns the column of the coordinate + MCTLASS_HOST_DEVICE + Index & w() { return this->at(kW); } + + /// Returns the channel of the coordinate + MCTLASS_HOST_DEVICE + Index const & c() const { return this->at(kC); } + + /// Returns the channel of the coordinate + MCTLASS_HOST_DEVICE + Index & c() { return this->at(kC); } + + // + // Coord operators + // + + /// Element-wise addition + MCTLASS_HOST_DEVICE + Tensor4DCoord operator+(Base const& b) const { + return Tensor4DCoord(Base::operator+(b)); + } + + /// Element-wise subtraction + MCTLASS_HOST_DEVICE + Tensor4DCoord operator-(Base const& b) const { + return Tensor4DCoord(Base::operator-(b)); + } + + /// Element-wise multiplication + MCTLASS_HOST_DEVICE + Tensor4DCoord operator*(Base const& b) const { + return Tensor4DCoord(Base::operator*(b)); + } + + /// Element-wise division + MCTLASS_HOST_DEVICE + Tensor4DCoord operator/(Base const& b) const { + return Tensor4DCoord(Base::operator/(b)); + } + + /// In-place addition + MCTLASS_HOST_DEVICE + Tensor4DCoord& operator+=(Base const& b) { + Base::operator+=(b); + return *this; + } + + /// In-place subtraction + MCTLASS_HOST_DEVICE + Tensor4DCoord& operator-=(Base const& b) { + Base::operator-=(b); + return *this; + } + + /// In-place multiplication + MCTLASS_HOST_DEVICE + Tensor4DCoord& operator*=(Base const& b) { + Base::operator*=(b); + return *this; + } + + /// In-place division + MCTLASS_HOST_DEVICE + Tensor4DCoord& operator/=(Base const& b) { + Base::operator/=(b); + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Defines a canonical 5D coordinate used by tensor operations. +struct Tensor5DCoord : public Coord<5> { + + /// Base class + using Base = Coord<5>; + + /// Index type + using Index = typename Base::Index; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// Batch dimension + static int const kN = 0; + + /// Depth dimension + static int const kD = 1; + + /// Height dimension + static int const kH = 2; + + /// Width dimension + static int const kW = 3; + + /// Channels dimension + static int const kC = 4; + + // + // Methods + // + + /// Default ctor + MCTLASS_HOST_DEVICE + Tensor5DCoord() { } + + /// Constructs from Coord<5> + MCTLASS_HOST_DEVICE + Tensor5DCoord(Coord<5> const &coord): Base(coord) { } + + /// Helper to construct from N, D, H, W, and C. + MCTLASS_HOST_DEVICE + Tensor5DCoord(Index n, Index d, Index h, Index w, Index c): Base(make_Coord(n, d, h, w, c)) { } + + /// Helper to construct from N, D, H, W, and C, which are LongIndex type + MCTLASS_HOST_DEVICE + Tensor5DCoord(LongIndex n, LongIndex d, LongIndex h, LongIndex w, LongIndex c) + : Base(make_Coord(Index(n), Index(d), Index(h), Index(w), Index(c))) { } + + /// Returns the batch of the coordinate + MCTLASS_HOST_DEVICE + Index const & n() const { return this->at(kN); } + + /// Returns the batch of the coordinate + MCTLASS_HOST_DEVICE + Index & n() { return this->at(kN); } + + /// Returns the batch of the coordinate + MCTLASS_HOST_DEVICE + Index const & d() const { return this->at(kD); } + + /// Returns the batch of the coordinate + MCTLASS_HOST_DEVICE + Index & d() { return this->at(kD); } + + /// Returns the row of the coordinate + MCTLASS_HOST_DEVICE + Index const & h() const { return this->at(kH); } + + /// Returns the row of the coordinate + MCTLASS_HOST_DEVICE + Index & h() { return this->at(kH); } + + /// Returns the column of the coordinate + MCTLASS_HOST_DEVICE + Index const & w() const { return this->at(kW); } + + /// Returns the column of the coordinate + MCTLASS_HOST_DEVICE + Index & w() { return this->at(kW); } + + /// Returns the channel of the coordinate + MCTLASS_HOST_DEVICE + Index const & c() const { return this->at(kC); } + + /// Returns the channel of the coordinate + MCTLASS_HOST_DEVICE + Index & c() { return this->at(kC); } + + // + // Coord operators + // + + /// Element-wise addition + MCTLASS_HOST_DEVICE + Tensor5DCoord operator+(Base const& b) const { + return Tensor5DCoord(Base::operator+(b)); + } + + /// Element-wise subtraction + MCTLASS_HOST_DEVICE + Tensor5DCoord operator-(Base const& b) const { + return Tensor5DCoord(Base::operator-(b)); + } + + /// Element-wise multiplication + MCTLASS_HOST_DEVICE + Tensor5DCoord operator*(Base const& b) const { + return Tensor5DCoord(Base::operator*(b)); + } + + /// Element-wise division + MCTLASS_HOST_DEVICE + Tensor5DCoord operator/(Base const& b) const { + return Tensor5DCoord(Base::operator/(b)); + } + + /// In-place addition + MCTLASS_HOST_DEVICE + Tensor5DCoord& operator+=(Base const& b) { + Base::operator+=(b); + return *this; + } + + /// In-place subtraction + MCTLASS_HOST_DEVICE + Tensor5DCoord& operator-=(Base const& b) { + Base::operator-=(b); + return *this; + } + + /// In-place multiplication + MCTLASS_HOST_DEVICE + Tensor5DCoord& operator*=(Base const& b) { + Base::operator*=(b); + return *this; + } + + /// In-place division + MCTLASS_HOST_DEVICE + Tensor5DCoord& operator/=(Base const& b) { + Base::operator/=(b); + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/tensor_ref.h b/csrc/mctlass/include/mctlass/tensor_ref.h new file mode 100644 index 0000000..91dd7c9 --- /dev/null +++ b/csrc/mctlass/include/mctlass/tensor_ref.h @@ -0,0 +1,418 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines a structure containing strides, bounds, and a pointer to tensor data. +*/ +#pragma once + + +#include "mctlass/mctlass.h" +#include "mctlass/coord.h" +#include "mctlass/platform/platform.h" +#include "mctlass/subbyte_reference.h" + +namespace mctlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Default layout function from coordinates in a tensor's index space into the n-D array held +/// in memory. +/// +/// All layout functions must define at least the members shown in IdentityTensorLayout<>. +template +class IdentityTensorLayout { +public: + /// Logical rank of tensor + static int const kRank = Rank; + + /// Rank of stride vector + static int const kStrideRank = Rank; + + /// Index type used for coordinates + using Index = int32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using TensorCoord = Coord; + + /// Stride vector + using Stride = Coord; + +private: + + // + // Data members + // + + /// Stride data member + Stride stride_; + +public: + + // + // Methods + // + + MCTLASS_HOST_DEVICE + IdentityTensorLayout(Stride const &stride = Stride()): stride_(stride) { } + + /// Returns the offset of a coordinate in linear memory + MCTLASS_HOST_DEVICE + LongIndex operator()(Coord const &coord) const { + return coord.dot(stride_); + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + MCTLASS_HOST_DEVICE + Stride & stride() { + return stride_; + } + + /// Compute the number of contiguous elements needed to store a tensor with the given size + MCTLASS_HOST_DEVICE + LongIndex capacity(TensorCoord const &size) const { + int idx = stride_.max_dim_index(); + return stride_[idx] * size[idx]; + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/* \brief TensorRef is a template for objects pointing to the start of tensors of arbitrary rank + and layout within memory. A TensorRef combines a pointer and a Layout concept + + Examples: + + (These examples use helpers for matrix layouts defined in mctlass/layout/matrix.h) + + 1. Column-major matrix may be represented as a rank=2 tensor: + + TensorRef A(ptr_A, ldm); + + 2. Row-major matrix may be represented as a rank=2 tensor: + + TensorRef B(ptr_A, ldm); + + 3. An interleaved matrix may be represented as a rank=2 tensor: + + TensorRef > C; + + 4. A helper exists to define a TensorRef for a contiguous matrix whose layout + is not known at compile time. + + int ldm; // leading dimension + layout::Matrix kind; // Could be layout::Matrix::kRowMajor or layout::Matrix::kColumnMajor + + + TensorRef E(ptr_E, {ldm, kind}); + +*/ +template < + /// Data type of element stored within tensor (concept: NumericType) + typename Element_, + /// Defines a mapping from logical coordinate to linear memory (concept: Layout) + typename Layout_ +> +class TensorRef { + public: + /// Data type of individual access + using Element = Element_; + + /// Mapping function from logical coordinate to linear memory + using Layout = Layout_; + + /// Reference type to an element + using Reference = typename platform::conditional< + sizeof_bits::value >= 8, + Element &, + SubbyteReference + >::type; + + /// Logical rank of tensor index space + static int const kRank = Layout::kRank; + + /// Index type + using Index = typename Layout::Index; + + /// Long index used for pointer offsets + using LongIndex = typename Layout::LongIndex; + + /// Coordinate in logical tensor space + using TensorCoord = typename Layout::TensorCoord; + + /// Layout's stride vector + using Stride = typename Layout::Stride; + + /// TensorRef to constant data + using ConstTensorRef = TensorRef< + typename platform::remove_const::type const, + Layout>; + + /// TensorRef to non-constant data + using NonConstTensorRef = TensorRef< + typename platform::remove_const::type, + Layout>; + + /// Require at least rank=1. Mathematically, a rank=0 tensor would be considered to be a + /// scalar, but degenerate cases such as these are difficult to accommodate without + /// extensive C++ metaprogramming or support for zero-length arrays. + static_assert(kRank > 0, "Cannot define a zero-rank TensorRef"); + + private: + + /// Pointer + Element* ptr_; + + /// Layout object maps logical coordinates to linear offsets + Layout layout_; + + public: + + // + // Methods + // + + /// Constructs a TensorRef with a pointer and layout object. + MCTLASS_HOST_DEVICE + TensorRef(): ptr_(nullptr) { + + } + + /// Constructs a TensorRef with a pointer and layout object. + MCTLASS_HOST_DEVICE + TensorRef( + Element *ptr, ///< pointer to start of tensor + Layout const &layout ///< layout object containing stride and mapping function + ): + ptr_(ptr), layout_(layout) { + + } + + /// Converting constructor from TensorRef to non-constant data. + template + MCTLASS_HOST_DEVICE + TensorRef( + NonConstTensorRef const &ref, ///< TensorRef to non-const data + ///SFINAE trick to avoid creating a copy-constructor when Element_ is already non-const + _Magic magic = (typename platform::enable_if< ! platform::is_same >::value, _Magic>::type)0 + ): + ptr_(ref.data()), layout_(ref.layout()) { } + + /// Returns a reference to constant-valued tensor. + MCTLASS_HOST_DEVICE + ConstTensorRef const_ref() const { + return ConstTensorRef(ptr_, layout_); + } + + MCTLASS_HOST_DEVICE + NonConstTensorRef non_const_ref() const { + return NonConstTensorRef(const_cast::type *>(ptr_), layout_); + } + + /// Updates only the pointer + MCTLASS_HOST_DEVICE + void reset(Element* ptr = nullptr) { + ptr_ = ptr; + } + + /// Updates the pointer and layout object + MCTLASS_HOST_DEVICE + void reset(Element* ptr, Layout const &layout) { + ptr_ = ptr; + layout_ = layout; + } + + /// Returns true if the TensorRef is non-null + MCTLASS_HOST_DEVICE + bool good() const { + return ptr_ != nullptr; + } + + /// Returns the pointer to referenced data + MCTLASS_HOST_DEVICE + Element * data() const { return ptr_; } + + /// Returns a reference to the element at a given linear index + MCTLASS_HOST_DEVICE + Reference data(LongIndex idx) const { + return ReferenceFactory::type, + (sizeof_bits::value < 8)>::get(ptr_, idx); + } + + /// Returns the layout object + MCTLASS_HOST_DEVICE + Layout & layout() { + return layout_; + } + + /// Returns the layout object + MCTLASS_HOST_DEVICE + Layout layout() const { + return layout_; + } + + /// Returns the layout object's stride vector + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the layout object's stride vector + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Returns the layout object's stride in a given physical dimension + MCTLASS_HOST_DEVICE + typename Layout::Stride::Index stride(int dim) const { + return layout_.stride().at(dim); + } + + /// Returns the layout object's stride in a given physical dimension + MCTLASS_HOST_DEVICE + typename Layout::Stride::Index & stride(int dim) { + return layout_.stride().at(dim); + } + + /// Computes the offset of an index from the origin of the tensor + MCTLASS_HOST_DEVICE + LongIndex offset(TensorCoord const& coord) const { + return layout_(coord); + } + + /// Returns a reference to the element at a given Coord + MCTLASS_HOST_DEVICE + Reference at(TensorCoord const& coord) const { + return data(offset(coord)); + } + + /// Returns a reference to the element at a given Coord + MCTLASS_HOST_DEVICE + Reference operator[](TensorCoord const& coord) const { + return data(offset(coord)); + } + + /// Adds an offset to each pointer + MCTLASS_HOST_DEVICE + TensorRef & add_pointer_offset(LongIndex offset_) { + ptr_ += offset_; + return *this; + } + + /// Adds an offset to each pointer + MCTLASS_HOST_DEVICE + TensorRef & add_coord_offset(TensorCoord const &coord) { + add_pointer_offset(offset(coord)); + return *this; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorRef operator+(TensorCoord const& b) const { + TensorRef result(*this); + result.add_coord_offset(b); + return result; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorRef & operator+=(TensorCoord const& b) { + add_coord_offset(b); + return *this; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorRef operator-(TensorCoord const& b) const { + TensorRef result(*this); + result.add_pointer_offset(-offset(b)); + return result; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorRef & operator-=(TensorCoord const& b) { + add_pointer_offset(-offset(b)); + return *this; + } +}; + +/// Constructs a TensorRef, deducing types from arguments. +template < + typename Element, + typename Layout +> +MCTLASS_HOST_DEVICE +TensorRef make_TensorRef(Element *ptr, Layout const &layout) { + return TensorRef(ptr, layout); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Partial specializations to handle degenerate and sub-byte cases. +// +/////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + typename Element, + typename Layout +> +MCTLASS_HOST_DEVICE +bool TensorRef_aligned(TensorRef const &ref, int alignment) { + + int const kStrideRank = Layout::kStrideRank; + + if (reinterpret_cast(ref.data()) % alignment) { + return false; + } + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kStrideRank; ++i) { + if (ref.stride(i) % alignment) { + return false; + } + } + + return true; +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/tensor_ref_planar_complex.h b/csrc/mctlass/include/mctlass/tensor_ref_planar_complex.h new file mode 100644 index 0000000..df58c70 --- /dev/null +++ b/csrc/mctlass/include/mctlass/tensor_ref_planar_complex.h @@ -0,0 +1,374 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines a structure containing strides, bounds, and a pointer to tensor data. +*/ +#pragma once + +#include +#include "mctlass/mctlass.h" +#include "mctlass/complex.h" +#include "mctlass/tensor_ref.h" + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct PlanarComplexReference { + + // + // Type definitions + // + + using Element = Element_; + using ComplexElement = complex; + + // + // Data members + // + + Element *real; + Element *imag; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + PlanarComplexReference( + Element *real_ = nullptr, + Element *imag_ = nullptr + ): + real(real_), imag(imag_) { } + + /// Loads the complex element + MCTLASS_HOST_DEVICE + operator complex() const { + return complex{*real, *imag}; + } + + /// Stores a complex element to the location pointed to by the reference + MCTLASS_HOST_DEVICE + PlanarComplexReference &operator=(complex const &rhs) { + *real = rhs.real(); + *imag = rhs.imag(); + return *this; + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/* \brief TensorRef is a template for objects pointing to the start of tensors of arbitrary rank + and layout within memory. A TensorRef combines a pointer and a Layout concept + +*/ +template < + /// Data type of element stored within tensor (concept: NumericType) + typename Element_, + /// Defines a mapping from logical coordinate to linear memory (concept: Layout) + typename Layout_ +> +class TensorRefPlanarComplex { + public: + /// Data type of individual access + using Element = Element_; + + /// Complex element type + using ComplexElement = complex; + + /// Mapping function from logical coordinate to linear memory + using Layout = Layout_; + + static_assert(sizeof_bits::value >= 8, + "Planar complex not suitable for subbyte elements at this time"); + + /// Reference type to an element + using Reference = PlanarComplexReference; + + /// Logical rank of tensor index space + static int const kRank = Layout::kRank; + + /// Index type + using Index = typename Layout::Index; + + /// Long index used for pointer offsets + using LongIndex = typename Layout::LongIndex; + + /// Coordinate in logical tensor space + using TensorCoord = typename Layout::TensorCoord; + + /// Layout's stride vector + using Stride = typename Layout::Stride; + + /// TensorRef to constant data + using ConstTensorRef = TensorRefPlanarComplex< + typename platform::remove_const::type const, + Layout>; + + /// TensorRef to non-constant data + using NonConstTensorRef = TensorRefPlanarComplex< + typename platform::remove_const::type, + Layout>; + + /// Require at least rank=1. Mathematically, a rank=0 tensor would be considered to be a + /// scalar, but degenerate cases such as these are difficult to accommodate without + /// extensive C++ metaprogramming or support for zero-length arrays. + static_assert(kRank > 0, "Cannot define a zero-rank TensorRef"); + + private: + + /// Pointer + Element* ptr_; + + /// Layout object maps logical coordinates to linear offsets + Layout layout_; + + /// Offset to imaginary part + LongIndex imaginary_stride_; + + public: + + // + // Methods + // + + /// Constructs a TensorRef with a pointer and layout object. + MCTLASS_HOST_DEVICE + TensorRefPlanarComplex( + Element *ptr = nullptr, ///< pointer to start of tensor + Layout const &layout = Layout(), ///< layout object containing stride and mapping function + LongIndex imaginary_stride = 0 + ): + ptr_(ptr), layout_(layout), imaginary_stride_(imaginary_stride) { + + } + + /// Converting constructor from TensorRef to non-constant data. + MCTLASS_HOST_DEVICE + TensorRefPlanarComplex( + NonConstTensorRef const &ref ///< TensorRef to non-const data + ): + ptr_(ref.data()), layout_(ref.layout()), imaginary_stride_(ref.imaginary_stride_) { } + + /// Returns a reference to constant-valued tensor. + MCTLASS_HOST_DEVICE + ConstTensorRef const_ref() const { + return ConstTensorRef(ptr_, layout_, imaginary_stride_); + } + + MCTLASS_HOST_DEVICE + NonConstTensorRef non_const_ref() const { + return NonConstTensorRef( + const_cast::type *>(ptr_), + layout_, + imaginary_stride_); + } + + /// Updates only the pointer + MCTLASS_HOST_DEVICE + void reset(Element* ptr = nullptr, LongIndex imaginary_stride = 0) { + ptr_ = ptr; + imaginary_stride_ = imaginary_stride; + } + + /// Updates the pointer and layout object + MCTLASS_HOST_DEVICE + void reset(Element* ptr, Layout const &layout, LongIndex imaginary_stride) { + ptr_ = ptr; + layout_ = layout; + imaginary_stride_ = imaginary_stride; + } + + /// Returns true if the TensorRef is non-null + MCTLASS_HOST_DEVICE + bool good() const { + return ptr_ != nullptr; + } + + /// Returns the pointer to referenced data + MCTLASS_HOST_DEVICE + Element * data() const { return ptr_; } + + /// Returns the pointer to referenced data + MCTLASS_HOST_DEVICE + Element * imaginary_data() const { return ptr_ + imaginary_stride_; } + + /// Returns a reference to the element at a given linear index + MCTLASS_HOST_DEVICE + Reference data(LongIndex idx) const { + return Reference(ptr_ + idx, ptr_ + idx + imaginary_stride_); + } + + /// Returns the layout object + MCTLASS_HOST_DEVICE + Layout & layout() { + return layout_; + } + + /// Returns the layout object + MCTLASS_HOST_DEVICE + Layout layout() const { + return layout_; + } + + /// Gets the stride to an imaginary element + LongIndex imaginary_stride() const { + return imaginary_stride_; + } + + /// Gets the stride to an imaginary element + LongIndex &imaginary_stride() { + return imaginary_stride_; + } + + /// Returns the layout object's stride vector + MCTLASS_HOST_DEVICE + Stride stride() const { + return layout_.stride(); + } + + /// Returns the layout object's stride vector + MCTLASS_HOST_DEVICE + Stride & stride() { + return layout_.stride(); + } + + /// Returns the layout object's stride in a given physical dimension + MCTLASS_HOST_DEVICE + Index stride(int dim) const { + return layout_.stride().at(dim); + } + + /// Returns the layout object's stride in a given physical dimension + MCTLASS_HOST_DEVICE + Index & stride(int dim) { + return layout_.stride().at(dim); + } + + /// Computes the offset of an index from the origin of the tensor + MCTLASS_HOST_DEVICE + LongIndex offset(TensorCoord const& coord) const { + return layout_(coord); + } + + /// Returns a reference to the element at a given Coord + MCTLASS_HOST_DEVICE + Reference at(TensorCoord const& coord) const { + return data(offset(coord)); + } + + /// Returns a reference to the element at a given Coord + MCTLASS_HOST_DEVICE + Reference operator[](TensorCoord const& coord) const { + return data(offset(coord)); + } + + /// Adds an offset to each pointer + MCTLASS_HOST_DEVICE + TensorRefPlanarComplex & add_pointer_offset(LongIndex offset_) { + ptr_ += offset_; + return *this; + } + + /// Adds an offset to each pointer + MCTLASS_HOST_DEVICE + TensorRefPlanarComplex & add_coord_offset(TensorCoord const &coord) { + add_pointer_offset(offset(coord)); + return *this; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorRefPlanarComplex operator+(TensorCoord const& b) const { + TensorRefPlanarComplex result(*this); + result.add_coord_offset(b); + return result; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorRefPlanarComplex & operator+=(TensorCoord const& b) { + add_coord_offset(b); + return *this; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorRefPlanarComplex operator-(TensorCoord const& b) const { + TensorRefPlanarComplex result(*this); + result.add_pointer_offset(-offset(b)); + return result; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorRefPlanarComplex & operator-=(TensorCoord const& b) { + add_pointer_offset(-offset(b)); + return *this; + } + + /// TensorRef to real-valued tensor + MCTLASS_HOST_DEVICE + mctlass::TensorRef ref_real() const { + return mctlass::TensorRef(data(), layout()); + } + + /// TensorRef to real-valued tensor + MCTLASS_HOST_DEVICE + mctlass::TensorRef ref_imag() const { + return mctlass::TensorRef(imaginary_data(), layout()); + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Constructs a TensorRef, deducing types from arguments. +template < + typename Element, + typename Layout +> +MCTLASS_HOST_DEVICE +TensorRefPlanarComplex make_TensorRefPlanarComplex( + Element *ptr, + Layout const &layout, + int64_t imaginary_stride) { + + return TensorRefPlanarComplex(ptr, layout, imaginary_stride); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +/////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/tensor_view.h b/csrc/mctlass/include/mctlass/tensor_view.h new file mode 100644 index 0000000..8c188d3 --- /dev/null +++ b/csrc/mctlass/include/mctlass/tensor_view.h @@ -0,0 +1,297 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines a structure containing strides and a pointer to tensor data. + + TensorView is derived from TensorRef and contributes bounds to the tensor's index space. Thus, + it is a complete mathematical object and may be used in tensor algorithms. It is decoupled from + data storage and is therefore lightweight and may be embedded in larger tensor objects or + memory structures. + + See mctlass/tensor_ref.h for more details about the mapping of the logical tensor index space to + linear memory. +*/ + +#pragma once + +#if !defined(__MACACC_RTC__) +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/tensor_ref.h" + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Data type of element stored within tensor + typename Element_, + /// Maps a Coord in the logical tensor index space to the internal n-D array + typename Layout_ +> +class TensorView : public TensorRef { + public: + + /// Base tensor reference + using Base = mctlass::TensorRef; + + /// Mapping function from logical coordinate to internal n-D array + using Layout = Layout_; + + /// TensorRef pointing to constant memory + using ConstTensorRef = typename Base::ConstTensorRef; + + /// Underlying TensorRef type + using TensorRef = Base; + + /// Data type of individual access + using Element = Element_; + + /// Reference type to an element + using Reference = Element &; + + /// Logical rank of tensor index space + static int const kRank = Layout::kRank; + + /// Index type + using Index = typename Layout::Index; + + /// Long index used for pointer offsets + using LongIndex = typename Layout::LongIndex; + + /// Coordinate in logical tensor space + using TensorCoord = typename Layout::TensorCoord; + + /// Coordinate in storage n-D array + using Stride = typename Layout::Stride; + + /// TensorView pointing to constant memory + using ConstTensorView = TensorView< + typename platform::remove_const::type const, + Layout>; + + /// TensorView pointing to non-constant memory + using NonConstTensorView = TensorView< + typename platform::remove_const::type, + Layout>; + + /// Require at least rank=1. Mathematically, a rank=0 tensor would be considered to be a + /// scalar, but degenerate cases such as these are difficult to accommodate without + /// extensive C++ metaprogramming or support for zero-length arrays. + static_assert(kRank > 0, "Cannot define a zero-rank TensorRef"); + + private: + + /// View extent + TensorCoord extent_; + + public: + + // + // Methods + // + + /// Constructs a TensorView object + MCTLASS_HOST_DEVICE + TensorView() { } + + /// Constructs a TensorView object + MCTLASS_HOST_DEVICE + TensorView( + Element *ptr, ///< pointer to start of tensor + Layout const &layout, ///< layout object containing stride and mapping function + TensorCoord const &extent ///< size of the view in logical coordinates + ): + Base(ptr, layout), extent_(extent) { + + } + + /// Constructs a TensorView object + MCTLASS_HOST_DEVICE + TensorView( + TensorRef const &ref, ///< pointer and layout object referencing a tensor + TensorCoord const &extent ///< logical size of tensor + ): + Base(ref), extent_(extent) { + + } + + /// Converting constructor from TensorRef to non-constant data. + MCTLASS_HOST_DEVICE + TensorView( + NonConstTensorView const &view ///< TensorView to non-const data + ): + Base(view), extent_(view.extent_) { } + + /// Updates the pointer and layout object + MCTLASS_HOST_DEVICE + void reset(Element* ptr, Layout const &layout, TensorCoord const &extent) { + Base::reset(ptr, layout); + this->resize(extent); + } + + /// Updates the pointer + MCTLASS_HOST_DEVICE + void reset(Element* ptr) { + Base::reset(ptr); + } + + /// Changes the size of the view without affecting pointer or layout + MCTLASS_HOST_DEVICE + void resize(TensorCoord const &extent) { + this->extent_ = extent; + } + + /// Returns the extent of the view (the size along each logical dimension). + MCTLASS_HOST_DEVICE + TensorCoord const& extent() const { return extent_; } + + /// Returns the extent along a particular logical dimension. + MCTLASS_HOST_DEVICE + Index extent(int dim) const { return extent_.at(dim); } + + /// Returns the number of logical elements + MCTLASS_HOST_DEVICE + LongIndex size() const { + return extent_.product(); + } + + /// Determines whether a location is within a tensor + MCTLASS_HOST_DEVICE + bool contains(TensorCoord const& coord) const { + MCTLASS_PRAGMA_UNROLL + for (int dim = 0; dim < kRank; ++dim) { + if (!(coord[dim] >= 0 && coord[dim] < extent(dim))) { + return false; + } + } + return true; + } + + /// Returns a TensorRef pointing to the first element of the tensor. + MCTLASS_HOST_DEVICE + TensorRef ref() const { + return TensorRef(this->data(), this->layout()); + } + + /// Returns a TensorRef pointing to the first element of the tensor. + MCTLASS_HOST_DEVICE + ConstTensorRef const_ref() const { + return ConstTensorRef(this->data(), this->layout()); + } + + /// Returns a TensorView to const data + MCTLASS_HOST_DEVICE + ConstTensorView const_view() const { + return ConstTensorView(const_ref(), extent_); + } + + /// Returns a Tensor_view given location and size quantities + MCTLASS_HOST_DEVICE + TensorView subview( + TensorCoord extent, ///< extent of the resulting view + TensorCoord const& location = TensorCoord() ///< resulting view's origin within the old view + ) const { + + TensorView result(this->ref(), extent.clamp(extent_ - location)); + result.add_coord_offset(location); + return result; + } + + /// Returns the number of scalar elements needed to store tensor. + MCTLASS_HOST_DEVICE + size_t capacity() const { + return Base::layout().capacity(extent_); + } + + /// Returns a TensorView offset by a given amount + MCTLASS_HOST_DEVICE + TensorView operator+( + TensorCoord const& b ///< offset in the logical coordinate space of the tensor + ) const { + + TensorView result(*this); + result.add_pointer_offset(this->offset(b)); + return result; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorView& operator+=( + TensorCoord const& b ///< offset in the logical coordinate space of the tensor + ) { + + this->add_pointer_offset(this->offset(b)); + return *this; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorView operator-( + TensorCoord const& b ///< offset in the logical coordinate space of the tensor + ) const { + + TensorRef result(*this); + result.add_pointer_offset(-this->offset(b)); + return result; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorView& operator-=( + TensorCoord const& b ///< offset in the logical coordinate space of the tensor + ) { + + this->add_pointer_offset(-this->offset(b)); + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Constructs a TensorRef, deducing types from arguments. +template < + typename Element, + typename Layout +> +MCTLASS_HOST_DEVICE TensorView make_TensorView( + Element *ptr, + Layout const &layout, + typename Layout::TensorCoord const &extent) { + + return TensorView(ptr, layout, extent); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/tensor_view_planar_complex.h b/csrc/mctlass/include/mctlass/tensor_view_planar_complex.h new file mode 100644 index 0000000..7c46830 --- /dev/null +++ b/csrc/mctlass/include/mctlass/tensor_view_planar_complex.h @@ -0,0 +1,301 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines a structure containing strides and a pointer to tensor data. + + TensorView is derived from TensorRef and contributes bounds to the tensor's index space. Thus, + it is a complete mathematical object and may be used in tensor algorithms. It is decoupled from + data storage and is therefore lightweight and may be embedded in larger tensor objects or + memory structures. + + See mctlass/tensor_ref.h for more details about the mapping of the logical tensor index space to + linear memory. +*/ + +#pragma once + +#if !defined(__MACACC_RTC__) +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/tensor_ref_planar_complex.h" + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template < + /// Data type of element stored within tensor + typename Element_, + /// Maps a Coord in the logical tensor index space to the internal n-D array + typename Layout_ +> +class TensorViewPlanarComplex : public TensorRefPlanarComplex { + public: + + /// Base tensor reference + using Base = mctlass::TensorRefPlanarComplex; + + /// Mapping function from logical coordinate to internal n-D array + using Layout = Layout_; + + /// TensorRef pointing to constant memory + using ConstTensorRef = typename Base::ConstTensorRef; + + /// Underlying TensorRef type + using TensorRef = Base; + + /// Data type of individual access + using Element = Element_; + + /// Reference type to an element + using Reference = Element &; + + /// Logical rank of tensor index space + static int const kRank = Layout::kRank; + + /// Index type + using Index = typename Layout::Index; + + /// Long index used for pointer offsets + using LongIndex = typename Layout::LongIndex; + + /// Coordinate in logical tensor space + using TensorCoord = typename Layout::TensorCoord; + + /// Coordinate in storage n-D array + using Stride = typename Layout::Stride; + + /// TensorView pointing to constant memory + using ConstTensorView = TensorViewPlanarComplex< + typename platform::remove_const::type const, + Layout>; + + /// TensorView pointing to non-constant memory + using NonConstTensorView = TensorViewPlanarComplex< + typename platform::remove_const::type, + Layout>; + + /// Require at least rank=1. Mathematically, a rank=0 tensor would be considered to be a + /// scalar, but degenerate cases such as these are difficult to accommodate without + /// extensive C++ metaprogramming or support for zero-length arrays. + static_assert(kRank > 0, "Cannot define a zero-rank TensorRef"); + + private: + + /// View extent + TensorCoord extent_; + + public: + + // + // Methods + // + + /// Constructs a TensorView object + MCTLASS_HOST_DEVICE + TensorViewPlanarComplex(TensorCoord const &extent = TensorCoord()): extent_(extent) { + + } + + /// Constructs a TensorView object + MCTLASS_HOST_DEVICE + TensorViewPlanarComplex( + Element *ptr, ///< pointer to start of tensor + Layout const &layout, ///< layout object containing stride and mapping function + LongIndex imaginary_stride, ///< stride between real and imaginary part + TensorCoord const &extent ///< size of the view in logical coordinates + ): + Base(ptr, layout, imaginary_stride), extent_(extent) { + + } + + /// Constructs a TensorView object + MCTLASS_HOST_DEVICE + TensorViewPlanarComplex( + TensorRef const &ref, ///< pointer and layout object referencing a tensor + TensorCoord const &extent ///< logical size of tensor + ): + Base(ref), extent_(extent) { + + } + + /// Converting constructor from TensorRef to non-constant data. + MCTLASS_HOST_DEVICE + TensorViewPlanarComplex( + NonConstTensorView const &view ///< TensorView to non-const data + ): + Base(view), extent_(view.extent_) { } + + /// Updates the pointer and layout object + MCTLASS_HOST_DEVICE + void reset(Element* ptr, Layout const &layout, LongIndex imaginary_stride, TensorCoord size) { + Base::reset(ptr, layout, imaginary_stride); + this->resize(extent_); + } + + /// Changes the size of the view without affecting pointer or layout + MCTLASS_HOST_DEVICE + void resize(TensorCoord extent) { + this->extent_ = extent; + } + + /// Returns the extent of the view (the size along each logical dimension). + MCTLASS_HOST_DEVICE + TensorCoord const& extent() const { return extent_; } + + /// Returns the extent along a particular logical dimension. + MCTLASS_HOST_DEVICE + Index extent(int dim) const { return extent_.at(dim); } + + /// Determines whether a location is within a tensor + MCTLASS_HOST_DEVICE + bool contains(TensorCoord const& coord) const { + MCTLASS_PRAGMA_UNROLL + for (int dim = 0; dim < kRank; ++dim) { + if (!(coord[dim] >= 0 && coord[dim] < extent(dim))) { + return false; + } + } + return true; + } + + /// Returns a TensorRef pointing to the first element of the tensor. + MCTLASS_HOST_DEVICE + Base ref() const { + return Base(this->data(), this->layout(), this->imaginary_stride()); + } + + /// Returns a TensorRef pointing to the first element of the tensor. + MCTLASS_HOST_DEVICE + ConstTensorRef const_ref() const { + return ConstTensorRef(this->data(), this->layout()); + } + + /// Returns a TensorView to const data + MCTLASS_HOST_DEVICE + ConstTensorView const_view() const { + return ConstTensorView(const_ref(), extent_); + } + + /// Returns a Tensor_view given location and size quantities + MCTLASS_HOST_DEVICE + TensorViewPlanarComplex subview( + TensorCoord extent, ///< extent of the resulting view + TensorCoord const& location = TensorCoord() ///< resulting view's origin within the old view + ) const { + + TensorViewPlanarComplex result(this->ref(), extent.clamp(extent_ - location)); + result.add_coord_offset(location); + return result; + } + + /// Returns the number of scalar elements needed to store tensor. + MCTLASS_HOST_DEVICE + size_t capacity() const { + return Base::layout().capacity(extent_); + } + + /// Returns a TensorView offset by a given amount + MCTLASS_HOST_DEVICE + TensorViewPlanarComplex operator+( + TensorCoord const& b ///< offset in the logical coordinate space of the tensor + ) const { + + TensorViewPlanarComplex result(*this); + result.add_pointer_offset(this->offset(b)); + return result; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorViewPlanarComplex& operator+=( + TensorCoord const& b ///< offset in the logical coordinate space of the tensor + ) { + + this->add_pointer_offset(this->offset(b)); + return *this; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorViewPlanarComplex operator-( + TensorCoord const& b ///< offset in the logical coordinate space of the tensor + ) const { + + TensorRef result(*this); + result.add_pointer_offset(-this->offset(b)); + return result; + } + + /// Returns a TensorRef offset by a given amount + MCTLASS_HOST_DEVICE + TensorViewPlanarComplex& operator-=( + TensorCoord const& b ///< offset in the logical coordinate space of the tensor + ) { + + this->add_pointer_offset(-this->offset(b)); + return *this; + } + + /// TensorRef to real-valued tensor + MCTLASS_HOST_DEVICE + mctlass::TensorView view_real() const { + return mctlass::TensorView(this->data(), this->layout(), extent_); + } + + /// TensorRef to real-valued tensor + MCTLASS_HOST_DEVICE + mctlass::TensorView view_imag() const { + return mctlass::TensorView(this->imaginary_data(), this->layout(), extent_); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Constructs a TensorRef, deducing types from arguments. +template < + typename Element, + typename Layout +> +MCTLASS_HOST_DEVICE TensorViewPlanarComplex make_TensorViewPlanarComplex( + Element *ptr, + Layout const &layout, + typename Layout::LongIndex imaginary_stride, + typename Layout::TensorCoord const &extent) { + + return TensorViewPlanarComplex(ptr, layout, imaginary_stride, extent); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/tfloat32.h b/csrc/mctlass/include/mctlass/tfloat32.h new file mode 100644 index 0000000..ad865b0 --- /dev/null +++ b/csrc/mctlass/include/mctlass/tfloat32.h @@ -0,0 +1,478 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Defines a proxy class for storing Tensor Float 32 data type. +*/ +#pragma once + +#if defined(__MACACC_RTC__) +#include "mctlass/floating_point_nvrtc.h" +#else +#include +#include +#include +#endif + +#include "mctlass/mctlass.h" + +namespace mctlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tensor Float 32 data type +struct alignas(4) tfloat32_t { + + // + // Data members + // + + /// Storage type + uint32_t storage; + + // + // Methods + // + + /// Constructs from an unsigned int + MCTLASS_HOST_DEVICE + static tfloat32_t bitcast(uint32_t x) { + tfloat32_t h; + h.storage = x; + return h; + } + + /// Emulated rounding is fast in device code + MCTLASS_HOST_DEVICE + static tfloat32_t round_half_ulp_truncate(float const &s) { + uint32_t x = reinterpret_cast(s); + + #if defined(__MACA_ARCH__) + if (::isfinite(s)) { + x += 0x1000u; + } + #else + if (std::isfinite(s)) { + x += 0x1000u; + } + #endif + + return tfloat32_t::bitcast(x); + } + + /// Default constructor + MCTLASS_HOST_DEVICE + tfloat32_t() = default; + + /// Floating-point conversion - round toward nearest even + MCTLASS_HOST_DEVICE + // explicit tfloat32_t(float x): storage(round_half_ulp_truncate(x).storage) { } + tfloat32_t(float x): storage(round_half_ulp_truncate(x).storage) { } + + /// Floating-point conversion - round toward nearest even + MCTLASS_HOST_DEVICE + // explicit tfloat32_t(double x): tfloat32_t(float(x)) { + tfloat32_t(double x): tfloat32_t(float(x)) { + } + + /// Integer conversion - round toward zero + MCTLASS_HOST_DEVICE + // explicit tfloat32_t(int x) { + tfloat32_t(int x) { + float flt = static_cast(x); + #if defined(__MACA_ARCH__) + storage = reinterpret_cast(flt); + #else + std::memcpy(&storage, &flt, sizeof(storage)); + #endif + } + + /// Converts to float + MCTLASS_HOST_DEVICE + operator float() const { + + // Conversions to IEEE single-precision requires clearing dont-care bits + // of the mantissa. + unsigned bits = (storage & ~0x1fffu); + + #if defined(__MACA_ARCH__) + return reinterpret_cast(bits); + #else + float flt; + std::memcpy(&flt, &bits, sizeof(flt)); + return flt; + #endif + } + + /// Converts to float + MCTLASS_HOST_DEVICE + explicit operator double() const { + return double(float(*this)); + } + + /// Converts to int + MCTLASS_HOST_DEVICE + explicit operator int() const { + return int(float(*this)); + } + + /// Casts to bool + MCTLASS_HOST_DEVICE + explicit operator bool() const { + return (float(*this) != 0.0f); + } + + /// Obtains raw bits + MCTLASS_HOST_DEVICE + uint32_t raw() const { + return storage; + } + + /// Returns the sign bit + MCTLASS_HOST_DEVICE + bool signbit() const { + return ((raw() & 0x80000000) != 0); + } + + /// Returns the biased exponent + MCTLASS_HOST_DEVICE + int exponent_biased() const { + return int((raw() >> 23) & 0x0ff); + } + + /// Returns the unbiased exponent + MCTLASS_HOST_DEVICE + int exponent() const { + return exponent_biased() - 127; + } + + /// Returns the mantissa + MCTLASS_HOST_DEVICE + int mantissa() const { + return int(raw() & 0x7fffff); + } +}; + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +MCTLASS_HOST_DEVICE +bool signbit(mctlass::tfloat32_t const& h) { + return h.signbit(); +} + +MCTLASS_HOST_DEVICE +mctlass::tfloat32_t abs(mctlass::tfloat32_t const& h) { + return mctlass::tfloat32_t::bitcast(h.raw() & 0x7fffffff); +} + +MCTLASS_HOST_DEVICE +bool isnan(mctlass::tfloat32_t const& h) { + return (h.exponent_biased() == 0x0ff) && h.mantissa(); +} + +MCTLASS_HOST_DEVICE +bool isfinite(mctlass::tfloat32_t const& h) { + return (h.exponent_biased() != 0x0ff); +} + +MCTLASS_HOST_DEVICE +mctlass::tfloat32_t nan_tf32(const char*) { + // NVIDIA canonical NaN + return mctlass::tfloat32_t::bitcast(0x7fffffff); +} + +MCTLASS_HOST_DEVICE +bool isinf(mctlass::tfloat32_t const& h) { + return (h.exponent_biased() == 0x0ff) && !h.mantissa(); +} + +MCTLASS_HOST_DEVICE +bool isnormal(mctlass::tfloat32_t const& h) { + return h.exponent_biased() && h.exponent_biased() != 0x0ff; +} + +MCTLASS_HOST_DEVICE +int fpclassify(mctlass::tfloat32_t const& h) { + int exp = h.exponent_biased(); + int mantissa = h.mantissa(); + if (exp == 0x0ff) { + if (mantissa) { + return FP_NAN; + } + else { + return FP_INFINITE; + } + } + else if (!exp) { + if (mantissa) { + return FP_SUBNORMAL; + } + else { + return FP_ZERO; + } + } + return FP_NORMAL; +} + +MCTLASS_HOST_DEVICE +mctlass::tfloat32_t sqrt(mctlass::tfloat32_t const& h) { +#if defined(__MACACC_RTC__) + return mctlass::tfloat32_t(sqrtf(float(h))); +#else + return mctlass::tfloat32_t(std::sqrt(float(h))); +#endif +} + +MCTLASS_HOST_DEVICE +tfloat32_t copysign(tfloat32_t const& a, tfloat32_t const& b) { + + uint32_t a_mag = (reinterpret_cast(a) & 0x7fffffff); + uint32_t b_sign = (reinterpret_cast(b) & 0x80000000); + uint32_t result = (a_mag | b_sign); + + return reinterpret_cast(result); +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Standard Library operations and definitions +// +/////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace std { + +#if !defined(__MACACC_RTC__) +/// Numeric limits +template <> +struct numeric_limits { + static bool const is_specialized = true; + static bool const is_signed = true; + static bool const is_integer = false; + static bool const is_exact = false; + static bool const has_infinity = true; + static bool const has_quiet_NaN = true; + static bool const has_signaling_NaN = false; + static std::float_denorm_style const has_denorm = std::denorm_present; + static bool const has_denorm_loss = true; + static std::float_round_style const round_style = std::round_to_nearest; + static bool const is_iec559 = false; + static bool const is_bounded = true; + static bool const is_modulo = false; + static int const digits = 19; + + /// Least positive value + static mctlass::tfloat32_t min() { return mctlass::tfloat32_t::bitcast(0x01); } + + /// Minimum finite value + static mctlass::tfloat32_t lowest() { return mctlass::tfloat32_t::bitcast(0xff7fffff); } + + /// Maximum finite value + static mctlass::tfloat32_t max() { return mctlass::tfloat32_t::bitcast(0x7f7fffff); } + + /// Returns smallest finite value + static mctlass::tfloat32_t epsilon() { return mctlass::tfloat32_t::bitcast(0x1000); } + + /// Returns smallest finite value + static mctlass::tfloat32_t round_error() { return mctlass::tfloat32_t(0.5f); } + + /// Returns smallest finite value + static mctlass::tfloat32_t infinity() { return mctlass::tfloat32_t::bitcast(0x7f800000); } + + /// Returns smallest finite value + static mctlass::tfloat32_t quiet_NaN() { return mctlass::tfloat32_t::bitcast(0x7fffffff); } + + /// Returns smallest finite value + static mctlass::tfloat32_t signaling_NaN() { return mctlass::tfloat32_t::bitcast(0x7fffffff); } + + /// Returns smallest finite value + static mctlass::tfloat32_t denorm_min() { return mctlass::tfloat32_t::bitcast(0x1); } +}; +#endif + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace std + +/////////////////////////////////////////////////////////////////////////////////////////////////// +// +// Arithmetic operators +// +/////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +MCTLASS_HOST_DEVICE +bool operator==(tfloat32_t const& lhs, tfloat32_t const& rhs) { + return float(lhs) == float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator!=(tfloat32_t const& lhs, tfloat32_t const& rhs) { + return float(lhs) != float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator<(tfloat32_t const& lhs, tfloat32_t const& rhs) { + return float(lhs) < float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator<=(tfloat32_t const& lhs, tfloat32_t const& rhs) { + return float(lhs) <= float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator>(tfloat32_t const& lhs, tfloat32_t const& rhs) { + return float(lhs) > float(rhs); +} + +MCTLASS_HOST_DEVICE +bool operator>=(tfloat32_t const& lhs, tfloat32_t const& rhs) { + return float(lhs) >= float(rhs); +} + +MCTLASS_HOST_DEVICE +tfloat32_t operator+(tfloat32_t const& lhs, tfloat32_t const& rhs) { + return tfloat32_t(float(lhs) + float(rhs)); +} + + +MCTLASS_HOST_DEVICE +tfloat32_t operator-(tfloat32_t const& lhs) { + union u_tff32 { + float val_f32; + tfloat32_t val_tf; + MCTLASS_HOST_DEVICE u_tff32() : val_f32(0) { } + }; + union u_tff32 x; x.val_f32 = -reinterpret_cast(lhs); + return x.val_tf; +} + +MCTLASS_HOST_DEVICE +tfloat32_t operator-(tfloat32_t const& lhs, tfloat32_t const& rhs) { + return tfloat32_t(float(lhs) - float(rhs)); +} + +MCTLASS_HOST_DEVICE +tfloat32_t operator*(tfloat32_t const& lhs, tfloat32_t const& rhs) { + return tfloat32_t(float(lhs) * float(rhs)); +} + +MCTLASS_HOST_DEVICE +tfloat32_t operator/(tfloat32_t const& lhs, tfloat32_t const& rhs) { + return tfloat32_t(float(lhs) / float(rhs)); +} + +MCTLASS_HOST_DEVICE +tfloat32_t& operator+=(tfloat32_t & lhs, tfloat32_t const& rhs) { + lhs = tfloat32_t(float(lhs) + float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +tfloat32_t& operator-=(tfloat32_t & lhs, tfloat32_t const& rhs) { + lhs = tfloat32_t(float(lhs) - float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +tfloat32_t& operator*=(tfloat32_t & lhs, tfloat32_t const& rhs) { + lhs = tfloat32_t(float(lhs) * float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +tfloat32_t& operator/=(tfloat32_t & lhs, tfloat32_t const& rhs) { + lhs = tfloat32_t(float(lhs) / float(rhs)); + return lhs; +} + +MCTLASS_HOST_DEVICE +tfloat32_t& operator++(tfloat32_t & lhs) { + float tmp(lhs); + ++tmp; + lhs = tfloat32_t(tmp); + return lhs; +} + +MCTLASS_HOST_DEVICE +tfloat32_t& operator--(tfloat32_t & lhs) { + float tmp(lhs); + --tmp; + lhs = tfloat32_t(tmp); + return lhs; +} + +MCTLASS_HOST_DEVICE +tfloat32_t operator++(tfloat32_t & lhs, int) { + tfloat32_t ret(lhs); + float tmp(lhs); + tmp++; + lhs = tfloat32_t(tmp); + return ret; +} + +MCTLASS_HOST_DEVICE +tfloat32_t operator--(tfloat32_t & lhs, int) { + tfloat32_t ret(lhs); + float tmp(lhs); + tmp--; + lhs = tfloat32_t(tmp); + return ret; +} + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// +// User-defined literals +// + +MCTLASS_HOST_DEVICE +mctlass::tfloat32_t operator "" _tf32(long double x) { + return mctlass::tfloat32_t(float(x)); +} + +MCTLASS_HOST_DEVICE +mctlass::tfloat32_t operator "" _tf32(unsigned long long int x) { + return mctlass::tfloat32_t(int(x)); +} + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/thread/matrix.h b/csrc/mctlass/include/mctlass/thread/matrix.h new file mode 100644 index 0000000..c38ae13 --- /dev/null +++ b/csrc/mctlass/include/mctlass/thread/matrix.h @@ -0,0 +1,198 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Defines a matrix object intended for storing data in registers and operations within + a CUDA thread. +*/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/matrix_coord.h" + +namespace mctlass { +namespace thread { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Per-thread matrix object storing a packed matrix +template < + typename Element, + int Rows, + int Columns, + typename Layout = layout::RowMajor +> +class Matrix : public Array { +public: + + // Verify layout refers to a rank=2 matrix. + static_assert( + Layout::kRank == 2, + "Layout type must refer to a rank=2 matrix"); + + /// Base type + using Base = Array; + + /// Element type + using Element = Element_; + + /// Number of rows + static int const kRows = Rows; + + /// Number of columns + static int const kColumns = Columns; + + /// Layout within the array + using Layout = Layout_; + + /// Reference type to an element + using Reference = Element &; + + /// Logical rank of tensor index space + static int const kRank = 2; + + /// Index type + using Index = typename Layout::Index; + + /// Long index used for pointer offsets + using LongIndex = typename Layout::LongIndex; + + /// Coordinate in logical tensor space + using TensorCoord = typename Layout::TensorCoord; + + /// Stride type + using Stride = typename Layout::Stride; + + /// TensorRef to matrix object + using TensorRef = TensorRef; + + /// TensorRef to constant matrix object + using ConstTensorRef = typename TensorRef::ConstTensorRef; + + /// TensorRef to matrix object + using TensorView = TensorView; + + /// TensorRef to constant matrix object + using ConstTensorView = typename TensorView::ConstTensorView; + + /// Diagonal vector + using Diagonal = Vector; + +private: + + +public: + + // + // Methods + // + + /// Returns the size of the object + MCTLASS_HOST_DEVICE + static MatrixCoord extent() { + return make_Coord(kRows, kColumns); + } + + /// Returns the layout object + MCTLASS_HOST_DEVICE + static Layout layout() { + return Layout::packed(extent()); + } + + /// Ctor + MCTLASS_HOST_DEVICE + Matrix() { } + + /// Ctor + MCTLASS_HOST_DEVICE + Matrix(Diagonal const &diag) { + } + + /// Returns a TensorRef pointing to the first element of the tensor. + MCTLASS_HOST_DEVICE + TensorRef ref() { + return TensorRef(this->data(), layout()); + } + + /// Returns a TensorRef pointing to the first element of the tensor. + MCTLASS_HOST_DEVICE + ConstTensorRef const_ref() const { + return ConstTensorRef(this->data(), layout()); + } + + /// Returns a TensorRef pointing to the first element of the tensor. + MCTLASS_HOST_DEVICE + TensorView view() { + return TensorView(ref(), extent()); + } + + /// Returns a TensorView to const data + MCTLASS_HOST_DEVICE + ConstTensorView const_view() const { + return ConstTensorView(const_ref(), extent()); + } + + /// Returns a reference to the element at a given Coord + MCTLASS_HOST_DEVICE + Reference at(MatrixCoord const& coord) const { + typename Base::size_type offset_(layout().offset(coord)); + return Base::at(offset_); + } + + /// Returns the number of scalar elements needed to store tensor. + MCTLASS_HOST_DEVICE + LongIndex capacity() const { + return LongIndex(Base::size()); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Column vector defined as a matrix with exactly one column +template < + typename Element, + int Rows, + typename Layout = layout::ColumnMajor +> +using ColumnVector = Matrix; + +/// Row vector defined as a matrix with exactly one row +template < + typename Element, + int Columns, + typename Layout = layout::RowMajor +> +using RowVector = Matrix; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace thread +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/trace.h b/csrc/mctlass/include/mctlass/trace.h new file mode 100644 index 0000000..e857686 --- /dev/null +++ b/csrc/mctlass/include/mctlass/trace.h @@ -0,0 +1,59 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Helpers for optionally tracing through code when debugging. + + This file is to be included after all other headers. +*/ + +#pragma once + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Tracing options +#ifndef MCTLASS_DEBUG_TRACE_LEVEL +#define MCTLASS_DEBUG_TRACE_LEVEL 0 +#endif + +#if MCTLASS_DEBUG_TRACE_LEVEL +#include +#include "mctlass/core_io.h" +#if defined(__MACA_ARCH__) +#define MCTLASS_TRACE_HOST(x) +#else +#define MCTLASS_TRACE_HOST(x) { std::cout << __FILE__ << ":" << __LINE__ << " " << x << std::endl; } +#endif +#else +#define MCTLASS_TRACE_HOST(x) +#endif + +//////////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/transform/collective/sm90_wgmma_transpose.hpp b/csrc/mctlass/include/mctlass/transform/collective/sm90_wgmma_transpose.hpp new file mode 100644 index 0000000..de169e8 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/collective/sm90_wgmma_transpose.hpp @@ -0,0 +1,336 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing how threads are mapped to a given tile. +*/ + +#pragma once + +#include "cute/arch/mma_sm90_gmma.hpp" +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace collective { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace detail { +using namespace cute; + +template +constexpr auto +gmma_smem_transpose_or_passthrough() { + if constexpr (Transpose) { + if constexpr (cute::is_same_v, SmemLayoutAtom>) { + return GMMA::Layout_K_SW128_Atom{}; + } + else if constexpr (cute::is_same_v, SmemLayoutAtom>) { + return GMMA::Layout_K_SW64_Atom{}; + } + else if constexpr (cute::is_same_v, SmemLayoutAtom>) { + return GMMA::Layout_K_SW32_Atom{}; + } + else if constexpr (cute::is_same_v, SmemLayoutAtom>) { + return GMMA::Layout_K_INTER_Atom{}; + } + else { + static_assert(mctlass::detail::dependent_false, "Unsupported Layout_SW_Atom for B SMEM transposition"); + } + } + else { + return SmemLayoutAtom{}; + } +} + +template +constexpr auto +use_universal_transposition() { + if constexpr (sizeof(ElementType) == 1) { + return !cute::is_same_v, SmemCopyAtom>; + } + else if constexpr (sizeof(ElementType) == 4){ + // Only universal transposition can handle SW64 and Non swizzle SMEM layout + if constexpr (cute::is_same_v, SmemCopyAtom> || + cute::is_same_v, SmemCopyAtom>) { + return true; + } + else { + return false; + } + } + else { + static_assert(mctlass::detail::dependent_false, "Unsupported ElementType for B SMEM transposition"); + } +} + +/// Transpose B operand in SMEM +template < + class TensorSmemB, + class TensorTransposedSmemB, + class PipelineState, + class TiledMma, + class SmemLayoutB, + class SmemLayoutAtomB, + class ElementB> +MCTLASS_DEVICE void +transpose_b_operand ( + TensorSmemB const& sB, + TensorTransposedSmemB const& gmma_sB, + PipelineState const& smem_pipe_read, + int warp_idx, int warp_group_thread_idx, + TiledMma, SmemLayoutB, SmemLayoutAtomB, ElementB) +{ + ////////////////////////////////////////////////////////////////////////////////////////////////////////////// + /// Important terms: + /// WarpgroupTileSize : The warp_group_tile size (WarpgroupTileSize x WarpgroupTileSize) a warp group would transpose + /// WarpTileSize : The warp_tile size (WarpTile x WarpTile) a warp would transpose + /// Step : The number of steps a warp group takes to complete the entire warp_group_tile transposition. + /// WarpTileNCoordLUT : The look up table to store the n-dim coords used by the warps + /// WarpTileKCoordLUT : The look up table to store the k-dim coords used by the warps + ////////////////////////////////////////////////////////////////////////////////////////////////////////////// + static_assert(size(TiledMma{}) == NumThreadsPerWarpGroup, "Wrong thread number for TransposeB"); + constexpr int WarpgroupTileSize = size<1>(SmemLayoutB{}); // A warp group tile would process entire Smem K. + constexpr int NumWarpsPerWarpGroup = NumThreadsPerWarpGroup / NumThreadsPerWarp; + + constexpr int BytesPerSmemSwizzleUnit = 16; + constexpr int WarpThreadShapeN = BytesPerSmemSwizzleUnit / sizeof(ElementB); + constexpr int WarpThreadShapeK = NumThreadsPerWarp / WarpThreadShapeN; + + ////////////////////////////////////////////////////////////////////////////////////////////////////////////// + /// Optimized transposition, less regs per thread than universal approach, need warp sync between load and store + /// TF32/FP32 would use the 2-steps approach. Fp8/Int8 would use 8-steps approach. + ////////////////////////////////////////////////////////////////////////////////////////////////////////////// + if constexpr (!detail::use_universal_transposition()) { + constexpr int Steps = sizeof(ElementB) == 1 ? 8 : 2; + constexpr int NumWarpTilePerWarpgroupTile = NumWarpsPerWarpGroup * (Steps == 8 ? 2 : 1); + + constexpr int WarpTileSize = WarpgroupTileSize / NumWarpTilePerWarpgroupTile; + static_assert(WarpTileSize >= WarpThreadShapeN && WarpTileSize >= WarpThreadShapeK, "Invaild warp thread shape." ); + constexpr auto WarpThreadLayout = make_layout(make_shape(Int{}, Int{})); + constexpr int TilesPerWarp = 2; // Each Warp would process 2 warp_tiles in one step. + ////////////////////////////////////////////////////////////////////////////////////////////////////////////// + /// A warp group uses 2 or 8 steps to transpose the whole WarpgroupTileSize x WarpgroupTileSize. + /// In each step, one warp would hold two warp_tiles. + /// Step 0: Step 1: + /// W0 W1 W2 W3 -- -- -- -- + /// W1 W0 -- -- -- -- W3 W2 + /// W2 -- -- -- -- W3 W0 W1 + /// W3 -- -- -- -- W2 W1 W1 + /// OR: + /// Divide a warp_group_tile into 8x8 warp_tiles to futher reduce the reg usage. + /// Step 0: Step 1: Step 2: Step 3: + /// W0 W1 W2 W3 -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- + /// W1 W0 -- -- -- -- -- -- -- -- W3 W2 -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- + /// W2 -- -- -- -- -- -- -- -- W3 W0 W1 -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- + /// W3 -- -- -- -- -- -- -- -- W2 W1 W0 -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- + /// -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- W0 W1 W2 W3 -- -- -- -- -- -- -- -- + /// -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- W1 W0 -- -- -- -- -- -- -- -- W3 W2 + /// -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- W2 -- -- -- -- -- -- -- -- W3 W0 W1 + /// -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- W3 -- -- -- -- -- -- -- -- W2 W1 W0 + /// + /// Step 4: Step 5: Step 6: Step 7: + /// -- -- -- -- W0 W1 W2 W3 -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- + /// -- -- -- -- -- -- -- -- -- -- -- -- W0 W1 W2 W3 -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- + /// -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- W0 W1 W2 W3 -- -- -- -- -- -- -- -- + /// -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- -- W0 W1 W2 W3 + /// W0 -- -- -- -- -- -- -- -- W0 -- -- -- -- -- -- -- -- W0 -- -- -- -- -- -- -- -- W0 -- -- -- -- + /// W1 -- -- -- -- -- -- -- -- W1 -- -- -- -- -- -- -- -- W1 -- -- -- -- -- -- -- -- W1 -- -- -- -- + /// W2 -- -- -- -- -- -- -- -- W2 -- -- -- -- -- -- -- -- W2 -- -- -- -- -- -- -- -- W2 -- -- -- -- + /// W3 -- -- -- -- -- -- -- -- W3 -- -- -- -- -- -- -- -- W3 -- -- -- -- -- -- -- -- W3 -- -- -- -- + /// + ///////////////////////////////////////////////////////////////////////////////////////////////////////////// + /// + /// Fully static coord LUT to avoid extra register use. + /// [warp_id][step][warp_tile][n / k] + /// Step 0 Step 1 Step 2 Step 3 Step 4 Step 5 Step 6 Step 7 + /// {{{0,0}, {1,1}}, {{2,2}, {3,3}}, {{4,4}, {5,5}}, {{6,6}, {7,7}}, {{4,0}, {0,4}}, {{4,1}, {1,4}}, {{4,2}, {2,4}}, {{4,3}, {3,4}}}, // W0 + /// {{{1,0}, {0,1}}, {{3,2}, {2,3}}, {{5,4}, {4,5}}, {{7,6}, {6,7}}, {{5,0}, {0,5}}, {{5,1}, {1,5}}, {{5,2}, {2,5}}, {{5,3}, {3,5}}}, // W1 + /// {{{2,0}, {0,2}}, {{3,1}, {1,3}}, {{6,4}, {4,6}}, {{7,5}, {5,7}}, {{6,0}, {0,6}}, {{6,1}, {1,6}}, {{6,2}, {2,6}}, {{6,3}, {3,6}}}, // W2 + /// {{{3,0}, {0,3}}, {{2,1}, {1,2}}, {{7,4}, {4,7}}, {{6,5}, {5,6}}, {{7,0}, {0,7}}, {{7,1}, {1,7}}, {{7,2}, {2,7}}, {{7,3}, {3,7}}}, // W3 + /// + /// Encoding the coord of warp tile0 into two int64_t values. + /// Only encoding Step 0 ~ Step 4, since Step 5 ~ Step 7 have a straightforward pattern. + /// Only encoding warp tile0, since the coords of warp tile1 could be easily deduced from warp tile0. + /// The 2-step transposition and the 8-step transposition share the same encoding. + /// + ////////////////////////////////////////////////////////////////////////////////////////////////////////////// + constexpr int64_t WarpTileNCoordLUT = 06723763275316420; + constexpr int64_t WarpTileKCoordLUT = 05410541064206420; + constexpr int NumStepsEncoded = 4; // Only encoding first 4 steps into LUT. + constexpr int MaskPerStep = 07; // Each step is encoded into 3bits, + constexpr int NumBitsPerStep = 3; + constexpr int MaskPerWarp = 07777; // Each warp has 4 steps(12 bits) + constexpr int NumBitsPerWarp = 12; + + const int current_warp_tile_n_coord_LUT = (WarpTileNCoordLUT >> (warp_idx * NumBitsPerWarp)) & MaskPerWarp; + const int current_warp_tile_k_coord_LUT = (WarpTileKCoordLUT >> (warp_idx * NumBitsPerWarp)) & MaskPerWarp; + + // Number of warp_group_tiles + static_assert(size<0>(SmemLayoutB{}) % WarpgroupTileSize == 0, + "Copy size must evenly divide SMEM tile."); + constexpr int WarpgroupTileNum = size<0>(SmemLayoutB{}) / WarpgroupTileSize; + + // Divide entire SMEM to multiple warp_tiles + constexpr auto WarpTileShape = make_shape(Int(), Int()); + Tensor s_tile = zipped_divide( sB(_,_,smem_pipe_read.index()), WarpTileShape); + Tensor s_tile_transposed = zipped_divide(gmma_sB(_,_,smem_pipe_read.index()), WarpTileShape); + + // Get copy tile + auto sB_tiled_copy = make_tiled_copy( + Copy_Atom{}, + WarpThreadLayout, // thr_layout + Layout<_1>{} // val_layout + ); + static_assert(size(sB_tiled_copy) * NumWarpsPerWarpGroup == size(TiledMma{}), "Wrong thread number in TiledCopy."); + auto sB_thr_copy = sB_tiled_copy.get_thread_slice(warp_group_thread_idx % NumThreadsPerWarp); // slice based on lane_idx + + // Construct fragments for transposition + Tensor tmp_tCsB = sB_thr_copy.partition_S(flatten(s_tile(_, make_coord(_0{}, _0{})))); + decltype(make_fragment_like(tmp_tCsB)) transpose_fragments[TilesPerWarp] = { + make_fragment_like(tmp_tCsB), + make_fragment_like(tmp_tCsB) + }; + + MCTLASS_PRAGMA_NO_UNROLL + for (int warp_group_tile = 0; warp_group_tile < WarpgroupTileNum; ++warp_group_tile) { + int tmp_warp_tile_n_coord_LUT = current_warp_tile_n_coord_LUT; + int tmp_warp_tile_k_coord_LUT = current_warp_tile_k_coord_LUT; + + MCTLASS_PRAGMA_NO_UNROLL + for (int step = 0; step < Steps; ++step) { + // decoding the warp tile coord. + int warp_tile0_n = step < NumStepsEncoded ? (tmp_warp_tile_n_coord_LUT & MaskPerStep) : 4 + warp_idx; + int warp_tile0_k = step < NumStepsEncoded ? (tmp_warp_tile_k_coord_LUT & MaskPerStep) : step - 4; + int warp_tile1_n = warp_tile0_n == warp_tile0_k ? warp_tile0_n + 1 : warp_tile0_k; + int warp_tile1_k = warp_tile0_n == warp_tile0_k ? warp_tile0_k + 1 : warp_tile0_n; + + tmp_warp_tile_n_coord_LUT >>= NumBitsPerStep; + tmp_warp_tile_k_coord_LUT >>= NumBitsPerStep; + + // [warp_tile][n/k] + const int warp_tile_coord[TilesPerWarp][2] = { + // n k + {warp_group_tile * NumWarpTilePerWarpgroupTile + warp_tile0_n, warp_tile0_k}, // warp_tile 0 + {warp_group_tile * NumWarpTilePerWarpgroupTile + warp_tile1_n, warp_tile1_k} // warp_tile 1 + }; + + MCTLASS_PRAGMA_UNROLL + for (int warp_tile = 0; warp_tile < TilesPerWarp; ++warp_tile) { + Tensor tCsB = sB_thr_copy.partition_S( + flatten(s_tile(_, make_coord(warp_tile_coord[warp_tile][0], warp_tile_coord[warp_tile][1]))) + ); // (CPY, CPY_N, CPY_K) + copy(sB_tiled_copy, tCsB, transpose_fragments[warp_tile]); + } + + // Make sure elements in two 8x8 warp tiles are all consumed + __syncwarp(); + + MCTLASS_PRAGMA_UNROLL + for (int warp_tile = 0; warp_tile < TilesPerWarp; ++warp_tile) { + Tensor tCsB_transposed = sB_thr_copy.partition_D( + flatten(s_tile_transposed(_, make_coord(warp_tile_coord[warp_tile][0], warp_tile_coord[warp_tile][1]))) + ); // (CPY, CPY_N, CPY_K) + copy(sB_tiled_copy, transpose_fragments[warp_tile], tCsB_transposed); + } + + } // lock step + } // loop warp_group_tile + } // if not use universal transposition + + ////////////////////////////////////////////////////////////////////////////////////////////////////////////// + /// Universal transposition, need warp_group sync between load and store. + /// The number of reg used depends on the input elementB. + ////////////////////////////////////////////////////////////////////////////////////////////////////////////// + else { + /* + In one copy step, a warp group would load WarpgroupTileSize * WarpgroupTileSize tile then store to transposed location. + In warp_group_tile, each warp holds Four WarpTileSize x WarpTileSize elements: + K + ------------ + | W0 W1 W2 W3 --- + | W0 W1 W2 W3 | + | W0 W1 W2 W3 | --> Copy Step 0 + | W0 W1 W2 W3 --- + .... + | W0 W1 W2 W3 --- + | W0 W1 W2 W3 | + | W0 W1 W2 W3 | --> Copy Step n + | W0 W1 W2 W3 --- + */ + static_assert((NumThreadsPerWarpGroup % WarpThreadShapeN == 0), "Unsupported warp thread layout."); + constexpr auto WarpgroupThreadLayout = make_layout(make_shape(Int{}, Int{})); + + // Get copy tile and partition to each thread + auto sB_tiled_copy = make_tiled_copy( + Copy_Atom{}, + WarpgroupThreadLayout, // thr_layout + Layout<_1>{} // val_layout + ); + static_assert(size(sB_tiled_copy) == size(TiledMma{}), "Wrong thread number in TiledCopy."); + + auto sB_thr_copy = sB_tiled_copy.get_thread_slice(warp_group_thread_idx); + Tensor tCsB = sB_thr_copy.partition_S( sB(_,_,smem_pipe_read.index())); // (CPY, CPY_N, CPY_K) + Tensor tCsB_transposed = sB_thr_copy.partition_D(gmma_sB(_,_,smem_pipe_read.index())); // (CPY, CPY_N, CPY_K) + + // Divide partitioned tile to limit register usage + constexpr int CopySteps = size<0>(SmemLayoutB{}) / WarpgroupTileSize; + constexpr auto CopyTileShape = make_shape(size<0>(tCsB), Int< size<1>(tCsB) / CopySteps >{}, size<2>(tCsB)); + static_assert(size<1>(tCsB) % CopySteps == 0, "CopySteps must evenly divide rank 1 size of partitioned SMEM."); + + Tensor tCsB_copy_tile = zipped_divide(tCsB, CopyTileShape); + Tensor tCsB_copy_tile_transposed = zipped_divide(tCsB_transposed, CopyTileShape); + auto transpose_fragment = make_fragment_like(tCsB_copy_tile(_,_0{})); + + MCTLASS_PRAGMA_NO_UNROLL + for (int step = 0; step < CopySteps; ++step) { + copy(sB_tiled_copy, tCsB_copy_tile(_,step), transpose_fragment); + + // Make sure all elements are read before being overwritten + __syncthreads(); + + copy(sB_tiled_copy, transpose_fragment, tCsB_copy_tile_transposed(_,step)); + } + } // if use universal transposition + + // SMEM fence to make sure B is transposed before math + mctlass::arch::fence_view_async_shared(); +} + +}; // namespace detail + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace collective +} // namespace transform +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/pitch_linear_thread_map.h b/csrc/mctlass/include/mctlass/transform/pitch_linear_thread_map.h new file mode 100644 index 0000000..6cfb43f --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/pitch_linear_thread_map.h @@ -0,0 +1,931 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing how threads are mapped to a given tile. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/layout/pitch_linear.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { + +//////////////////////////////////////////////////////////////////////////////// + +/// Strip-mines a pitch-linear tile among a given number of threads, first along +/// the contiguous dimension then along the strided dimension. +/// +/// The tile must be divisible by the thread count such that all threads may +/// execute the same number of iterations with the same delta to exhaustively +/// cover the tile. +/// +/// This class satisfies the "RegularThreadMapping" concept. +/// +/// This ThreadMap is used by SIMT kernels and operand E of the sparse tensor +/// kernels. +template < + typename Shape_, + int Threads, + int ElementsPerAccess = 1 +> +struct PitchLinearStripminedThreadMap { + + /// Tensor coordinate + using TensorCoord = layout::PitchLinearCoord; + + /// Tile shape + using Shape = Shape_; + + /// Number of threads total + static int const kThreads = Threads; + + /// Extract vector length from Layout + static int const kElementsPerAccess = ElementsPerAccess; + + /// Shape of access by each thread + using ThreadAccessShape = layout::PitchLinearShape; + + /// Internal implementation details + struct Detail { + + static_assert(!(Shape::kContiguous % kElementsPerAccess), ""); + + /// Shape of the tile in units of vectors + using ShapeVec = layout::PitchLinearShape< + Shape::kContiguous / kElementsPerAccess, + Shape::kStrided + >; + + static_assert((Threads < ShapeVec::kContiguous && !(ShapeVec::kContiguous % kThreads)) || + (!(kThreads % ShapeVec::kContiguous)), + "Shape must be divisible by number of iterations of each thread."); + }; + + /// Number of iterations by each thread + using Iterations = typename platform::conditional< + Threads >= Detail::ShapeVec::kContiguous, + layout::PitchLinearShape< + 1, + // Redo the comparison here to work around divide by zero compiler + // error. The compiler evaluates both path of platform::conditional. + (Threads >= Detail::ShapeVec::kContiguous + ? (Detail::ShapeVec::kStrided + (kThreads / Detail::ShapeVec::kContiguous - 1)) / + (kThreads / Detail::ShapeVec::kContiguous) + : 0)>, + layout::PitchLinearShape>::type; + + + /// Interval between accesses along each dimension of the tensor's logical coordinate space + /// (in units of Elements) + using Delta = typename platform::conditional< + Threads >= Detail::ShapeVec::kContiguous, + layout::PitchLinearShape< + 1, + kThreads / Detail::ShapeVec::kContiguous + >, + layout::PitchLinearShape< + kThreads * kElementsPerAccess, + 1 + > + >::type; + + /// Shape of the tile in units of vectors + using StorageShape = typename platform::conditional< + Threads >= Detail::ShapeVec::kContiguous, + layout::PitchLinearShape, + layout::PitchLinearShape>::type; + + /// Maps thread ID to a coordinate offset within the tensor's logical coordinate space + /// (in units of Elements) + MCTLASS_HOST_DEVICE + static TensorCoord initial_offset(int thread_id) { + return TensorCoord( + (thread_id % Detail::ShapeVec::kContiguous) * kElementsPerAccess, + thread_id / Detail::ShapeVec::kContiguous); + } +}; + +/// This ThreadMap is used by GEMV +template < + typename Shape, + int Threads, + int ElementsPerAccess = 1 +> +struct PitchLinearTilePolicyStripminedThreadContiguous +{ + static_assert((Shape::kContiguous % (Threads * ElementsPerAccess)) == 0, + "Contiguous shape must divide number of threads"); + + using TensorCoord = layout::PitchLinearCoord; + + static int const kThreads = Threads; + static int const kElementsPerAccess = ElementsPerAccess; + + using Iterations = layout::PitchLinearShape< + Shape::kContiguous / (kThreads * kElementsPerAccess), + Shape::kStrided>; + + using Delta = layout::PitchLinearShape<1, 1>; + + MCTLASS_HOST_DEVICE + static TensorCoord initial_offset(int thread_id) + { + return TensorCoord(thread_id * Iterations::kContiguous * kElementsPerAccess, 0); + } +}; + +template < + typename Shape, + int Threads, + int ElementsPerAccess = 1 +> +struct PitchLinearTilePolicyStripminedThreadStrided +{ + static_assert((Shape::kStrided % Threads == 0), + "Strided shape must divide number of threads"); + + using TensorCoord = layout::PitchLinearCoord; + + static int const kThreads = Threads; + static int const kElementsPerAccess = ElementsPerAccess; + + using Iterations = layout::PitchLinearShape< + Shape::kContiguous / kElementsPerAccess, + Shape::kStrided / kThreads>; + + using Delta = layout::PitchLinearShape<1, 1>; + + using ShapeVec = Shape; + + MCTLASS_HOST_DEVICE + static TensorCoord initial_offset(int thread_id) + { + + return TensorCoord(0, thread_id * Iterations::kStrided); + } +}; + + +//////////////////////////////////////////////////////////////////////////////// + +/// Policy defining a warp-raked arrangement in which a shape is partitioned into contiguous +/// elements. +/// +/// This ThreadMap is used by tensor core kernels. +template < + typename Shape_, + int Threads, + typename WarpThreadArrangement_, + int ElementsPerAccess = 1 +> +struct PitchLinearWarpRakedThreadMap { + + /// Tensor coordinate + using TensorCoord = layout::PitchLinearCoord; + + /// Tile shape + using Shape = Shape_; + + /// Number of threads total + static int const kThreads = Threads; + + /// Extract vector length from Layout + static int const kElementsPerAccess = ElementsPerAccess; + + /// Shape of access by each thread + using ThreadAccessShape = layout::PitchLinearShape; + + /// Internal details made public to facilitate introspection + struct Detail { + + /// Fixed arrangement of threads within a warp (units of threads). + using WarpThreadArrangement = WarpThreadArrangement_; + + /// Number of threads per warp + static int const kWarpSize = WarpThreadArrangement::kCount; + + /// Number of participating warps + static int const kWarpCount = kThreads / kWarpSize; + + static_assert( + !(Shape::kContiguous % kElementsPerAccess), + "Shape must be divisible by vector length."); + + /// Compute the 'shape' of the overall tile in units of vectors + using ShapeInAccesses = layout::PitchLinearShape< + Shape::kContiguous / kElementsPerAccess, + Shape::kStrided + >; + + static_assert( + !(ShapeInAccesses::kContiguous % WarpThreadArrangement::kContiguous), + "ShapeInAccesses must be divisible by WarpThreadArrangement."); + + static_assert( + !(ShapeInAccesses::kStrided % WarpThreadArrangement::kStrided), + "ShapeInAccesses must be divisible by WarpThreadArrangement."); + + // compute number of warp-level accesses total + using WarpAccessIterations = layout::PitchLinearShape< + ShapeInAccesses::kContiguous / WarpThreadArrangement::kContiguous, + ShapeInAccesses::kStrided / WarpThreadArrangement::kStrided + >; + + // Divide it into the number of warps, first partitioning the strided dimension then the + // contiguous. + static int const kWarpsStrided = + (WarpAccessIterations::kStrided >= kWarpCount + ? kWarpCount + : WarpAccessIterations::kStrided); + + static int const kWarpsContiguous_0 = + (kWarpCount > WarpAccessIterations::kStrided + ? kWarpCount / kWarpsStrided + : 1); + static int const kWarpsContiguous_1 = + (kWarpCount > (2 * WarpAccessIterations::kStrided) + ? (kWarpCount / (2 * kWarpsStrided)) + : 1); + static int const kWarpsContiguous = ((WarpThreadArrangement::kContiguous == 4 && WarpThreadArrangement::kStrided == 16) + ? kWarpsContiguous_1 + : kWarpsContiguous_0); + /// Arrangement of warps within a threadblock-scoped tile + using WarpArrangement = layout::PitchLinearShape< + kWarpsContiguous, kWarpsStrided + >; + }; + + ///< Iterations along each dimension (concept: PitchLinearShape) + using Iterations = layout::PitchLinearShape< + Detail::WarpAccessIterations::kContiguous / Detail::kWarpsContiguous, + Detail::WarpAccessIterations::kStrided / Detail::kWarpsStrided + >; + + static_assert(Iterations::kCount, + "Number of iterations must be non-zero"); + + ///< Delta betweeen accesses (units of elements, concept: PitchLinearShape) + using Delta = layout::PitchLinearShape< + Detail::WarpThreadArrangement::kContiguous * kElementsPerAccess, + Detail::WarpThreadArrangement::kStrided + >; + + /// Maps thread ID to a coordinate offset within the tensor's logical coordinate space + MCTLASS_HOST_DEVICE + static TensorCoord initial_offset(int thread_id) { + + int warp_id = ((thread_id % kThreads) / Detail::kWarpSize); + int lane_id = (thread_id % Detail::kWarpSize); + + // + // compute warp-level offset + // + + // This is the shape of the entire area covered by a warp's memory access (in units of vectors) + layout::PitchLinearCoord warp_footprint{ + Detail::WarpThreadArrangement::kContiguous * Iterations::kContiguous, + Detail::WarpThreadArrangement::kStrided * Iterations::kStrided + }; + + // This is the offset of a specific warp (in units of vectors) + layout::PitchLinearCoord warp_offset{ + (warp_id % Detail::kWarpsContiguous), + (warp_id / Detail::kWarpsContiguous) + }; + + // This is the offset of a specific thread within a warp (units of vectors) + layout::PitchLinearCoord thread_offset_in_warp{ + lane_id % Detail::WarpThreadArrangement::kContiguous, + lane_id / Detail::WarpThreadArrangement::kContiguous + }; + + // This is the offset of a thread within a threadblock tile (units of vectors) + layout::PitchLinearCoord thread_offset_in_threadblock_tile_vec = + warp_footprint * warp_offset + thread_offset_in_warp; + + // This is the offset of a thread within a threadblock tile (units of elements) + layout::PitchLinearCoord thread_offset_in_threadblock_tile_base{ + thread_offset_in_threadblock_tile_vec.contiguous() * kElementsPerAccess, + thread_offset_in_threadblock_tile_vec.strided() + }; + return thread_offset_in_threadblock_tile_base; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Policy defining a warp-raked arrangement in which a shape is partitioned into contiguous +/// elements. Warps are arranged based on a stride. +/// +/// This ThreadMap is used by tensor core kernels for NCxHWx layout. +template < + typename Shape_, + int Threads, + typename WarpThreadArrangement_, + int ElementsPerAccess = 1 +> +struct PitchLinearStridedWarpRakedThreadMap { + + /// Tensor coordinate + using TensorCoord = layout::PitchLinearCoord; + + /// Tile shape + using Shape = Shape_; + + /// Number of threads total + static int const kThreads = Threads; + + using WarpThreadArrangement = WarpThreadArrangement_; + + /// Extract vector length from Layout + static int const kElementsPerAccess = ElementsPerAccess; + + /// Base ThreadMap + using BaseThreadMap = PitchLinearWarpRakedThreadMap< + Shape, + kThreads, + WarpThreadArrangement, + kElementsPerAccess + >; + + /// Shape of access by each thread + using ThreadAccessShape = typename BaseThreadMap::ThreadAccessShape; + + + struct Detail { + + using WarpThreadArrangement = WarpThreadArrangement_; + + using WarpAccessIterations = typename BaseThreadMap::Detail::WarpAccessIterations; + + static int const kWarpSize = BaseThreadMap::Detail::kWarpSize; + + static int const kWarpCount = BaseThreadMap::Detail::kWarpCount; + + using ShapeInAccesses = typename BaseThreadMap::Detail::ShapeInAccesses; + + // Divide it into the number of warps, first partitioning the contiguous dimension then the + // stride. + static int const kWarpsContiguous = + (WarpAccessIterations::kContiguous >= kWarpCount + ? kWarpCount + : WarpAccessIterations::kContiguous); + + static int const kWarpsStrided = + (kWarpCount > WarpAccessIterations::kContiguous + ? kWarpCount / kWarpsContiguous + : 1); + + /// Arrangement of warps within a threadblock-scoped tile + using WarpArrangement = layout::PitchLinearShape< + kWarpsContiguous, kWarpsStrided + >; + + }; + + ///< Iterations along each dimension (concept: PitchLinearShape) + using Iterations = layout::PitchLinearShape< + Detail::WarpAccessIterations::kContiguous / Detail::kWarpsContiguous, + Detail::WarpAccessIterations::kStrided / Detail::kWarpsStrided + >; + + static_assert(Iterations::kCount, + "Number of iterations must be non-zero"); + + ///< Delta betweeen accesses (units of elements, concept: PitchLinearShape) + using Delta = typename BaseThreadMap::Delta; + + /// Maps thread ID to a coordinate offset within the tensor's logical coordinate space + MCTLASS_HOST_DEVICE + static TensorCoord initial_offset(int thread_id) { + + int warp_id = (thread_id / Detail::kWarpSize); + int lane_id = (thread_id % Detail::kWarpSize); + + // + // compute warp-level offset + // + + // This is the shape of the entire area covered by a warp's memory access (in units of vectors) + layout::PitchLinearCoord warp_footprint{ + Detail::WarpThreadArrangement::kContiguous * Iterations::kContiguous, + Detail::WarpThreadArrangement::kStrided * Iterations::kStrided + }; + + // This is the offset of a specific warp (in units of vectors) + layout::PitchLinearCoord warp_offset{ + (warp_id % Detail::kWarpsContiguous), + (warp_id / Detail::kWarpsContiguous) + }; + + // This is the offset of a specific thread within a warp (units of vectors) + layout::PitchLinearCoord thread_offset_in_warp{ + lane_id % Detail::WarpThreadArrangement::kContiguous, + lane_id / Detail::WarpThreadArrangement::kContiguous + }; + + // This is the offset of a thread within a threadblock tile (units of vectors) + layout::PitchLinearCoord thread_offset_in_threadblock_tile_vec = + warp_footprint * warp_offset + thread_offset_in_warp; + + // This is the offset of a thread within a threadblock tile (units of elements) + layout::PitchLinearCoord thread_offset_in_threadblock_tile_base{ + thread_offset_in_threadblock_tile_vec.contiguous() * kElementsPerAccess, + thread_offset_in_threadblock_tile_vec.strided() + }; + + return thread_offset_in_threadblock_tile_base; + } + + +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Transpose the existing ThreadMap. For example, interleaved layout is like +/// congruous in the global memory and crosswise in the shared memory. We need +/// to transpose the coordinates between two. + +template +struct TransposePitchLinearThreadMap { + /// Underlying ThreadMap + using ThreadMap = ThreadMap_; + + /// Tensor coordinate + using TensorCoord = typename ThreadMap::TensorCoord; + + /// Tile shape + using Shape = typename ThreadMap::Shape; + + /// Number of threads total + static int const kThreads = ThreadMap::kThreads; + + /// Extract vector length from Layout + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + + /// Shape of access by each thread + using ThreadAccessShape = layout::PitchLinearShape; + + /// Internal details made public to facilitate introspection + struct Detail { + /// Fixed arrangement of threads within a warp (units of threads). + using WarpThreadArrangement = WarpThreadArrangement_; + + /// Number of threads per warp + static int const kWarpSize = WarpThreadArrangement::kCount; + + /// Number of participating warps + static int const kWarpCount = kThreads / kWarpSize; + + static_assert(!(Shape::kContiguous % kElementsPerAccess), + "Shape must be divisible by vector length."); + + /// Arrangement of warps within a threadblock-scoped tile + using WarpArrangement = + layout::PitchLinearShape; + }; + + ///< Iterations along each dimension (concept: PitchLinearShape) + using Iterations = + layout::PitchLinearShape; + + static_assert(Iterations::kContiguous == 1, + "Contiguous iteration has to be one to reuse the same shared store function with those that don't need transpose"); + + static_assert(Iterations::kCount, "Number of iterations must be non-zero"); + + ///< Delta betweeen accesses (units of elements, concept: PitchLinearShape) + using Delta = + layout::PitchLinearShape; + + /// Maps thread ID to a coordinate offset within the tensor's logical + /// coordinate space Note this is slightly different from the one of + /// PitchLinearWarpRakedThreadMap. + MCTLASS_HOST_DEVICE + static TensorCoord initial_offset(int thread_id) { + + int warp_id = (thread_id / Detail::kWarpSize); + int lane_id = (thread_id % Detail::kWarpSize); + + // + // compute warp-level offset + // + + // This is the shape of the entire area covered by a warp's memory access + // (in units of vectors) + layout::PitchLinearCoord warp_footprint{ + Detail::WarpThreadArrangement::kContiguous * Iterations::kContiguous, + Detail::WarpThreadArrangement::kStrided * Iterations::kStrided}; + + // This is the offset of a specific warp (in units of vectors) + // Note the order of / and %. Also the 2nd operand is kStrided. + layout::PitchLinearCoord warp_offset{ + (warp_id / Detail::WarpArrangement::kStrided), + (warp_id % Detail::WarpArrangement::kStrided)}; + + // This is the offset of a specific thread within a warp (units of vectors) + layout::PitchLinearCoord thread_offset_in_warp{ + lane_id % Detail::WarpThreadArrangement::kContiguous, + lane_id / Detail::WarpThreadArrangement::kContiguous}; + + // This is the offset of a thread within a threadblock tile (units of + // vectors) + layout::PitchLinearCoord thread_offset_in_threadblock_tile_vec = + warp_footprint * warp_offset + thread_offset_in_warp; + + // This is the offset of a thread within a threadblock tile (units of + // elements) + layout::PitchLinearCoord thread_offset_in_threadblock_tile_base{ + thread_offset_in_threadblock_tile_vec.contiguous() * kElementsPerAccess, + thread_offset_in_threadblock_tile_vec.strided()}; + + return thread_offset_in_threadblock_tile_base; + } +}; + +template +struct TransposePitchLinearThreadMapSimt { + /// Underlying ThreadMap + using ThreadMap = ThreadMap_; + + /// Tensor coordinate + using TensorCoord = typename ThreadMap::TensorCoord; + + /// Tile shape + using Shape = typename ThreadMap::Shape; + + /// Number of threads total + static int const kThreads = ThreadMap::kThreads; + + /// Extract vector length from Layout + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + + static_assert(kElementsPerAccess == 1 , "Simt transpose requires elements per access to be 1"); + ///< Iterations along each dimension (concept: PitchLinearShape) + using Iterations = + layout::PitchLinearShape; + + static_assert(Iterations::kCount, "Number of iterations must be non-zero"); + + static_assert(Iterations::kStrided == 1, + "Strided iteration has to be one to reuse the same shared store function with those that don't need transpose"); + + /// Shape of access by each thread + using ThreadAccessShape = typename ThreadMap::ThreadAccessShape; + + ///< Delta betweeen accesses (units of elements, concept: PitchLinearShape) + using Delta = + layout::PitchLinearShape; + + + /// Maps thread ID to a coordinate offset within the tensor's logical + /// coordinate space Note this is slightly different from the one of + /// PitchLinearWarpRakedThreadMap. + MCTLASS_HOST_DEVICE + static TensorCoord initial_offset(int thread_id) { + + TensorCoord coord = ThreadMap::initial_offset(thread_id); + + return TensorCoord( + coord.strided(), + coord.contiguous() + ); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + + +/// Policy defining a warp-striped arrangement. This partitions a tile into vectorized memory +/// accesses performed by each warp then distributes warps across them. Warps are striped in the +/// strided dimension and raked across the contiguous dimension. +template < + typename Shape_, /// Overall shape to partition in units of elements + int Threads, /// Number of partiticipation threads + typename WarpThreadArrangement_, /// Describes the shape of one memory access per warp + int ElementsPerAccess = 1 /// Number of elements accessed by each thread per memory operation (i.e. vector size) +> +struct PitchLinearWarpStripedThreadMap { + + /// Tensor coordinate + using TensorCoord = layout::PitchLinearCoord; + + /// Tile shape + using Shape = Shape_; + + /// Number of threads total + static int const kThreads = Threads / 2; + + /// Extract vector length from Layout + static int const kElementsPerAccess = ElementsPerAccess; + + /// Shape of access by each thread + using ThreadAccessShape = layout::PitchLinearShape; + + /// Internal details made public to facilitate introspection + struct Detail { + + /// Fixed arrangement of threads within a warp (units of threads). + using WarpThreadArrangement = WarpThreadArrangement_; + + /// Number of threads per warp + static int const kWarpSize = WarpThreadArrangement::kCount; + + /// Number of participating warps + static int const kWarpCount = kThreads / kWarpSize; + + static_assert( + !(Shape::kContiguous % kElementsPerAccess), + "Shape must be divisible by vector length."); + + /// Compute the 'shape' of the overall tile in units of vectors + using ShapeInAccesses = layout::PitchLinearShape< + Shape::kContiguous / kElementsPerAccess, + Shape::kStrided + >; + + // compute number of warp-level accesses total + using WarpAccessIterations = layout::PitchLinearShape< + ShapeInAccesses::kContiguous / WarpThreadArrangement::kContiguous, + ShapeInAccesses::kStrided / WarpThreadArrangement::kStrided + >; + + // Divide it into the number of warps, first partitioning the strided dimension then the + // contiguous. + static int const kWarpsStrided = + (WarpAccessIterations::kStrided >= kWarpCount + ? kWarpCount : (kWarpCount / WarpAccessIterations::kStrided)); + + static int const kWarpsContiguous = + (kWarpCount > WarpAccessIterations::kStrided ? + WarpAccessIterations::kContiguous / kWarpsStrided : 1); + + /// Arrangement of warps within a threadblock-scoped tile + using WarpArrangement = layout::PitchLinearShape< + kWarpsContiguous, kWarpsStrided + >; + }; + + ///< Iterations along each dimension (concept: PitchLinearShape) + using Iterations = layout::PitchLinearShape< + Detail::WarpAccessIterations::kContiguous / Detail::kWarpsContiguous, + Detail::WarpAccessIterations::kStrided / Detail::kWarpsStrided + >; + + static_assert(Iterations::kCount, + "Number of iterations must be non-zero"); + + ///< Delta betweeen accesses (units of elements, concept: PitchLinearShape) + using Delta = layout::PitchLinearShape< + Detail::WarpThreadArrangement::kContiguous * kElementsPerAccess, + Detail::WarpThreadArrangement::kStrided * Detail::WarpArrangement::kStrided + >; + + /// Maps thread ID to a coordinate offset within the tensor's logical coordinate space + MCTLASS_HOST_DEVICE + static TensorCoord initial_offset(int thread_id) { + + int warp_id = ((thread_id % kThreads) / Detail::kWarpSize); + int lane_id = (thread_id % Detail::kWarpSize); + + // + // compute warp-level offset + // + + // This is the shape of the entire area covered by a warp's memory access (in units of vectors) + layout::PitchLinearCoord warp_footprint{ + Detail::WarpThreadArrangement::kContiguous * Iterations::kContiguous, + Detail::WarpThreadArrangement::kStrided + }; + + // This is the offset of a specific warp (in units of vectors) + layout::PitchLinearCoord warp_offset{ + (warp_id % Detail::kWarpsContiguous), + (warp_id / Detail::kWarpsContiguous) + }; + + // This is the offset of a specific thread within a warp (units of vectors) + layout::PitchLinearCoord thread_offset_in_warp{ + lane_id % Detail::WarpThreadArrangement::kContiguous, + lane_id / Detail::WarpThreadArrangement::kContiguous + }; + + // This is the offset of a thread within a threadblock tile (units of vectors) + layout::PitchLinearCoord thread_offset_in_threadblock_tile_vec = + warp_footprint * warp_offset + thread_offset_in_warp; + + // This is the offset of a thread within a threadblock tile (units of elements) + layout::PitchLinearCoord thread_offset_in_threadblock_tile_base{ + thread_offset_in_threadblock_tile_vec.contiguous() * kElementsPerAccess, + thread_offset_in_threadblock_tile_vec.strided() + }; + + return thread_offset_in_threadblock_tile_base; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Strip-mines a pitch-linear tile among a given number of threads, first along the contiguous +/// dimension then along the strided dimension, while each thread access a 2D thread-tile. +/// +/// The tile must be divisible by the thread count such that all threads may execute the same +/// number of iterations with the same delta to exhaustively cover the tile. +/// +/// This class satisfies the "RegularThreadMapping" concept. +template < + typename Shape_, + int Threads, + typename ThreadTileShape +> +struct PitchLinear2DThreadTileStripminedThreadMap; + + +template < + typename Shape_, + int Threads +> +struct PitchLinear2DThreadTileStripminedThreadMap >{ + + /// Tensor coordinate + using TensorCoord = layout::PitchLinearCoord; + + /// Tile shape + using Shape = Shape_; + + /// Access Shape of each thread + using ThreadAccessShape = mctlass::layout::PitchLinearShape<4, 4>; + //using ThreadAccessShape = ThreadTileShape; + + /// Number of threads total + static int const kThreads = Threads; + + /// Extract length of each access from Layout + static int const kElementsPerAccess = ThreadAccessShape::kContiguous; + + static_assert(!(kElementsPerAccess % 4) , "kElementsPerAccess, needs to be multiple of 4 (32bits)"); + + /// Internal implementation details + struct Detail { + + static_assert(!(ThreadAccessShape::kContiguous % 4), "ThreadAccessShape, needs to be multiple of 4"); + + static_assert(!(Shape::kContiguous % ThreadAccessShape::kContiguous), ""); + + static_assert(!((Shape::kContiguous * Shape::kStrided) % (kThreads * ThreadAccessShape::kCount)), + "Shape must be divisible thread count * accesses per thread."); + + /// Shape of the tile in units of vectors + using ShapeVec = layout::PitchLinearShape< + Shape::kContiguous / ThreadAccessShape::kContiguous, + Shape::kStrided / ThreadAccessShape::kStrided + >; + + static_assert( + (Threads < ShapeVec::kContiguous && !(ShapeVec::kContiguous % kThreads)) || + (!(kThreads % ShapeVec::kContiguous) && !(ShapeVec::kStrided % (kThreads / ShapeVec::kContiguous))), + "Shape must be divisible by number of iterations of each thread." + ); + }; + + /// Number of iterations by each thread + using Iterations = typename platform::conditional< + Threads >= Detail::ShapeVec::kContiguous, + layout::PitchLinearShape< + 1, + // Redo the comparison here to work around divide by zero compiler + // error. The compiler evaluates both path of platform::conditional. + (Threads >= Detail::ShapeVec::kContiguous + ? Detail::ShapeVec::kStrided / + (kThreads / Detail::ShapeVec::kContiguous) + : 0)>, + layout::PitchLinearShape>::type; + + /// Interval between accesses along each dimension of the tensor's logical coordinate space + /// (in units of Elements) + using Delta = typename platform::conditional< + Threads >= Detail::ShapeVec::kContiguous, + layout::PitchLinearShape< + Shape::kContiguous, + kThreads * ThreadAccessShape::kStrided / Detail::ShapeVec::kContiguous + >, + layout::PitchLinearShape< + kThreads * ThreadAccessShape::kContiguous, + 1 + > + >::type; + + /// Maps thread ID to a coordinate offset within the tensor's logical coordinate space + /// (in units of Elements) + MCTLASS_HOST_DEVICE + static TensorCoord initial_offset(int thread_id) { + + return TensorCoord( + (thread_id % Detail::ShapeVec::kContiguous) * ThreadAccessShape::kContiguous, + (thread_id / Detail::ShapeVec::kContiguous) * ThreadAccessShape::kStrided); + } +}; + +/// Thread Mapping a 2D threadtiled mapping as a tranposed Pitchlinear2DThreadTile mapping +template +struct TransposePitchLinearThreadMap2DThreadTile { + /// Underlying ThreadMap + using ThreadMap = ThreadMap_; + + /// Tensor coordinate + using TensorCoord = typename ThreadMap::TensorCoord; + + /// Tile shape + using Shape = typename ThreadMap::Shape; + + /// Number of threads total + static int const kThreads = ThreadMap::kThreads; + + /// Extract vector length from Layout + static int const kElementsPerAccess = ThreadMap::kElementsPerAccess; + + + static_assert(kElementsPerAccess > 1 , "Simt transpose requires elements per access to be 1"); + ///< Iterations along each dimension (concept: PitchLinearShape) + using Iterations = + layout::PitchLinearShape; + + static_assert(Iterations::kCount, "Number of iterations must be non-zero"); + + /// Shape of access by each thread + using ThreadAccessShape = typename ThreadMap::ThreadAccessShape; + + ///< Delta betweeen accesses (units of elements, concept: PitchLinearShape) + using Delta = + layout::PitchLinearShape; + + + /// Maps thread ID to a coordinate offset within the tensor's logical + /// coordinate space Note this is slightly different from the one of + /// PitchLinearWarpRakedThreadMap. + MCTLASS_HOST_DEVICE + static TensorCoord initial_offset(int thread_id) { + + TensorCoord coord = ThreadMap::initial_offset(thread_id); + return TensorCoord( + coord.strided(), + coord.contiguous() + ); + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/thread/transpose.h b/csrc/mctlass/include/mctlass/transform/thread/transpose.h new file mode 100644 index 0000000..8e7aec7 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/thread/transpose.h @@ -0,0 +1,164 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Basic copy routines for tensor views +*/ + +#pragma once + +namespace mctlass { +namespace transform { +namespace thread { + +#if !defined(__MACA_ARCH__) + struct ucharHolder + { + union + { + unsigned char c[4]; + unsigned int ui; + }; + }; + + struct uchar2Holder + { + union + { + unsigned int ui[2]; + unsigned char c[8]; + }; + }; + + MCTLASS_DEVICE + unsigned int __byte_perm_maca(unsigned int x,unsigned int y,unsigned int s)//only found uint type declaration in cuda,so we not implemented int type + { + + struct uchar2Holder cHoldVal; + struct ucharHolder cHoldOut; + + cHoldVal.ui[0]=x; + cHoldVal.ui[1]=y; + + cHoldOut.c[0] = cHoldVal.c[((s << (31 - 2)) >> (31 - 2 + 0))]; + cHoldOut.c[1] = cHoldVal.c[((s << (31 - 6)) >> (31 - 6 + 4))]; + cHoldOut.c[2] = cHoldVal.c[((s << (31 - 10)) >> (31 - 10 + 8))]; + cHoldOut.c[3] = cHoldVal.c[((s << (31 - 14)) >> (31 - 14 + 12))]; + + return cHoldOut.ui; + + } + +#endif + +/// Transforms a fragment by doing a transpose +template < + int ElementCount, + typename TransposeShape, + typename Element +> struct Transpose; + +/// Specialization for int8_t 4x4 transpose +template +struct Transpose , int8_t> { + + static const int kElementCount = ElementCount_; + using TransposeShape = layout::PitchLinearShape<4,4>; + using Element = int8_t; + using Fragment = mctlass::Array; + + static_assert(!(kElementCount % TransposeShape::kCount), "Shape needs to be multiple of 16 elements to do a 4x4 transpose"); + + MCTLASS_DEVICE + void transform(Fragment& dst, Fragment& src) { + + // Expose src/dst as int arrays. + int* src_int = reinterpret_cast(&src); + int* dst_int = reinterpret_cast(&dst); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kElementCount / TransposeShape::kCount; i++){ + + int const i0 = 4 * i + 0; + int const i1 = 4 * i + 1; + int const i2 = 4 * i + 2; + int const i3 = 4 * i + 3; + + int a0 = src_int[i0]; + int a1 = src_int[i1]; + int a2 = src_int[i2]; + int a3 = src_int[i3]; + + int b0, b1, b2, b3, c0; + #if defined(__MACA_ARCH__) + b0 = __byte_perm(a0, a1, 0x0040); + c0 = __byte_perm(a2, a3, 0x0040); + b0 = __byte_perm(b0, c0, 0x5410); + + b1 = __byte_perm(a0, a1, 0x0051); + c0 = __byte_perm(a2, a3, 0x0051); + b1 = __byte_perm(b1, c0, 0x5410); + + b2 = __byte_perm(a0, a1, 0x0062); + c0 = __byte_perm(a2, a3, 0x0062); + b2 = __byte_perm(b2, c0, 0x5410); + + b3 = __byte_perm(a0, a1, 0x0073); + c0 = __byte_perm(a2, a3, 0x0073); + b3 = __byte_perm(b3, c0, 0x5410); + #else + b0 = __byte_perm_maca(a0, a1, 0x0040); + c0 = __byte_perm_maca(a2, a3, 0x0040); + b0 = __byte_perm_maca(b0, c0, 0x5410); + + b1 = __byte_perm_maca(a0, a1, 0x0051); + c0 = __byte_perm_maca(a2, a3, 0x0051); + b1 = __byte_perm_maca(b1, c0, 0x5410); + + b2 = __byte_perm_maca(a0, a1, 0x0062); + c0 = __byte_perm_maca(a2, a3, 0x0062); + b2 = __byte_perm_maca(b2, c0, 0x5410); + + b3 = __byte_perm_maca(a0, a1, 0x0073); + c0 = __byte_perm_maca(a2, a3, 0x0073); + b3 = __byte_perm_maca(b3, c0, 0x5410); + #endif + dst_int[i0] = b0; + dst_int[i1] = b1; + dst_int[i2] = b2; + dst_int[i3] = b3; + } + } +}; + +} // namespace thread +} // namespace layout +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/transform/thread/unary_op.h b/csrc/mctlass/include/mctlass/transform/thread/unary_op.h new file mode 100644 index 0000000..dc293e6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/thread/unary_op.h @@ -0,0 +1,105 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/complex.h" + +namespace mctlass { +namespace transform { +namespace thread { + +namespace UnaryTransform { + struct Identity; ///< None (i.e., identity) + struct Conjugate; ///< Complex conjugate +} + +/// Element-wise unary operator that transforms one element of a fragment at a time +template< + typename FragmentIn, ///< Input Fragment + typename FragmentOut,///< Output Fragment + typename Transform> ///< Unary transform operator +class UnaryOp +{ + public: + MCTLASS_DEVICE + static FragmentOut execute(FragmentIn &in) + { + static_assert(FragmentIn::kElements == FragmentOut::kElements, "Number of elements must match."); + static_assert(platform::is_same::value || + platform::is_same::value, + "Unary Operator not supported."); + + FragmentOut out; + if (platform::is_same::value ) + { + MCTLASS_PRAGMA_UNROLL + for (int i=0; i < FragmentIn::kElements; ++i){ + out[i] = static_cast(in[i]); + } + } + else if (platform::is_same::value ) + { + for (int i=0; i < FragmentIn::kElements; ++i){ + out[i] = conj(static_cast(in[i])); + } + } + return out; + } +}; + +template +class UnaryOp +{ + public: + MCTLASS_DEVICE + static FragmentIn execute(FragmentIn &in) + { + static_assert(platform::is_same::value || + platform::is_same::value, + "Unary Operator not supported."); + + if (platform::is_same::value ) + { + return in; + } + else if (platform::is_same::value ) + { + for(int i=0; i < FragmentIn::kElements; ++i){ + in[i] = conj(in[i]); + } + } + return in; + } + }; + } + } +} diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/ell_iterator.h b/csrc/mctlass/include/mctlass/transform/threadblock/ell_iterator.h new file mode 100644 index 0000000..87931a9 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/ell_iterator.h @@ -0,0 +1,199 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Ell iterator for matrix of indices (ellColInd matrix) +*/ + +#pragma once + +namespace mctlass { +namespace transform { +namespace threadblock { + +namespace ell{ + +constexpr unsigned int SmemPow = 8; +constexpr unsigned int SmemStages = 2; +constexpr unsigned int SmemSize = 1 << SmemPow; +constexpr unsigned int SmemMask = (SmemSize*SmemStages-1); + +class SharedStorage{ + public: + Array array; +}; + +class Iterator{ + public: + using Layout = layout::PitchLinear; + using LongIndex = typename Layout::LongIndex; + + private: + const int *gmem_col_idx_; + int *smem_col_idx_; + const int block_size_; + const int base_idx_; + const int k_shape_; + const int ell_increment_; + const int array_length_; + int col_idx_base_; + int residue_; + int counter_; + + int pow2_; + int residue_shape_; + + int smem_offset_; + int smem_stage_; + int gmem_offset_; + + int lane_; + + bool is_pow2_; + bool is_residue_tile_; + + public: + MCTLASS_DEVICE + void load_ell_indices(){ + for(int i=threadIdx.x; i= 0) ? gmem_col_idx : -1; + } + gmem_offset_ += SmemSize; + smem_stage_ ^= 1; + } + + MCTLASS_DEVICE + Iterator( + SharedStorage& shared_storage_base, + const int* col_idx, + const int& block_size, + const int& base_idx, + const int k_shape, + const int& problem_size_k, + const int& ell_stride, + const int& thread_idx) + : residue_(0), + counter_(0), + smem_offset_(0), + smem_stage_(0), + gmem_offset_(0), + block_size_(block_size), + base_idx_(base_idx), + k_shape_(k_shape), + ell_increment_(ell_stride * block_size), + array_length_((problem_size_k + block_size_ - 1) / block_size_), + residue_shape_(problem_size_k % k_shape_), + is_residue_tile_(residue_shape_ != 0), + smem_col_idx_(reinterpret_cast(&shared_storage_base.array)), + gmem_col_idx_(const_cast(col_idx)), + lane_(thread_idx % 32) { + + load_ell_indices(); + __syncthreads(); + + is_pow2_ = ((block_size_ & (block_size_ - 1)) == 0); + if( is_pow2_ && k_shape <= block_size_ ) lane_ = 0; + + col_idx_base_ = smem_col_idx_[(smem_offset_ + lane_) & SmemMask] * ell_increment_; + + pow2_ = 0; + while(block_size_ >> (pow2_ + 1)) ++pow2_; + } + + MCTLASS_DEVICE + int get_blocksize(){ + return block_size_; + } + + MCTLASS_DEVICE + Iterator &operator++(){ + if(is_residue_tile_){ + residue_ += residue_shape_; + is_residue_tile_ = false; + } else { + residue_ += k_shape_; + } + + if(residue_ < block_size_){ + return *this; + } + + if((array_length_ > SmemSize) && (((smem_offset_ >> SmemPow) & 1) != smem_stage_)) + load_ell_indices(); + + if(residue_ == block_size_){ + ++smem_offset_; + counter_ += ell_increment_; + residue_ = 0; + col_idx_base_ = smem_col_idx_[(smem_offset_ + lane_) & SmemMask] * ell_increment_ - counter_; + return *this; + } + + if(is_pow2_){ + smem_offset_ += residue_ >> pow2_; + counter_ += (residue_ >> pow2_) * ell_increment_; + residue_ = residue_ & ((1 << pow2_) - 1); + } + else { + smem_offset_ += residue_ / block_size_; + counter_ += (residue_ / block_size_) * ell_increment_; + residue_ %= block_size_; + } + + col_idx_base_ = smem_col_idx_[(smem_offset_ + lane_) & SmemMask] * ell_increment_ - counter_; + + return *this; + } + + MCTLASS_DEVICE + LongIndex get_offset(const int& idx) { + int num_jump_tiles; + if(is_pow2_) + num_jump_tiles = (idx + residue_) >> pow2_; + else + num_jump_tiles = (idx + residue_) / block_size_; + + int tmp = __shfl_sync(0xffffffff, col_idx_base_, num_jump_tiles); + return tmp - num_jump_tiles * ell_increment_; + } + + MCTLASS_DEVICE + LongIndex get_offset_fast() { + return col_idx_base_; + } +}; + +} +} +} +} diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/ell_predicated_tile_access_iterator.h b/csrc/mctlass/include/mctlass/transform/threadblock/ell_predicated_tile_access_iterator.h new file mode 100644 index 0000000..50883c4 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/ell_predicated_tile_access_iterator.h @@ -0,0 +1,1350 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Ell iterator for Blocked-Ell matrix (ellValue matrix) used with EllMmaMultistage +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/mctlass.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// EllPredicatedTileAccessIterator +/// +template +class EllPredicatedTileAccessIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of EllPredicatedTileAccessIterator for pitch-linear data. +/// +template +class EllPredicatedTileAccessIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + static int const kPredicatesPerByte = 4; + static int const kPredicatesPerWord = 4 * kPredicatesPerByte; + + static int const kPredicateCount = ThreadMap::Iterations::kCount * kAccessesPerVector; + + /// Number of 32b words containing predicates + static int const kPredicateByteCount = + (kPredicateCount + kPredicatesPerByte - 1) / kPredicatesPerByte; + static int const kPredicateWordCount = (kPredicateByteCount + 3) / 4; + + static unsigned const kPredicateMask = (1u << kPredicatesPerByte) - 1u; + + static_assert(kPredicateWordCount <= 4, "Too many predicates."); + + /// Predicate vector stores mask to guard accesses + using Mask = Array; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + friend EllPredicatedTileAccessIterator; + + private: + /// stride of pitch-linear layout (units of Element) + LongIndex stride_; + /// amount (in byte) to increment pointer to move to next access along + /// strided dimension + LongIndex inc_strided_; + /// amount (in byte) to increment pointer from last access to first access + /// of next tile + LongIndex inc_next_; + /// amount (in byte) to increment pointer from first access of current tile + /// to first access of next tile + LongIndex inc_advance_; + + public: + + // Default ctor + MCTLASS_HOST_DEVICE + Params(): stride_(0), inc_strided_(0), inc_next_(0), inc_advance_(0) { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) : stride_(layout.stride(0)) { + inc_strided_ = (LongIndex(stride_) * ThreadMap::Delta::kStrided) * + sizeof_bits::value / 8; + + if (kAdvanceRank) { + // advance along strided dimension + inc_advance_ = + Shape::kStrided * LongIndex(stride_) * sizeof_bits::value / 8; + } else { + // advance along contiguous dimension + inc_advance_ = Shape::kContiguous * sizeof_bits::value / 8; + } + + inc_next_ = inc_advance_ - LongIndex(ThreadMap::Iterations::kStrided - 1) * + ThreadMap::Delta::kStrided * LongIndex(stride_) * + sizeof_bits::value / 8; + }; + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + /// Parameters object with precomputed internal state + Params const ¶ms_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + + /// Guard predicates + uint32_t predicates_[kPredicateWordCount]; + + /// Size of tensor + TensorCoord extent_; + + /// Initial offset for each thread + TensorCoord thread_offset_; + + /// Offset to the first steady-state tile + TensorCoord residue_offset_; + + /// Initial offset to define ELL block + TensorCoord ell_offset_; + + /// Used for out-of-order visitation + bool is_residue_tile_; + + /// Iteration along vectors implied by the thread map + int iteration_vector_; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + public: + /// Computes predicates based on internally tracked per-thread offset. + MCTLASS_DEVICE + void compute_predicates_( + /// Extent of the matrix window + TensorCoord extent, + /// optionally, simplify predicate calculation during 'steady state' phase + bool is_steady_state = false) { + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = 0u; + } + + MCTLASS_PRAGMA_UNROLL + for (int access_idx = 0; access_idx < ThreadMap::Iterations::kCount * kAccessesPerVector; ++access_idx) { + + int s = access_idx / (ThreadMap::Iterations::kContiguous * kAccessesPerVector); + + int access_residual = access_idx % (ThreadMap::Iterations::kContiguous * kAccessesPerVector); + + int c = access_residual / kAccessesPerVector; + int v = access_residual % kAccessesPerVector; + + TensorCoord iteration_coord(c * ThreadMap::Delta::kContiguous + v * AccessType::kElements, + s * ThreadMap::Delta::kStrided); + + TensorCoord coord = thread_offset_ + iteration_coord; + + bool guard; + + if (is_steady_state) { + if (kAdvanceRank == 0) { + guard = (coord.strided() < extent.strided()); + } else { + guard = (coord.contiguous() < extent.contiguous()); + } + } else { + guard = (coord.strided() < extent.strided() && + coord.contiguous() < extent.contiguous()); + } + + int pred_idx = v + kAccessesPerVector * (c + ThreadMap::Iterations::kContiguous * s); + + int word_idx = pred_idx / kPredicatesPerWord; + int residual = pred_idx % kPredicatesPerWord; + int byte_idx = residual / kPredicatesPerByte; + int bit_idx = residual % kPredicatesPerByte; + + predicates_[word_idx] |= (unsigned(guard) << (byte_idx * 8 + bit_idx)); + + } + + } + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : params_(params), + pointer_(reinterpret_cast( + const_cast(pointer))), + extent_(extent), + is_residue_tile_(true) { + + TensorCoord residue_extent; + if (kAdvanceRank) { + + typename TensorCoord::Index residue_size = (extent_[kAdvanceRank] - threadblock_offset.strided()) % Shape::kStrided; + if (!residue_size) { + residue_size = Shape::kStrided; + } + + residue_offset_ = make_Coord(0, residue_size); + residue_extent = make_Coord( + extent_.contiguous(), + min(threadblock_offset.strided() + residue_size, extent_.strided()) + ); + } else { + + typename TensorCoord::Index residue_size = (extent_[kAdvanceRank] - threadblock_offset.contiguous()) % Shape::kContiguous; + if (!residue_size) { + residue_size = Shape::kContiguous; + } + + residue_offset_ = make_Coord(residue_size, 0); + + residue_extent = make_Coord( + min(extent_.contiguous(), threadblock_offset.contiguous() + residue_size), + extent_.strided() + ); + } + + // Per-thread offset in logical coordinates of tensor + ell_offset_ = ThreadMap::initial_offset(thread_id); + thread_offset_ = threadblock_offset + ThreadMap::initial_offset(thread_id); + + // update internal pointers + Layout layout(params_.stride_); + add_pointer_offset(layout(thread_offset_)); + + compute_predicates_(residue_extent, false); + + set_iteration_index(0); + } + + /// Construct a EllPredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id) + : EllPredicatedTileAccessIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += sizeof_bits::value * pointer_offset / 8; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_DEVICE + void add_tile_offset( + TensorCoord const &tile_offset) { + if (is_residue_tile_) { + + thread_offset_ += residue_offset_; + + Layout layout(params_.stride_); + add_pointer_offset(layout(residue_offset_)); + + compute_predicates_(extent_, true); + + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.strided() - 1); + pointer_ += Shape::kContiguous * tile_offset.contiguous(); + } else { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.contiguous() - 1); + pointer_ += Shape::kStrided * tile_offset.strided(); + } + } else { + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.strided()); + pointer_ += Shape::kContiguous * tile_offset.contiguous(); + } else { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.contiguous()); + pointer_ += Shape::kStrided * tile_offset.strided(); + } + } + is_residue_tile_ = false; + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast( + pointer_ + + iteration_contiguous_ * (ThreadMap::Delta::kContiguous * sizeof_bits::value) / 8) + iteration_vector_; + } + + /// Returns a k_location + MCTLASS_HOST_DEVICE + int get_k() const { + if(kAdvanceRank){ //strided + return ell_offset_.strided() + iteration_strided_ * ThreadMap::Delta::kStrided; + }else{ + return ell_offset_.contiguous() + iteration_contiguous_ * ThreadMap::Delta::kContiguous + iteration_vector_ * AccessType::kElements; + } + } + + MCTLASS_HOST_DEVICE + int get_stride() const { + if(kAdvanceRank) + return params_.stride_; + else + return 1; + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator &operator++() { + + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + + iteration_vector_ = 0; + ++iteration_contiguous_; + + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + iteration_contiguous_ = 0; + ++iteration_strided_; + + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + pointer_ += params_.inc_strided_; + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + iteration_strided_ = 0; + + // advance to next tile + pointer_ += params_.inc_next_; + + // now return to start tile - if the iterator is subsequently advanced, this + // subtraction as well as the subsequent integer addition are both elided by + // the compiler. + pointer_ -= params_.inc_advance_; + + return *this; + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator operator++(int) { + EllPredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = enable ? 0u : predicates_[i]; + } + + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = 0xffffffff; + } + + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = mask[i]; + } + + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + mask[i] = predicates_[i]; + } + } + + /// add mask for small tiles in ELL + MCTLASS_DEVICE + void ell_add_mask(int blocksize) { + + Mask mask; + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + mask[i] = 0u; + } + + MCTLASS_PRAGMA_UNROLL + for (int access_idx = 0; access_idx < ThreadMap::Iterations::kCount * kAccessesPerVector; ++access_idx) { + + int s = access_idx / (ThreadMap::Iterations::kContiguous * kAccessesPerVector); + + int access_residual = access_idx % (ThreadMap::Iterations::kContiguous * kAccessesPerVector); + + int c = access_residual / kAccessesPerVector; + int v = access_residual % kAccessesPerVector; + + TensorCoord iteration_coord(c * ThreadMap::Delta::kContiguous + v * AccessType::kElements, + s * ThreadMap::Delta::kStrided); + + TensorCoord coord = ell_offset_ + iteration_coord; + + bool guard; + + if (kAdvanceRank == 0) { + guard = (coord.strided() < blocksize); + } else { + guard = (coord.contiguous() < blocksize); + } + + int pred_idx = v + kAccessesPerVector * (c + ThreadMap::Iterations::kContiguous * s); + + int word_idx = pred_idx / kPredicatesPerWord; + int residual = pred_idx % kPredicatesPerWord; + int byte_idx = residual / kPredicatesPerByte; + int bit_idx = residual % kPredicatesPerByte; + + mask[word_idx] |= (unsigned(guard) << (byte_idx * 8 + bit_idx)); + + } + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + mask[i] &= predicates_[i]; + } + set_mask(mask); + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + + int pred_idx = + iteration_vector_ + kAccessesPerVector * (iteration_contiguous_ + iteration_strided_ * ThreadMap::Iterations::kContiguous); + + int word_idx = pred_idx / kPredicatesPerWord; + int residual = pred_idx % kPredicatesPerWord; + int byte_idx = residual / kPredicatesPerByte; + int bit_idx = residual % kPredicatesPerByte; + + bool pred = (predicates_[word_idx] & (1u << (byte_idx * 8 + bit_idx))) != 0; + return pred; + + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of EllPredicatedTileAccessIterator for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class EllPredicatedTileAccessIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = EllPredicatedTileAccessIterator< + layout::PitchLinearShape, Element, + layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap, AccessType>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend EllPredicatedTileAccessIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default ctor + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))){}; + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator( + ///< Precomputed parameters object + Params const ¶ms, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), + threadblock_offset.column())) {} + + /// Construct a EllPredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : EllPredicatedTileAccessIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + MCTLASS_HOST_DEVICE + int get_k() const { + return iterator_.get_k(); + } + + MCTLASS_HOST_DEVICE + int get_stride() const { + return iterator_.get_stride(); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator operator++(int) { + EllPredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// add mask for small tiles in ELL + MCTLASS_DEVICE + void ell_add_mask(int blocksize) { + iterator_.ell_add_mask(blocksize); + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of EllPredicatedTileAccessIterator for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class EllPredicatedTileAccessIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = EllPredicatedTileAccessIterator< + layout::PitchLinearShape, Element, + layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap, AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend EllPredicatedTileAccessIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default ctor + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))){}; + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator( + ///< Precomputed parameters object + Params const ¶ms, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a EllPredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : EllPredicatedTileAccessIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + MCTLASS_HOST_DEVICE + int get_k() const { + return iterator_.get_k(); + } + + MCTLASS_HOST_DEVICE + int get_stride() const { + return iterator_.get_stride(); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator operator++(int) { + EllPredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// add mask for small tiles in ELL + MCTLASS_DEVICE + void ell_add_mask(int blocksize) { + iterator_.ell_add_mask(blocksize); + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of EllPredicatedTileAccessIterator for column-major interleaved data. +/// It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// + +template +class EllPredicatedTileAccessIterator, + AdvanceRank, ThreadMap_, AccessType_> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::ColumnMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = EllPredicatedTileAccessIterator< + layout::PitchLinearShape, + Element, layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap, + AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend EllPredicatedTileAccessIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + MCTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.row() * kInterleavedK, + extent.column() / kInterleavedK), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.row() * kInterleavedK, + threadblock_offset.column() / kInterleavedK)) {} + + /// Construct a EllPredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : EllPredicatedTileAccessIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + MCTLASS_HOST_DEVICE + int get_k() const { + return iterator_.get_k(); + } + + MCTLASS_HOST_DEVICE + int get_stride() const { + return iterator_.get_stride(); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator operator++(int) { + EllPredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// add mask for small tiles in ELL + MCTLASS_DEVICE + void ell_add_mask(int blocksize) { + iterator_.ell_add_mask(blocksize); + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { return iterator_.valid(); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of EllPredicatedTileAccessIterator for row-major interleaved data. +/// It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class EllPredicatedTileAccessIterator, + AdvanceRank, ThreadMap_, AccessType_> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::RowMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = EllPredicatedTileAccessIterator< + layout::PitchLinearShape, + Element, layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap, + AccessType>; + + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend EllPredicatedTileAccessIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + MCTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.column() * kInterleavedK, + extent.row() / kInterleavedK), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column() * kInterleavedK, + threadblock_offset.row() / kInterleavedK)) {} + + /// Construct a EllPredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : EllPredicatedTileAccessIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + MCTLASS_HOST_DEVICE + int get_k() const { + return iterator_.get_k(); + } + + MCTLASS_HOST_DEVICE + int get_stride() const { + return iterator_.get_stride(); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileAccessIterator operator++(int) { + EllPredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// add mask for small tiles in ELL + MCTLASS_DEVICE + void ell_add_mask(int blocksize) { + iterator_.ell_add_mask(blocksize); + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { return iterator_.valid(); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/ell_predicated_tile_iterator.h b/csrc/mctlass/include/mctlass/transform/threadblock/ell_predicated_tile_iterator.h new file mode 100644 index 0000000..fe413e9 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/ell_predicated_tile_iterator.h @@ -0,0 +1,1315 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Ell iterator for Blocked-Ell matrix (ellValue matrix) used with EllMmaPipelined +*/ + +#pragma once + +#include "mctlass/arch/memory.h" +#include "mctlass/transform/threadblock/predicated_tile_access_iterator.h" + +#include "mctlass/transform/threadblock/ell_predicated_tile_access_iterator.h" +#include "mctlass/transform/threadblock/ell_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// EllPredicatedTileIterator +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +/// Regular tile iterator using a precomputed control structure to minimize register liveness +/// and integer arithmetic. +/// +/// Layout is assumed to be invariant at the time the precomputed "Params" object is constructed. +/// +/// Base pointer and tensor extents may be specified at the time the iterator is constructed. +/// Subsequently, they are assumed to be immutable. +/// +/// Adding a logical coordinate offset may be performed at the time the iterator is constructed. +/// Subsequent additions to logical coordinate offset may be performed but are relatively expensive. +/// +/// Visitation order is intended to first visit a "residual" tile that may be partially full in +/// both the advance dimension and the steady-state dimension. This is assumed to be the last +/// tile in the iteration sequence. Advancing an iterator that has just been constructed moves to +/// the first tile that is full in the advance dimension and recomputes predicates. Subsequent +/// accesses may be performed without updating internal predicates and are efficient in terms of +/// live register state and pointer arithmetic instructions. +/// +/// To be efficient, this assumes the iterator will be dereferenced and advanced at least once +/// outside any looping structure to minimize integer arithmetic. +/// +/// Acceses out of bounds are safe so long as `clear_mask()` is called prior to dereferencing +/// the iterator. +/// +/// +/// Example: +/// +/// An efficient pipeline structure may be constructed as follows: +/// +// template +// __global__ void kernel( +// typename Iterator::Params params, +// typename Iterator::Element *ptr, +// TensorCoord extent) { +// +// typename Iterator::Fragment fragment; +// +// TensorCoord threadblock_offset(0, 0); +// +// Iterator iter(params, ptr, extent, threadIdx.x, threadblock_offsets); +// +// +// fragment = *iter; // load "residue" tile first +// ++iter; // advance to first "steady state" tile and update internal masks +// +// +// #pragma unroll +// for (int i = Remaining - 1; i >= 0; --i) { +// +// f(fragment); +// +// if (!i) { +// iter.clear_mask(); // light-weight operation to clear masks - subsequent loads become NO-OPs. +// } +// +// fragment = *iter; // load tile during "steady state" phase +// ++iter; // advance to next tile - lightweight due to steady-state masks +// } +// } +// +// void host(TensorView view) { +// +// using Iterator = transform::threadblock::EllPredicatedTileIterator; +// +// typename Iterator::Params params(view.layout()); +// +// kernel(params, view.data()); +// } +/// +/// +template < + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + int AccessSize = ThreadMap::kElementsPerAccess +> +class EllPredicatedTileIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of EllPredicatedTileIterator for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class EllPredicatedTileIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + /// Type used for internal memory accesses + using AccessType = AlignedArray::value / 8)>; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = + EllPredicatedTileAccessIterator; + + static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename TileAccessIterator::Mask; + + /// Iterator for ELL storage + using EllIterator = typename mctlass::transform::threadblock::ell::Iterator; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + friend EllPredicatedTileIterator; + + private: + /// Parameters object + typename TileAccessIterator::Params params_; + + public: + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) : params_(layout) { } + + MCTLASS_HOST_DEVICE + Params() { } + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : address_iterator_(params.params_, pointer, extent, thread_id, + threadblock_offset) {} + + /// Construct a EllPredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : EllPredicatedTileIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator &operator++() { + if (kAdvanceRank) + address_iterator_.add_tile_offset({0, 1}); + else + address_iterator_.add_tile_offset({1, 0}); + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator operator++(int) { + EllPredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Returns a stride + MCTLASS_HOST_DEVICE + int get_stride() const { return address_iterator_.get_stride(); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { address_iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { address_iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { address_iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { address_iterator_.get_mask(mask); } + + /// add mask for small tiles in ELL + MCTLASS_HOST_DEVICE + void ell_add_mask(int blocksize) { address_iterator_.ell_add_mask(blocksize); } + + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + load_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + address_iterator_.set_iteration_index(idx); + char const *byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + + AccessType const *access_ptr = reinterpret_cast(byte_ptr); + + mctlass::arch::global_load( + frag_ptr[idx], access_ptr, address_iterator_.valid()); + + ++address_iterator_; + } + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_byte_offset(frag, 0); } + + MCTLASS_DEVICE + void load_with_ell_index(Fragment &frag, EllIterator &ell_iter) { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + address_iterator_.set_iteration_index(idx); + LongIndex ell_offset = 0; + + int k_offset = address_iterator_.get_k(); + ell_offset = ell_iter.get_offset(k_offset) * sizeof(Element); + + char const *byte_ptr = reinterpret_cast(address_iterator_.get()) + ell_offset; + + AccessType const *access_ptr = reinterpret_cast(byte_ptr); + + bool is_valid = address_iterator_.valid(); + is_valid = is_valid && (ell_offset >= 0); + + mctlass::arch::global_load( + frag_ptr[idx], access_ptr, is_valid); + + ++address_iterator_; + } + } + } + } + + MCTLASS_DEVICE + void load_with_ell_index_fast(Fragment &frag, EllIterator &ell_iter) { + + LongIndex ell_offset = ell_iter.get_offset_fast() * sizeof(Element); + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + address_iterator_.set_iteration_index(idx); + char const *byte_ptr = reinterpret_cast(address_iterator_.get()) + ell_offset; + + AccessType const *access_ptr = reinterpret_cast(byte_ptr); + + bool is_valid = address_iterator_.valid(); + is_valid = is_valid && (ell_offset >= 0); + + mctlass::arch::global_load( + frag_ptr[idx], access_ptr, is_valid); + + ++address_iterator_; + } + } + } + } + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + store_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) { + address_iterator_.set_iteration_index(0); + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + char *byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType *access_ptr = reinterpret_cast(byte_ptr); + + if (address_iterator_.valid()) { + *access_ptr = frag_ptr[idx]; + } + ++address_iterator_; + } + } + } + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_byte_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of EllPredicatedTileIterator for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize +> +class EllPredicatedTileIterator { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = EllPredicatedTileIterator< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessSize + >; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Iterator for ELL storage + using EllIterator = typename mctlass::transform::threadblock::ell::Iterator; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + + friend EllPredicatedTileIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) { + + } + }; + + +private: + + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + + /// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const &threadblock_offset ///< Initial offset of threadblock + ): + iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column()) + ) { } + + /// Construct a EllPredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ): EllPredicatedTileIterator(params, pointer, extent, thread_id, make_Coord(0, 0)) { } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator operator++(int) { + EllPredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Returns a stride + MCTLASS_HOST_DEVICE + int get_stride() const { return iterator_.get_stride(); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + iterator_.get_mask(mask); + } + + /// add mask for small tiles in ELL + MCTLASS_HOST_DEVICE + void ell_add_mask(int blocksize) { + iterator_.ell_add_mask(blocksize); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + MCTLASS_DEVICE + void load_with_ell_index(Fragment &frag, EllIterator& ell_iter) { + iterator_.load_with_ell_index(frag, ell_iter); + } + + MCTLASS_DEVICE + void load_with_ell_index_fast(Fragment &frag, EllIterator& ell_iter) { + iterator_.load_with_ell_index_fast(frag, ell_iter); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of EllPredicatedTileIterator for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize +> +class EllPredicatedTileIterator { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = EllPredicatedTileIterator< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize + >; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Iterator for ELL storage + using EllIterator = typename mctlass::transform::threadblock::ell::Iterator; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + + friend EllPredicatedTileIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) { + + }; + }; + + +private: + + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + + /// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const &threadblock_offset ///< Initial offset of threadblock + ): + iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row()) + ) { } + + /// Construct a EllPredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ): EllPredicatedTileIterator(params, pointer, extent, thread_id, make_Coord(0, 0)) { } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator operator++(int) { + EllPredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Returns a stride + MCTLASS_HOST_DEVICE + int get_stride() const { return iterator_.get_stride(); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + iterator_.get_mask(mask); + } + + /// add mask for small tiles in ELL + MCTLASS_HOST_DEVICE + void ell_add_mask(int blocksize) { + iterator_.ell_add_mask(blocksize); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + MCTLASS_DEVICE + void load_with_ell_index(Fragment &frag, EllIterator& ell_iter) { + iterator_.load_with_ell_index(frag, ell_iter); + } + + MCTLASS_DEVICE + void load_with_ell_index_fast(Fragment &frag, EllIterator& ell_iter) { + iterator_.load_with_ell_index_fast(frag, ell_iter); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of EllPredicatedTileIterator for interleaved data. It is mapped +/// to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// + +template +class EllPredicatedTileIterator, + AdvanceRank, ThreadMap_, AccessSize> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::ColumnMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = EllPredicatedTileIterator< + layout::PitchLinearShape, + Element, layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap, AccessSize>; + + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Iterator for ELL storage + using EllIterator = typename mctlass::transform::threadblock::ell::Iterator; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend EllPredicatedTileIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + MCTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.row() * kInterleavedK, + extent.column() / kInterleavedK), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.row() * kInterleavedK, + threadblock_offset.column() / kInterleavedK)) {} + + /// Construct a EllPredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : EllPredicatedTileIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator operator++(int) { + EllPredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Returns a stride + MCTLASS_HOST_DEVICE + int get_stride() const { return iterator_.get_stride(); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// add mask for small tiles in ELL + MCTLASS_HOST_DEVICE + void ell_add_mask(int blocksize) { iterator_.ell_add_mask(blocksize); } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + MCTLASS_DEVICE + void load_with_ell_index(Fragment &frag, EllIterator& ell_iter) { + iterator_.load_with_ell_index(frag, ell_iter); + } + + MCTLASS_DEVICE + void load_with_ell_index_fast(Fragment &frag, EllIterator& ell_iter) { + iterator_.load_with_ell_index_fast(frag, ell_iter); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of EllPredicatedTileIterator for interleaved-32 data. It is +/// mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class EllPredicatedTileIterator, + AdvanceRank, ThreadMap_, AccessSize> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::RowMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = EllPredicatedTileIterator< + layout::PitchLinearShape, + Element, layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap, AccessSize>; + + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend EllPredicatedTileIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + MCTLASS_HOST_DEVICE + Params() {} + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.column() * kInterleavedK, + extent.row() / kInterleavedK), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column() * kInterleavedK, + threadblock_offset.row() / kInterleavedK)) {} + + /// Construct a EllPredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : EllPredicatedTileIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + EllPredicatedTileIterator operator++(int) { + EllPredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Returns a stride + MCTLASS_HOST_DEVICE + int get_stride() const { return iterator_.get_stride(); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// add mask for small tiles in ELL + MCTLASS_HOST_DEVICE + void ell_add_mask(int blocksize) { iterator_.ell_add_mask(blocksize); } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/predicated_scale_bias_vector_access_iterator.h b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_scale_bias_vector_access_iterator.h new file mode 100644 index 0000000..be7b1c3 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_scale_bias_vector_access_iterator.h @@ -0,0 +1,375 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates calculating the address and predicates to the load of scale and bias vectors. + + This iterator uses masks to guard out-of-bounds accesses. + + It can be used to load the gamma and beta vectors of layernorm which is loop variant. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/mctlass.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/conv/threadblock/conv2d_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedScaleBiasVectorAccessIterator +/// +template +class PredicatedScaleBiasVectorAccessIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for fprop pitch-linear data. +/// +template +class PredicatedScaleBiasVectorAccessIterator { + public: + + using ThreadblockShape = ThreadblockShape_; + using Element = Element_; + using Layout = layout::PitchLinear; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + static int const kElementsPerAccess = 128 / sizeof_bits::value; + static int const kThreads = ThreadblockShape::kContiguous / kElementsPerAccess; + + using AccessType = AlignedArray; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + /// Internal pointer to first access of tile + BytePointer pointer_; + + TensorCoord thread_offset_; + + int problem_size_k_; + + /// Used for out-of-order visitation + bool is_residue_tile_; + + bool guard_; + + TensorCoord::Index residue_size_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + /// Extent of tensor + int problem_size_k, + /// Pointer to the start of the scale vector + ConstPointer scale_pointer, + /// Pointer to the start of the bias vector + ConstPointer bias_pointer, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) { + pointer_ = (thread_id < kThreads) + ? reinterpret_cast( + const_cast(scale_pointer)) + : reinterpret_cast( + const_cast(bias_pointer)); + + // Per-thread offset in logical coordinates of tensor + int thread_base = (thread_id < kThreads) ? 0 : kThreads; + + problem_size_k_ = problem_size_k; + + is_residue_tile_ = true; + + residue_size_ = (problem_size_k_ - threadblock_offset.contiguous()) % ThreadblockShape::kContiguous; + + if (residue_size_ == 0) { + residue_size_ = ThreadblockShape::kContiguous; + } + + guard_ = ((thread_id - thread_base) * kElementsPerAccess) < residue_size_; + + thread_offset_ = + threadblock_offset + + TensorCoord((thread_id - thread_base) * kElementsPerAccess, 0); + + set_iteration_index(0); + } + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + /// Extent of tensor + int problem_size_k, + /// Pointer to start of scale vector + ConstPointer scale_pointer, + /// Pointer to start of scale vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id) + : PredicatedScaleBiasVectorAccessIterator(problem_size_k, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) {} + + /// Advances an iterator along logical dimensions of matrix in units of whole threadblock tiles + MCTLASS_DEVICE + void add_tile_offset( + TensorCoord const &tile_offset) { + + guard_ = threadIdx.x < kThreads * 2; + + TensorCoord offset = is_residue_tile_ ? + TensorCoord(residue_size_ + ThreadblockShape::kContiguous * (tile_offset.contiguous() - 1), 0) + : TensorCoord(ThreadblockShape::kContiguous * tile_offset.contiguous(), 0); + + thread_offset_ = + thread_offset_ + + offset; + + is_residue_tile_ = false; + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + + return reinterpret_cast( + pointer_ + + (thread_offset_.contiguous() * sizeof_bits::value / 8)); + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator &operator++() { + return *this; + } + + /// Increment and return an instance to self. + MCTLASS_DEVICE + PredicatedScaleBiasVectorAccessIterator operator++(int) { + PredicatedScaleBiasVectorAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + guard_ &= (!enable); + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return guard_; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedScaleBiasVectorAccessIterator { + public: + + using ThreadblockShape = ThreadblockShape_; + using Element = Element_; + using Layout = layout::RowMajor; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedScaleBiasVectorAccessIterator< + layout::PitchLinearShape, + Element, + layout::PitchLinear>; + + using AccessType = typename UnderlyingIterator::AccessType; + static int const kElementsPerAccess = UnderlyingIterator::kElementsPerAccess; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + ///< Extent of tensor + int problem_size_k, + ///< Pointer to the start of the scale vector + ConstPointer scale_pointer, + ///< Pointer to the start of the bias vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(problem_size_k, scale_pointer, bias_pointer, + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator( + int problem_size_k, ///< Extent of tensor + ConstPointer scale_pointer, ///< Pointer to the start of the scale vector + ConstPointer bias_pointer, ///< Pointer to the start of the bias vector + int thread_id ///< ID of each participating thread + ) + : PredicatedScaleBiasVectorAccessIterator(problem_size_k, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// threadblock tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorAccessIterator operator++(int) { + PredicatedScaleBiasVectorAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/predicated_scale_bias_vector_iterator.h b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_scale_bias_vector_iterator.h new file mode 100644 index 0000000..ee0f7c1 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_scale_bias_vector_iterator.h @@ -0,0 +1,335 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates calculating the address and predicates to the load of scale and bias vectors. + + This iterator uses masks to guard out-of-bounds accesses. + + This can be used to load var and mean vectors in layernorm which is loop invariant. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/mctlass.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedScaleBiasVectorIterator +/// +template +class PredicatedScaleBiasVectorIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for wgrad pitch-linear data. +/// +template +class PredicatedScaleBiasVectorIterator { + public: + + using WarpShape = WarpShape_; + using Element = Element_; + using Layout = layout::PitchLinear; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + static int const kElementsPerAccess = 1; + + using AccessType = AlignedArray; + + static int const kIterations = WarpShape::kContiguous / 8; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array<__half2, 2 * kIterations * kElementsPerAccess>; + + private: + // + // Data members + // + + /// Internal pointer to first access of tile + ConstPointer scale_pointer_; + ConstPointer bias_pointer_; + + /// Size of tensor + int problem_size_; + + int32_t thread_offset_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorIterator( + /// Extent of tensor + int problem_size, + /// Pointer to the start of the scale vector + ConstPointer scale_pointer, + /// Pointer to the start of the bias vector + ConstPointer bias_pointer, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : problem_size_(problem_size), + scale_pointer_(scale_pointer), + bias_pointer_(bias_pointer) { + + thread_offset_ = threadblock_offset.contiguous() + (thread_id % 32) / 4; + } + + /// Construct a PredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorIterator( + /// Extent of tensor + int problem_size, + /// Pointer to start of scale vector + ConstPointer scale_pointer, + /// Pointer to start of scale vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id) + : PredicatedScaleBiasVectorIterator(problem_size, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + /// Advances an iterator along logical dimensions of matrix in units of whole warp tiles + MCTLASS_DEVICE + void add_tile_offset( + TensorCoord const &tile_offset) { + + thread_offset_ += (WarpShape::kContiguous * tile_offset.contiguous()); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + #if 1 + frag.fill(__float2half2_rn(0.0f)); + #else + __half2 v; + v.x = 0; + v.y = 0; + frag.fill(v); + #endif + __half2 *frag_ptr = reinterpret_cast<__half2 *>(&frag); + + // load scale + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + + mctlass::arch::global_load< + __half, + sizeof(AccessType) + >( + frag_ptr[c * 2].x, + scale_pointer_ + thread_offset_ + c * 8, + (thread_offset_ + c * 8) < problem_size_ + ); + } + + // load bias + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + + mctlass::arch::global_load< + __half, + sizeof(AccessType) + >( + frag_ptr[c * 2 + 1].x, + bias_pointer_ + thread_offset_ + c * 8, + (thread_offset_ + c * 8) < problem_size_ + ); + } + + // duplicate scale + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + frag_ptr[c * 2].y = frag_ptr[c * 2].x; + } + + // duplicate bias + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + frag_ptr[c * 2 + 1].y = frag_ptr[c * 2 + 1].x; + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedScaleBiasVectorIterator { + public: + + using WarpShape = WarpShape_; + using Element = Element_; + using Layout = layout::RowMajor; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedScaleBiasVectorIterator< + layout::PitchLinearShape, + Element, + layout::PitchLinear>; + + using AccessType = typename UnderlyingIterator::AccessType; + static int const kElementsPerAccess = UnderlyingIterator::kElementsPerAccess; + using Fragment = typename UnderlyingIterator::Fragment; + + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorIterator( + ///< Extent of tensor + int problem_size, + ///< Pointer to the start of the scale vector + ConstPointer scale_pointer, + ///< Pointer to the start of the bias vector + ConstPointer bias_pointer, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(problem_size, scale_pointer, bias_pointer, + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedScaleBiasVectorIterator( + int problem_size, ///< Extent of tensor + ConstPointer scale_pointer, ///< Pointer to the start of the scale vector + ConstPointer bias_pointer, ///< Pointer to the start of the bias vector + int thread_id ///< ID of each participating thread + ) + : PredicatedScaleBiasVectorIterator(problem_size, + scale_pointer, bias_pointer, + thread_id, make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// threadblock tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + iterator_.load(frag); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator.h b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator.h new file mode 100644 index 0000000..8e06b1c --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator.h @@ -0,0 +1,2117 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates calculating the address and predicates to the load of tiles + from pitch-linear rank=2 tensors. + + This iterator uses masks to guard out-of-bounds accesses. The first tile this + iterator visits maybe partial, then the remaining tiles are complete. So, we + only need to compute the predicates twice, once before the first tile and + once for the remaining full tiles which can share the same predicates. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/mctlass.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/permute.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/transform/threadblock/predicated_tile_access_iterator_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileAccessIteratorPredicates +/// +template +class PredicatedTileAccessIteratorPredicates { + public: + using Shape = Shape_; + using Element = Element_; + using Layout = Layout_; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorCoord = typename Layout::TensorCoord; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + static int const kPredicatesPerByte = 4; + static int const kPredicatesPerWord = 4 * kPredicatesPerByte; + + static int const kPredicateCount = ThreadMap::Iterations::kCount * kAccessesPerVector; + + /// Number of 32b words containing predicates + static int const kPredicateByteCount = + (kPredicateCount + kPredicatesPerByte - 1) / kPredicatesPerByte; + static int const kPredicateWordCount = (kPredicateByteCount + 3) / 4; + + static unsigned const kPredicateMask = (1u << kPredicatesPerByte) - 1u; + + static_assert(kPredicateWordCount <= 4, "Too many predicates."); + + /// Predicate vector stores mask to guard accesses + using Mask = Array; + +// private: + /// Guard predicates + uint32_t predicates_[kPredicateWordCount]; + + /// Size of tensor + TensorCoord extent_; + + /// Initial offset for each thread + TensorCoord thread_offset_; + + /// Offset to the first steady-state tile + TensorCoord residue_offset_; + + /// Iteration along vectors implied by the thread map + int iteration_vector_; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + public: + /// Computes predicates based on internally tracked per-thread offset. + MCTLASS_DEVICE + void compute_predicates_( + /// Extent of the matrix window + TensorCoord extent, + /// optionally, simplify predicate calculation during 'steady state' phase + bool is_steady_state = false) { + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = 0u; + } + + MCTLASS_PRAGMA_UNROLL + for (int access_idx = 0; access_idx < ThreadMap::Iterations::kCount * kAccessesPerVector; ++access_idx) { + + int s = access_idx / (ThreadMap::Iterations::kContiguous * kAccessesPerVector); + + int access_residual = access_idx % (ThreadMap::Iterations::kContiguous * kAccessesPerVector); + + int c = access_residual / kAccessesPerVector; + int v = access_residual % kAccessesPerVector; + + TensorCoord iteration_coord(c * ThreadMap::Delta::kContiguous + v * AccessType::kElements, + s * ThreadMap::Delta::kStrided); + + TensorCoord coord = thread_offset_ + iteration_coord; + + bool guard; + + if (is_steady_state) { + if (kAdvanceRank == 0) { + guard = (coord.strided() < extent.strided()); + } else { + guard = (coord.contiguous() < extent.contiguous()); + } + } else { + guard = (coord.strided() < extent.strided() && + coord.contiguous() < extent.contiguous()); + } + + int pred_idx = v + kAccessesPerVector * (c + ThreadMap::Iterations::kContiguous * s); + + int word_idx = pred_idx / kPredicatesPerWord; + int residual = pred_idx % kPredicatesPerWord; + int byte_idx = residual / kPredicatesPerByte; + int bit_idx = residual % kPredicatesPerByte; + + predicates_[word_idx] |= (unsigned(guard) << (byte_idx * 8 + bit_idx)); + + } + + } + + MCTLASS_HOST_DEVICE + void set_predicates(int thread_id, TensorCoord const &threadblock_offset) { + + TensorCoord residue_extent; + if (kAdvanceRank) { + + typename TensorCoord::Index residue_size = (extent_[kAdvanceRank] - threadblock_offset.strided()) % Shape::kStrided; + if (!residue_size) { + residue_size = Shape::kStrided; + } + + residue_offset_ = make_Coord(0, residue_size); + residue_extent = make_Coord( + extent_.contiguous(), + min(threadblock_offset.strided() + residue_size, extent_.strided()) + ); + } else { + + typename TensorCoord::Index residue_size = (extent_[kAdvanceRank] - threadblock_offset.contiguous()) % Shape::kContiguous; + if (!residue_size) { + residue_size = Shape::kContiguous; + } + + residue_offset_ = make_Coord(residue_size, 0); + + residue_extent = make_Coord( + min(extent_.contiguous(), threadblock_offset.contiguous() + residue_size), + extent_.strided() + ); + } + + // Per-thread offset in logical coordinates of tensor + thread_offset_ = threadblock_offset + ThreadMap::initial_offset(thread_id); + + compute_predicates_(residue_extent, false); + + set_iteration_index(0); + } + + /// Default constructor + PredicatedTileAccessIteratorPredicates() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorPredicates( + /// Extent of tensor + TensorCoord extent) + : extent_(extent) { + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorPredicates &operator++() { + + return *this; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = enable ? 0u : predicates_[i]; + } + + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = 0xffffffff; + } + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = mask[i]; + } + + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + mask[i] = predicates_[i]; + } + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() const { + + + int pred_idx = + iteration_vector_ + kAccessesPerVector * (iteration_contiguous_ + iteration_strided_ * ThreadMap::Iterations::kContiguous); + + int word_idx = pred_idx / kPredicatesPerWord; + int residual = pred_idx % kPredicatesPerWord; + int byte_idx = residual / kPredicatesPerByte; + int bit_idx = residual % kPredicatesPerByte; + + bool pred = (predicates_[word_idx] & (1u << (byte_idx * 8 + bit_idx))) != 0; + return pred; + + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileAccessIterator +/// +template +class PredicatedTileAccessIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for pitch-linear data. +/// +template +class PredicatedTileAccessIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingPredicates = PredicatedTileAccessIteratorPredicates< + Shape, Element, Layout, AdvanceRank, ThreadMap, AccessType>; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + static bool constexpr Permute = !platform::is_same::value + && !platform::is_same>::value; + + using Mask = typename UnderlyingPredicates::Mask; + + /// Uses a non-template class + struct Params : PredicatedTileAccessIteratorParams { + + using Base = PredicatedTileAccessIteratorParams; + + // Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) : + Base(layout.stride(0), + MakePredicatedTileAccessIteratorDesc()() + ) { } + + MCTLASS_HOST_DEVICE + Params(Base const &base) : + Base(base) { } + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + UnderlyingPredicates the_predicates; + + /// Parameters object with precomputed internal state + //Params const ¶ms_; + Params params_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + + /// Used for out-of-order visitation + bool is_residue_tile_; + + /// Below is used when Gather is turned on. We need to record strided_offset + /// and contiguous_offset separated to compute the offset by using + /// + /// offset = contiguous_offset + indices[strided_offset] + + + /// Gather indices + int const *indices_; + + /// Function to perform layout permutation and offset computation + PermuteLayout permute_layout_; + + /// Tracks thread's coordinate offset in the matrix for current tile. + /// This is only used in the following cases: + /// - when Gather is true, strided coordinate needed to access indices (contiguous offset is tracked via pointer_) + /// - when Permute is true, both coordinates are neeeded as input into permutation function (pointer_ is fixed) + TensorCoord coord_offset_; + + private: + /// Computes predicates based on internally tracked per-thread offset. + MCTLASS_DEVICE + void compute_predicates_( + /// Extent of the matrix window + TensorCoord extent, + /// optionally, simplify predicate calculation during 'steady state' phase + bool is_steady_state = false) { + the_predicates.compute_predicates_(extent, is_steady_state); + } + + public: + + /// Default constructor + PredicatedTileAccessIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset, + /// Gather indices + int const *indices = nullptr) + : params_(params), + pointer_(reinterpret_cast( + const_cast(pointer))), + the_predicates(extent), + is_residue_tile_(true), + indices_(indices), + permute_layout_(TensorCoord(extent.contiguous(), extent.strided()), params.stride_) { + + the_predicates.set_predicates(thread_id, threadblock_offset); + + if (Gather) { + assert(indices_); + } + + // update internal pointers + Layout layout(params_.stride_); + + if (!Gather && !Permute) { + add_pointer_offset(layout(the_predicates.thread_offset_)); + } else { + coord_offset_ = the_predicates.thread_offset_; + if (!Permute) { + add_pointer_offset(layout(make_Coord(coord_offset_.contiguous(), 0))); + } + } + } + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id) + : PredicatedTileAccessIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + the_predicates.set_iteration_index(index); + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += sizeof_bits::value * pointer_offset / 8; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_DEVICE + void add_tile_offset( + TensorCoord const &tile_offset) { + if (is_residue_tile_) { + + the_predicates.thread_offset_ += the_predicates.residue_offset_; + + the_predicates.compute_predicates_(the_predicates.extent_, true); + + Layout layout(params_.stride_); + + if (!Gather && !Permute) { + add_pointer_offset(layout(the_predicates.residue_offset_)); + + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.strided() - 1); + pointer_ += Shape::kContiguous * tile_offset.contiguous() * sizeof_bits::value / 8; + } else { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.contiguous() - 1); + pointer_ += Shape::kStrided * tile_offset.strided() * sizeof_bits::value / 8; + } + } else { + coord_offset_.strided() = the_predicates.thread_offset_.strided() + Shape::kStrided * (tile_offset.strided() - kAdvanceRank); + if (!Permute) { + add_pointer_offset(layout(make_Coord(the_predicates.residue_offset_.contiguous(), 0))); + add_pointer_offset(Shape::kContiguous * (tile_offset.contiguous() - (1 - kAdvanceRank))); + } else { + coord_offset_.contiguous() = the_predicates.thread_offset_.contiguous() + Shape::kContiguous * (tile_offset.contiguous() - (1 - kAdvanceRank)); + } + } + } else { + if (!Gather && !Permute) { + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.strided()); + pointer_ += Shape::kContiguous * tile_offset.contiguous(); + } else { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.contiguous()); + pointer_ += Shape::kStrided * tile_offset.strided(); + } + } else { + coord_offset_.strided() += Shape::kStrided * tile_offset.strided(); + if (!Permute) { + add_pointer_offset(Shape::kContiguous * tile_offset.contiguous()); + } else { + coord_offset_.contiguous() += Shape::kContiguous * tile_offset.contiguous(); + } + } + } + + is_residue_tile_ = false; + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + if (Gather || Permute) + { + if (!valid()) { + return nullptr; + } + + Index coord_contig = (Permute ? coord_offset_.contiguous() : 0) + the_predicates.iteration_contiguous_ * ThreadMap::Delta::kContiguous + the_predicates.iteration_vector_ * AccessType::kElements; + Index coord_strided = coord_offset_.strided() + the_predicates.iteration_strided_ * ThreadMap::Delta::kStrided; + if (Gather) { + coord_strided = indices_[coord_strided]; + } + + LongIndex offset = Permute ? permute_layout_(TensorCoord(coord_contig, coord_strided)) : (coord_strided * LongIndex(params_.stride_) + coord_contig); + return reinterpret_cast(pointer_ + OffsetBytes(offset)); + } + + return reinterpret_cast( + pointer_ + + the_predicates.iteration_contiguous_ * (ThreadMap::Delta::kContiguous * sizeof_bits::value) / 8) + the_predicates.iteration_vector_; + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator &operator++() { + + the_predicates.operator++(); + + ++the_predicates.iteration_vector_; + if (the_predicates.iteration_vector_ < kAccessesPerVector) { + return *this; + } + + the_predicates.iteration_vector_ = 0; + ++the_predicates.iteration_contiguous_; + + if (the_predicates.iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + + // Enter here only if (iteration_contiguous_ == ThreadMap::Iteration::kContiguous) + the_predicates.iteration_contiguous_ = 0; + ++the_predicates.iteration_strided_; + + if (the_predicates.iteration_strided_ < ThreadMap::Iterations::kStrided) { + if (!Gather && !Permute) { + pointer_ += params_.inc_strided_; + } + + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + the_predicates.iteration_strided_ = 0; + + if (!Gather && !Permute) { + // advance to next tile + pointer_ += params_.inc_next_; + + // now return to start tile - if the iterator is subsequently advanced, this + // subtraction as well as the subsequent integer addition are both elided by + // the compiler. + pointer_ -= params_.inc_advance_; + } + + return *this; + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator operator++(int) { + PredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + the_predicates.clear_mask(enable); + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + the_predicates.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + the_predicates.set_mask(mask); + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + the_predicates.get_mask(mask); + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() const { + return the_predicates.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for column-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileAccessIterator< + layout::PitchLinearShape, Element, + layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap, AccessType, + Gather, PermuteLayout>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))){}; + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const &base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + + /// Default constructor + PredicatedTileAccessIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + ///< Precomputed parameters object + Params const ¶ms, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset, + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), + threadblock_offset.column()), + indices) {} + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator operator++(int) { + PredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileAccessIterator< + layout::PitchLinearShape, Element, + layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap, AccessType, + Gather, PermuteLayout>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))){}; + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const &base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + + /// Default constructor + PredicatedTileAccessIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + ///< Precomputed parameters object + Params const ¶ms, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset, + /// Gather indices + int const *indices = nullptr) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row()), + indices) {} + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator operator++(int) { + PredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for affine rank 2 data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIterator, + AdvanceRank, ThreadMap_, AccessType_, false, + layout::NoPermute> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRankN<2>; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingPredicates = PredicatedTileAccessIteratorPredicates< + Shape, Element, layout::PitchLinear, AdvanceRank, ThreadMap, AccessType>; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingPredicates::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + friend PredicatedTileAccessIterator; + + private: + /// stride of pitch-linear layout (units of Element) + Coord stride_; + /// amount (in byte) to increment pointer to move to next access along + /// contiguous dimension + LongIndex inc_contiguous_; + /// amount (in byte) to increment pointer from first access of current + /// contiguous dimension to first access of next one. + LongIndex inc_strided_; + /// amount (in byte) to increment pointer from last access of current + /// contiguous dimension to first access of next one. + LongIndex inc_next_strided_; + /// amount (in byte) to increment pointer from last access to first access + /// of next tile + LongIndex inc_next_; + /// amount (in byte) to increment pointer from first access of current tile + /// to first access of next tile + LongIndex inc_advance_; + + public: + + // Default ctor + MCTLASS_HOST_DEVICE + Params(): stride_(0), inc_contiguous_(0), inc_strided_(0), inc_next_(0), inc_advance_(0) { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) : stride_({layout.stride(0), layout.stride(1)}) { + inc_contiguous_ = (LongIndex(stride_[0]) * ThreadMap::Delta::kContiguous) * + sizeof_bits::value / 8; + + inc_strided_ = (LongIndex(stride_[1]) * ThreadMap::Delta::kStrided) * + sizeof_bits::value / 8; + + inc_next_strided_ = inc_strided_ - LongIndex(ThreadMap::Iterations::kContiguous - 1) * inc_contiguous_; + + if (kAdvanceRank) { + // advance along strided dimension + inc_advance_ = + Shape::kStrided * LongIndex(stride_[1]) * sizeof_bits::value / 8; + } else { + // advance along contiguous dimension + inc_advance_ = Shape::kContiguous * stride_[0] * sizeof_bits::value / 8; + } + + inc_next_ = inc_advance_ - LongIndex(ThreadMap::Iterations::kContiguous - 1) * inc_contiguous_ - LongIndex(ThreadMap::Iterations::kStrided - 1) * inc_strided_; + }; + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + // + // Data members + // + + /// Parameters object with precomputed internal state + Params params_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + + UnderlyingPredicates the_predicates; + + /// Used for out-of-order visitation + bool is_residue_tile_; + + private: + /// Computes predicates based on internally tracked per-thread offset. + MCTLASS_DEVICE + void compute_predicates_( + /// Extent of the matrix window + TensorCoord extent, + /// optionally, simplify predicate calculation during 'steady state' phase + bool is_steady_state = false) { + the_predicates.compute_predicates_(extent, is_steady_state); + } + + public: + + /// Default constructor + PredicatedTileAccessIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + ///< Precomputed parameters object + Params const ¶ms, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset, + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ) + : params_(params), + pointer_(reinterpret_cast( + const_cast(pointer))), + the_predicates(extent), + is_residue_tile_(true) { + + the_predicates.set_predicates(thread_id, threadblock_offset); + + // update internal pointers + Layout layout(params_.stride_); + add_pointer_offset(layout(the_predicates.thread_offset_)); + } + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { the_predicates.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += sizeof_bits::value * pointer_offset / 8; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + if (is_residue_tile_) { + + the_predicates.thread_offset_ += the_predicates.residue_offset_; + + Layout layout(params_.stride_); + add_pointer_offset(layout(the_predicates.residue_offset_)); + + the_predicates.compute_predicates_(the_predicates.extent_, true); + + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset[1] - 1); + pointer_ += Shape::kContiguous * tile_offset[0]; + } else { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset[0] - 1); + pointer_ += Shape::kStrided * tile_offset[1]; + } + } else { + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset[1]); + pointer_ += Shape::kContiguous * tile_offset[0]; + } else { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset[0]); + pointer_ += Shape::kStrided * tile_offset[1]; + } + } + is_residue_tile_ = false; + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(pointer_) + the_predicates.iteration_vector_; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator &operator++() { + the_predicates.operator++(); + ++the_predicates.iteration_vector_; + if (the_predicates.iteration_vector_ < kAccessesPerVector) { + return *this; + } + + the_predicates.iteration_vector_ = 0; + ++the_predicates.iteration_contiguous_; + + if (the_predicates.iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + pointer_ += params_.inc_contiguous_; + return *this; + } + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + the_predicates.iteration_contiguous_ = 0; + ++the_predicates.iteration_strided_; + + if (the_predicates.iteration_strided_ < ThreadMap::Iterations::kStrided) { + pointer_ += params_.inc_next_strided_; + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + the_predicates.iteration_strided_ = 0; + + // advance to next tile + pointer_ += params_.inc_next_; + + // now return to start tile - if the iterator is subsequently advanced, this + // subtraction as well as the subsequent integer addition are both elided by + // the compiler. + pointer_ -= params_.inc_advance_; + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator operator++(int) { + PredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { the_predicates.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { the_predicates.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { the_predicates.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { the_predicates.get_mask(mask); } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return the_predicates.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for affine rank 2 column-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = PredicatedTileAccessIterator< + layout::PitchLinearShape, Element, + layout::AffineRankN<2>, (kAdvanceRank == 0 ? 0 : 1), ThreadMap, AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default constructor + Params() = default; + + /// Construct the Params object given an AffineRankN<2> tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::AffineRankN<2>(layout.stride(0), layout.stride(1))){}; + }; + + private: + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + + public: + + /// Default constructor + PredicatedTileAccessIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + ///< Precomputed parameters object + Params const ¶ms, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset, + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), + threadblock_offset.column())) {} + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset(make_Coord(tile_offset.row(), tile_offset.column())); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator operator++(int) { + PredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for affine rank-2 row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = PredicatedTileAccessIterator< + layout::PitchLinearShape, Element, + layout::AffineRankN<2>, (kAdvanceRank == 0 ? 1 : 0), ThreadMap, AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default constructor + Params() = default; + + /// Construct the Params object given an AffineRankN<2> tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::AffineRankN<2>(layout.stride(1), layout.stride(0))){}; + }; + + private: + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + + public: + + /// Default constructor + PredicatedTileAccessIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + ///< Precomputed parameters object + Params const ¶ms, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset, + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset(make_Coord(tile_offset.column(), tile_offset.row())); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator operator++(int) { + PredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for column-major interleaved data. +/// It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// + +template +class PredicatedTileAccessIterator, + AdvanceRank, ThreadMap_, AccessType_, false, + layout::NoPermute> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::ColumnMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileAccessIterator< + layout::PitchLinearShape, + Element, layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap, + AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + + MCTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const &base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + + /// Default constructor + PredicatedTileAccessIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset, + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.row() * kInterleavedK, + extent.column() / kInterleavedK), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.row() * kInterleavedK, + threadblock_offset.column() / kInterleavedK)) {} + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator operator++(int) { + PredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { return iterator_.valid(); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for row-major interleaved data. +// It is mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIterator, + AdvanceRank, ThreadMap_, AccessType_, false, + layout::NoPermute> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::RowMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileAccessIterator< + layout::PitchLinearShape, + Element, layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap, + AccessType>; + + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + /// Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + + MCTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const &base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + + /// Default constructor + PredicatedTileAccessIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset, + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.column() * kInterleavedK, + extent.row() / kInterleavedK), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column() * kInterleavedK, + threadblock_offset.row() / kInterleavedK)) {} + + /// Construct a PredicatedTileAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator operator++(int) { + PredicatedTileAccessIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { return iterator_.valid(); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator_2dthreadtile.h b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator_2dthreadtile.h new file mode 100644 index 0000000..e3d678a --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator_2dthreadtile.h @@ -0,0 +1,835 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates calculating the address and predicates to the load of tiles + from pitch-linear rank=2 tensors. + + This iterator uses masks to guard out-of-bounds accesses and visits the last + "residue" tile first, with the objective of minimizing predicate mask updates + during steady-state operation. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/mctlass.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" +#include "mctlass/transform/threadblock/predicated_tile_access_iterator_params.h" + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileAccessIterator2dThreadTile +/// +template +class PredicatedTileAccessIterator2dThreadTile; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator2dThreadTile for pitch-linear data. +/// +template +class PredicatedTileAccessIterator2dThreadTile { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + static int const kPredicatesPerByte = 4; + static int const kPredicatesPerWord = 4 * kPredicatesPerByte; + + /// Number of 32b words containing predicates + static int const kPredicateByteCount = (ThreadMap::Iterations::kCount * ThreadMap::ThreadAccessShape::kStrided + kPredicatesPerByte - 1) / kPredicatesPerByte; + static int const kPredicateWordCount = (kPredicateByteCount + 3) / 4; + + static unsigned const kPredicateMask = (1u << kPredicatesPerByte) - 1u; + + static_assert(kPredicateWordCount <= 4, "Too many predicates."); + + /// Predicate vector stores mask to guard accesses + using Mask = Array; + + /// Uses a non-template class + struct Params : PredicatedTileAccessIteratorParams { + + public: + friend PredicatedTileAccessIterator2dThreadTile; + + using Base = PredicatedTileAccessIteratorParams; + + // Default ctor + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) : + Base(layout.stride(0), + MakePredicatedTileAccessIteratorDesc()() + ) { } + + MCTLASS_HOST_DEVICE + Params(Base const &base) : + Base(base) { } + }; + + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + /// Parameters object with precomputed internal state + //Params const ¶ms_; + Params const params_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + + /// Guard predicates + uint32_t predicates_[kPredicateWordCount]; + + /// Size of tensor + TensorCoord extent_; + + /// Initial offset for each thread + TensorCoord thread_offset_; + + /// Index of residue tile + int residue_tile_idx_; + + /// Used for out-of-order visitation + bool is_residue_tile_; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + /// Tracks iterations within the thread loop + int iteration_thread_; + + private: + /// Computes predicates based on internally tracked per-thread offset. + MCTLASS_HOST_DEVICE + void compute_predicates_( + /// optionally, simplify predicate calculation during 'steady state' phase + bool is_steady_state = false) { + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = 0u; + } + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + MCTLASS_PRAGMA_UNROLL + for (int ts = 0; ts < ThreadMap::ThreadAccessShape::kStrided; ts++) { + + TensorCoord iteration_coord(c * ThreadMap::Delta::kContiguous, + ts + s * ThreadMap::Delta::kStrided); + + TensorCoord coord = thread_offset_ + iteration_coord; + + bool guard; + + if (is_steady_state) { + if (kAdvanceRank == 0) { + guard = (coord.strided() < extent_.strided()); + } else { + guard = (coord.contiguous() < extent_.contiguous()); + } + } else { + guard = (coord.strided() < extent_.strided() && + coord.contiguous() < extent_.contiguous()); + } + + int pred_idx = ts + c * ThreadMap::ThreadAccessShape::kStrided + s * ThreadMap::Iterations::kContiguous * ThreadMap::ThreadAccessShape::kStrided; + int word_idx = pred_idx / kPredicatesPerWord; + int residual = pred_idx % kPredicatesPerWord; + int byte_idx = residual / kPredicatesPerByte; + int bit_idx = residual % kPredicatesPerByte; + + predicates_[word_idx] |= (unsigned(guard) << (byte_idx * 8 + bit_idx)); + + } + } + } + + } + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator2dThreadTile( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : params_(params), + pointer_(reinterpret_cast( + const_cast(pointer))), + extent_(extent), + is_residue_tile_(true) { + + + TensorCoord residue_offset; + if (kAdvanceRank) { + residue_tile_idx_ = + (extent_[kAdvanceRank] - threadblock_offset[kAdvanceRank] - 1) / + Shape::kStrided; + residue_offset = make_Coord(0, residue_tile_idx_ * Shape::kStrided); + } else { + residue_tile_idx_ = + (extent_[kAdvanceRank] - threadblock_offset[kAdvanceRank] - 1) / + Shape::kContiguous; + residue_offset = make_Coord(residue_tile_idx_ * Shape::kContiguous, 0); + } + + // Per-thread offset in logical coordinates of tensor + thread_offset_ = threadblock_offset + residue_offset + + ThreadMap::initial_offset(thread_id); + + // update internal pointers + Layout layout(params_.stride_); + add_pointer_offset(layout(thread_offset_)); + + compute_predicates_(false); + + set_iteration_index(0); + } + + /// Construct a PredicatedTileAccessIterator2dThreadTile with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator2dThreadTile( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id) + : PredicatedTileAccessIterator2dThreadTile(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + + int residual = index % (ThreadMap::Iterations::kContiguous * ThreadMap::ThreadAccessShape::kStrided); + iteration_strided_ = index / (ThreadMap::Iterations::kContiguous * ThreadMap::ThreadAccessShape::kStrided); + + iteration_contiguous_ = residual / ThreadMap::ThreadAccessShape::kStrided; + iteration_thread_ = residual % ThreadMap::ThreadAccessShape::kStrided; + + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += int(sizeof(Element)) * pointer_offset; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_DEVICE + void add_tile_offset( + TensorCoord const &tile_offset) { + if (is_residue_tile_) { + TensorCoord residue_offset; + if (kAdvanceRank) { + residue_offset = TensorCoord(0, residue_tile_idx_ * Shape::kStrided); + } else { + residue_offset = TensorCoord(residue_tile_idx_ * Shape::kContiguous, 0); + } + + thread_offset_ -= residue_offset; + + Layout layout(params_.stride_); + add_pointer_offset(-layout(residue_offset)); + + compute_predicates_(true); + + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * (tile_offset.strided() - 1); + pointer_ += Shape::kContiguous * tile_offset.contiguous(); + } else { + pointer_ += params_.inc_advance_ * (tile_offset.contiguous() - 1); + pointer_ += Shape::kStrided * tile_offset.strided(); + } + } else { + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * tile_offset.strided(); + pointer_ += Shape::kContiguous * tile_offset.contiguous(); + } else { + pointer_ += params_.inc_advance_ * tile_offset.contiguous(); + pointer_ += Shape::kStrided * tile_offset.strided(); + } + } + is_residue_tile_ = false; + } + + MCTLASS_HOST_DEVICE + AccessType *get() const { + + AccessType *ret_val = reinterpret_cast( + pointer_ + (iteration_thread_ * params_.stride_ + iteration_contiguous_ * ThreadMap::Delta::kContiguous) * int(sizeof(Element))); + + return ret_val; + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator2dThreadTile &operator++() { + + iteration_thread_++; + + if (iteration_thread_ < ThreadMap::ThreadAccessShape::kStrided) + return *this; + + iteration_thread_ = 0; + + ++iteration_contiguous_; + + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) + return *this; + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + iteration_contiguous_ = 0; + ++iteration_strided_; + + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + pointer_ += params_.inc_strided_; + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + iteration_strided_ = 0; + + // advance to next tile + pointer_ += params_.inc_next_; + + // now return to start tile - if the iterator is subsequently advanced, this + // subtraction as well as the subsequent integer addition are both elided by + // the compiler. + pointer_ -= params_.inc_advance_; + + return *this; + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator2dThreadTile operator++(int) { + PredicatedTileAccessIterator2dThreadTile self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = enable ? 0u : predicates_[i]; + } + + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = 0xffffffff; + } + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = mask[i]; + } + + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + mask[i] = predicates_[i]; + } + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + + int pred_idx = + iteration_thread_ + + iteration_contiguous_ * ThreadMap::ThreadAccessShape::kStrided + + iteration_strided_ * ThreadMap::Iterations::kContiguous * ThreadMap::ThreadAccessShape::kStrided; + + int word_idx = pred_idx / kPredicatesPerWord; + int residual = pred_idx % kPredicatesPerWord; + int byte_idx = residual / kPredicatesPerByte; + int bit_idx = residual % kPredicatesPerByte; + + bool pred = (predicates_[word_idx] & (1u << (byte_idx * 8 + bit_idx))) != 0; + + return pred; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator2dThreadTile for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIterator2dThreadTile { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileAccessIterator2dThreadTile< + layout::PitchLinearShape, Element, + layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap, AccessType>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIterator2dThreadTile; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default ctor + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))){} + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const &base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator2dThreadTile( + ///< Precomputed parameters object + Params const ¶ms, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), + threadblock_offset.column())) {} + + /// Construct a PredicatedTileAccessIterator2dThreadTile with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator2dThreadTile( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIterator2dThreadTile(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator2dThreadTile &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator2dThreadTile operator++(int) { + PredicatedTileAccessIterator2dThreadTile self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator2dThreadTile for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIterator2dThreadTile { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileAccessIterator2dThreadTile< + layout::PitchLinearShape, Element, + layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap, AccessType>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIterator2dThreadTile; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default ctor + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))){} + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const &base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator2dThreadTile( + ///< Precomputed parameters object + Params const ¶ms, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedTileAccessIterator2dThreadTile with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator2dThreadTile( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIterator2dThreadTile(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator2dThreadTile &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIterator2dThreadTile operator++(int) { + PredicatedTileAccessIterator2dThreadTile self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator_params.h b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator_params.h new file mode 100644 index 0000000..7be06e9 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator_params.h @@ -0,0 +1,289 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/layout/pitch_linear.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Predicated tile access iterator descriptor object containing template dependent state +struct PredicatedTileAccessIteratorDesc { + + int element_size_bits; + int advance_rank; + layout::PitchLinearCoord threadblock_shape; + layout::PitchLinearCoord threadmap_iterations; + layout::PitchLinearCoord threadmap_delta; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorDesc() { } + + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorDesc( + int element_size_bits_, + int advance_rank_, + layout::PitchLinearCoord threadblock_shape_, + layout::PitchLinearCoord threadmap_iterations_, + layout::PitchLinearCoord threadmap_delta_ + ): + element_size_bits(element_size_bits_), + advance_rank(advance_rank_), + threadblock_shape(threadblock_shape_), + threadmap_iterations(threadmap_iterations_), + threadmap_delta(threadmap_delta_) + { + #if 0 + printf("PredicatedTileAccessIteratorDesc(%d, %d, {%d, %d}, {%d, %d}, {%d, %d}})\n", + element_size_bits, + advance_rank, + threadblock_shape.contiguous(), threadblock_shape.strided(), + threadmap_iterations.contiguous(), threadmap_iterations.strided(), + threadmap_delta.contiguous(), threadmap_delta.strided()); + #endif + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +/// Helper template to construct an PredicatedTileAccessIteratorDesc from a template +// dependent state +template < + typename Shape, typename Element, typename Layout, + int AdvanceRank, typename ThreadMap> + struct MakePredicatedTileAccessIteratorDesc; +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for pitch-linear data. +template < + typename Shape, typename Element, int AdvanceRank, + typename ThreadMap> +struct MakePredicatedTileAccessIteratorDesc < + Shape, Element, layout::PitchLinear, AdvanceRank, ThreadMap> { + + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorDesc operator()() { + + return PredicatedTileAccessIteratorDesc( + sizeof_bits::value, + AdvanceRank, + {Shape::kContiguous, Shape::kStrided}, + {ThreadMap::Iterations::kContiguous, ThreadMap::Iterations::kStrided}, + {ThreadMap::Delta::kContiguous, ThreadMap::Delta::kStrided} + ); +} + +}; +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for column-major data. +template < + typename Shape, typename Element, int AdvanceRank, + typename ThreadMap> +struct MakePredicatedTileAccessIteratorDesc < + Shape, Element, layout::ColumnMajor, AdvanceRank, ThreadMap> { + + static int const kAdvanceRank = AdvanceRank; + + using UnderlyingMakeOperator = MakePredicatedTileAccessIteratorDesc< + layout::PitchLinearShape, Element, + layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap>; + + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorDesc operator()() { + + return UnderlyingMakeOperator()(); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for row-major data. +template < + typename Shape, typename Element, int AdvanceRank, + typename ThreadMap> +struct MakePredicatedTileAccessIteratorDesc < + Shape, Element, layout::RowMajor, AdvanceRank, ThreadMap> { + + static int const kAdvanceRank = AdvanceRank; + + using UnderlyingMakeOperator = MakePredicatedTileAccessIteratorDesc< + layout::PitchLinearShape, Element, + layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap>; + + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorDesc operator()() { + + return UnderlyingMakeOperator()(); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for column-major interleaved data. +template < + typename Shape, typename Element, int AdvanceRank, + typename ThreadMap, int InterleavedK> +struct MakePredicatedTileAccessIteratorDesc < + Shape, Element, layout::ColumnMajorInterleaved, AdvanceRank, ThreadMap> { + + static int const kAdvanceRank = AdvanceRank; + static int const kInterleavedK = InterleavedK; + + using UnderlyingMakeOperator = MakePredicatedTileAccessIteratorDesc< + layout::PitchLinearShape, Element, + layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap>; + + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorDesc operator()() { + + return UnderlyingMakeOperator()(); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIterator for roww-major interleaved data. +template < + typename Shape, typename Element, int AdvanceRank, + typename ThreadMap, int InterleavedK> +struct MakePredicatedTileAccessIteratorDesc < + Shape, Element, layout::RowMajorInterleaved, AdvanceRank, ThreadMap> { + + static int const kAdvanceRank = AdvanceRank; + static int const kInterleavedK = InterleavedK; + + using UnderlyingMakeOperator = MakePredicatedTileAccessIteratorDesc< + layout::PitchLinearShape, Element, + layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap>; + + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorDesc operator()() { + + return UnderlyingMakeOperator()(); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// +// Parameters struct +// + +struct PredicatedTileAccessIteratorParams { + + using Index = int32_t; + using LongIndex = int64_t; + + // + // Data members + // + /// stride of pitch-linear layout (units of Element) + LongIndex stride_; + /// amount (in byte) to increment pointer to move to next access along + /// strided dimension + LongIndex inc_strided_; + /// amount (in byte) to increment pointer from last access to first access + /// of next tile + LongIndex inc_next_; + /// amount (in byte) to increment pointer from first access of current tile + /// to first access of next tile + LongIndex inc_advance_; + + // + // Methods + // + + MCTLASS_HOST_DEVICE + Status initialize(LongIndex stride, PredicatedTileAccessIteratorDesc desc) { + + stride_ = stride; + + inc_strided_ = (LongIndex(stride_) * desc.threadmap_delta.strided()) * + desc.element_size_bits / 8; + + if (desc.advance_rank) { + // advance along strided dimension + inc_advance_ = + desc.threadblock_shape.strided() * LongIndex(stride_) * desc.element_size_bits / 8; + } else { + // advance along contiguous dimension + inc_advance_ = desc.threadblock_shape.contiguous() * desc.element_size_bits / 8; + } + + inc_next_ = inc_advance_ - LongIndex(desc.threadmap_iterations.strided() - 1) * + desc.threadmap_delta.strided() * LongIndex(stride_) * + desc.element_size_bits / 8; + + return Status::kSuccess; + } + + MCTLASS_HOST_DEVICE + Status initialize(Index stride, PredicatedTileAccessIteratorDesc desc) { + return initialize(LongIndex(stride), desc); + } + + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorParams() { + initialize(LongIndex(0), PredicatedTileAccessIteratorDesc()); + } + + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorParams(Index stride, PredicatedTileAccessIteratorDesc desc) { + initialize(stride, desc); + } + + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorParams(LongIndex stride, PredicatedTileAccessIteratorDesc desc) { + initialize(stride, desc); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator_triangular_matrix.h b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator_triangular_matrix.h new file mode 100644 index 0000000..4e16d71 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_access_iterator_triangular_matrix.h @@ -0,0 +1,892 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates calculating the address and predicates to the load of tiles + from pitch-linear rank=2 tensors. + + This iterator uses masks to guard out-of-bounds accesses and visits the last + "residue" tile first, with the objective of minimizing predicate mask updates + during steady-state operation. + + A precomputed "Params" object minimizes the amount of state that must be + stored in registers, and integer addition is used to advance the pointer + through memory. + + +*/ + +#pragma once + +#include "mctlass/blas3.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/predicate_vector.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/tensor_view.h" + +//////////////////////////////////////////////////////////////////////////////// + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileAccessIteratorTriangularMatrix +/// +template +class PredicatedTileAccessIteratorTriangularMatrix; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorTriangularMatrix for pitch-linear data. +/// +template +class PredicatedTileAccessIteratorTriangularMatrix { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + static int const kAccessesPerVector = ThreadMap::kElementsPerAccess / AccessType::kElements; + + using CompareOp = typename TrMatrixCompareOp::Type; + + static_assert( kFillMode == FillMode::kFull || + ((kFillMode == FillMode::kLower || kFillMode == FillMode::kUpper) && AccessType::kElements == 1), + "BLAS3 iterator for the triangular/symmetric matrix must use AccessType::kElements as 1"); + + static_assert(!(ThreadMap::kElementsPerAccess % AccessType::kElements), + "Vectors implied by the thread map must be divisible by the access type."); + + static int const kPredicatesPerByte = 4; + static int const kPredicatesPerWord = 4 * kPredicatesPerByte; + + static int const kPredicateCount = ThreadMap::Iterations::kCount * kAccessesPerVector; + + /// Number of 32b words containing predicates + static int const kPredicateByteCount = + (kPredicateCount + kPredicatesPerByte - 1) / kPredicatesPerByte; + static int const kPredicateWordCount = (kPredicateByteCount + 3) / 4; + + static unsigned const kPredicateMask = (1u << kPredicatesPerByte) - 1u; + + static_assert(kPredicateWordCount <= 4, "Too many predicates."); + + /// Predicate vector stores mask to guard accesses + using Mask = Array; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + friend PredicatedTileAccessIteratorTriangularMatrix; + + private: + /// stride of pitch-linear layout (units of Element) + StrideIndex stride_; + /// (true) pitch-linear layout is mapped to row-major matrix + /// (false) pitch-linear layout is mapped to column-major matrix + bool is_row_major_; + /// for vectorized access across the diagonal boundary guard condition is + /// checked for the element on the boundary + int access_diagonal_boundary_; + /// amount (in byte) to increment pointer to move to next access along + /// strided dimension + LongIndex inc_strided_; + /// amount (in byte) to increment pointer from last access to first access + /// of next tile + LongIndex inc_next_; + /// amount (in byte) to increment pointer from first access of current tile + /// to first access of next tile + LongIndex inc_advance_; + + public: + + // Default ctor + MCTLASS_HOST_DEVICE + Params(): stride_(0), inc_strided_(0), inc_next_(0), inc_advance_(0), is_row_major_(false), access_diagonal_boundary_(0) { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout, bool is_row_major, int access_diagonal_boundary) : + stride_(layout.stride(0)), is_row_major_(is_row_major), access_diagonal_boundary_(access_diagonal_boundary) { + + inc_strided_ = (LongIndex(stride_) * ThreadMap::Delta::kStrided) * + sizeof_bits::value / 8; + + if (kAdvanceRank) { + // advance along strided dimension + inc_advance_ = + Shape::kStrided * LongIndex(stride_) * sizeof_bits::value / 8; + } else { + // advance along contiguous dimension + inc_advance_ = Shape::kContiguous * sizeof_bits::value / 8; + } + + inc_next_ = inc_advance_ - LongIndex(ThreadMap::Iterations::kStrided - 1) * + ThreadMap::Delta::kStrided * LongIndex(stride_) * + sizeof_bits::value / 8; + + }; + + + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + /// Parameters object with precomputed internal state + Params const ¶ms_; + + /// Internal pointer to first access of tile + BytePointer pointer_; + + /// Guard predicates + uint32_t predicates_[kPredicateWordCount]; + + /// Track global memory addresses on the diagonal + /// To ignore imag part for diagonal elements of hermitian matrices + uint32_t predicates_onDiag_[kPredicateWordCount]; + + /// Size of tensor + TensorCoord extent_; + + /// Initial offset for each thread + TensorCoord thread_offset_; + + /// Iteration along vectors implied by the thread map + int iteration_vector_; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + private: + /// Computes predicates based on internally tracked per-thread offset. + MCTLASS_DEVICE + void compute_predicates_( + /// Extent of the matrix window + TensorCoord extent) { + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = 0u; + predicates_onDiag_[i] = 0u; + } + + CompareOp compare_op; + + MCTLASS_PRAGMA_UNROLL + for (int access_idx = 0; access_idx < ThreadMap::Iterations::kCount * kAccessesPerVector; ++access_idx) { + + int s = access_idx / (ThreadMap::Iterations::kContiguous * kAccessesPerVector); + + int access_residual = access_idx % (ThreadMap::Iterations::kContiguous * kAccessesPerVector); + + int c = access_residual / kAccessesPerVector; + int v = access_residual % kAccessesPerVector; + + TensorCoord iteration_coord(c * ThreadMap::Delta::kContiguous + v * AccessType::kElements, + s * ThreadMap::Delta::kStrided); + + TensorCoord coord = thread_offset_ + iteration_coord; + + bool guard; + bool onDiag = false; + + guard = ((coord.strided() < extent.strided()) && + (coord.contiguous() < extent.contiguous())); + + + // guard access on the wrong side of the triagular matrix diagonal + if (kFillMode == FillMode::kLower || kFillMode == FillMode::kUpper) { + coord += TensorCoord{params_.access_diagonal_boundary_, 0}; + + bool triagular_guard_row_major = compare_op(coord.strided(), coord.contiguous()) | !params_.is_row_major_; + bool triagular_guard_col_major = compare_op(coord.contiguous(), coord.strided()) | params_.is_row_major_; + + guard = guard && triagular_guard_row_major && triagular_guard_col_major; + + if (kDiagType == DiagType::kUnit) { + onDiag = (guard && coord.strided() == coord.contiguous()) ? true : false; + } + } + + int pred_idx_onDiag = v + kAccessesPerVector * (c + ThreadMap::Iterations::kContiguous * s); + int word_idx_onDiag = pred_idx_onDiag / kPredicatesPerWord; + int residual_onDiag = pred_idx_onDiag % kPredicatesPerWord; + int byte_idx_onDiag = residual_onDiag / kPredicatesPerByte; + int bit_idx_onDiag = residual_onDiag % kPredicatesPerByte; + + predicates_onDiag_[word_idx_onDiag] |= (unsigned(onDiag) << (byte_idx_onDiag * 8 + bit_idx_onDiag)); + + int pred_idx = v + kAccessesPerVector * (c + ThreadMap::Iterations::kContiguous * s); + + int word_idx = pred_idx / kPredicatesPerWord; + int residual = pred_idx % kPredicatesPerWord; + int byte_idx = residual / kPredicatesPerByte; + int bit_idx = residual % kPredicatesPerByte; + + predicates_[word_idx] |= (unsigned(guard) << (byte_idx * 8 + bit_idx)); + + } + + } + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorTriangularMatrix( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : params_(params), + pointer_(reinterpret_cast(const_cast(pointer))), + extent_(extent) { + + + // Per-thread offset in logical coordinates of tensor + thread_offset_ = threadblock_offset + ThreadMap::initial_offset(thread_id); + + // update internal pointers + Layout layout(params_.stride_); + add_pointer_offset(layout(thread_offset_)); + + compute_predicates_(extent_); + + set_iteration_index(0); + } + + /// Construct a PredicatedTileAccessIteratorTriangularMatrix with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorTriangularMatrix( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id) + : PredicatedTileAccessIteratorTriangularMatrix(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + + iteration_vector_ = index % kAccessesPerVector; + int residual_access = index / kAccessesPerVector; + + iteration_contiguous_ = residual_access % ThreadMap::Iterations::kContiguous; + iteration_strided_ = residual_access / ThreadMap::Iterations::kContiguous; + + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += sizeof_bits::value * pointer_offset / 8; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + + if (kAdvanceRank) { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.strided()); + pointer_ += Shape::kContiguous * tile_offset.contiguous(); + thread_offset_ += TensorCoord{0, Shape::kStrided * tile_offset.strided()}; + } else { + pointer_ += params_.inc_advance_ * LongIndex(tile_offset.contiguous()); + pointer_ += Shape::kStrided * tile_offset.strided(); + thread_offset_ += TensorCoord{Shape::kContiguous * tile_offset.contiguous(), 0}; + } + + compute_predicates_(extent_); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast( + pointer_ + + iteration_contiguous_ * (ThreadMap::Delta::kContiguous * sizeof_bits::value) / 8) + iteration_vector_; + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorTriangularMatrix &operator++() { + + ++iteration_vector_; + if (iteration_vector_ < kAccessesPerVector) { + return *this; + } + + iteration_vector_ = 0; + ++iteration_contiguous_; + + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) { + return *this; + } + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + iteration_contiguous_ = 0; + ++iteration_strided_; + + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + pointer_ += params_.inc_strided_; + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + iteration_strided_ = 0; + + // advance to next tile + pointer_ += params_.inc_next_; + + // now return to start tile - if the iterator is subsequently advanced, this + // subtraction as well as the subsequent integer addition are both elided by + // the compiler. + pointer_ -= params_.inc_advance_; + + return *this; + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorTriangularMatrix operator++(int) { + PredicatedTileAccessIteratorTriangularMatrix self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = enable ? 0u : predicates_[i]; + } + + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = 0xffffffff; + } + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + predicates_[i] = mask[i]; + } + + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kPredicateWordCount; ++i) { + mask[i] = predicates_[i]; + } + } + + /// Return if the address in on the diagonal + MCTLASS_HOST_DEVICE + bool getOnDiag() { + int pred_idx = + iteration_vector_ + kAccessesPerVector * (iteration_contiguous_ + iteration_strided_ * ThreadMap::Iterations::kContiguous); + + int word_idx = pred_idx / kPredicatesPerWord; + int residual = pred_idx % kPredicatesPerWord; + int byte_idx = residual / kPredicatesPerByte; + int bit_idx = residual % kPredicatesPerByte; + + bool pred = (predicates_onDiag_[word_idx] & (1u << (byte_idx * 8 + bit_idx))) != 0; + return pred; + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + + + int pred_idx = + iteration_vector_ + kAccessesPerVector * (iteration_contiguous_ + iteration_strided_ * ThreadMap::Iterations::kContiguous); + + int word_idx = pred_idx / kPredicatesPerWord; + int residual = pred_idx % kPredicatesPerWord; + int byte_idx = residual / kPredicatesPerByte; + int bit_idx = residual % kPredicatesPerByte; + + bool pred = (predicates_[word_idx] & (1u << (byte_idx * 8 + bit_idx))) != 0; + return pred; + + + //return true; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorTriangularMatrix for column-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIteratorTriangularMatrix { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileAccessIteratorTriangularMatrix< + layout::PitchLinearShape, Element, + layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap, + kSideMode, kFillMode, kDiagType, AccessType>; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + static int const kAccessDiagonalBoundary = + (kFillMode == FillMode::kLower) ? (AccessType::kElements - 1) : 0; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorTriangularMatrix; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default ctor + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0)), false, kAccessDiagonalBoundary){}; + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorTriangularMatrix( + ///< Precomputed parameters object + Params const ¶ms, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), + threadblock_offset.column())) {} + + /// Construct a PredicatedTileAccessIteratorTriangularMatrix with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorTriangularMatrix( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorTriangularMatrix(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.row(), tile_offset.column()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorTriangularMatrix &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorTriangularMatrix operator++(int) { + PredicatedTileAccessIteratorTriangularMatrix self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// Return if the address in on the diagonal + MCTLASS_HOST_DEVICE + bool getOnDiag() { + return iterator_.getOnDiag(); + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileAccessIteratorTriangularMatrix for row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileAccessIteratorTriangularMatrix { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + using AccessType = AccessType_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileAccessIteratorTriangularMatrix< + layout::PitchLinearShape, Element, + layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap, + kSideMode, kFillMode, kDiagType, AccessType>; + + static int const kAccessesPerVector = UnderlyingIterator::kAccessesPerVector; + + static int const kAccessDiagonalBoundary = + (kFillMode == FillMode::kUpper) ? (AccessType::kElements - 1) : 0; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileAccessIteratorTriangularMatrix; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default ctor + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0)), true, kAccessDiagonalBoundary){}; + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorTriangularMatrix( + ///< Precomputed parameters object + Params const ¶ms, + ///< Pointer to start of tensor + Pointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedTileAccessIteratorTriangularMatrix with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorTriangularMatrix( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileAccessIteratorTriangularMatrix(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorTriangularMatrix &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileAccessIteratorTriangularMatrix operator++(int) { + PredicatedTileAccessIteratorTriangularMatrix self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// Return if the address in on the diagonal + MCTLASS_HOST_DEVICE + bool getOnDiag() { + return iterator_.getOnDiag(); + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_iterator.h b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_iterator.h new file mode 100644 index 0000000..87440af --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_iterator.h @@ -0,0 +1,1888 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of tiles from pitch-linear rank=2 tensors. + + This iterator uses masks to guard out-of-bounds accesses. The first tile this + iterator visits maybe partial, then the remaining tiles are complete. So, we + only need to compute the predicates twice, once before the first tile and + once for the remaining full tiles which can share the same predicates. + + A precomputed "Params" object minimizes the amount of state that must be stored in registers, + and integer addition is used to advance the pointer through memory. +*/ + +#pragma once + +#include "mctlass/arch/memory.h" +#include "mctlass/transform/threadblock/predicated_tile_access_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileIterator +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +/// Regular tile iterator using a precomputed control structure to minimize register liveness +/// and integer arithmetic. +/// +/// Layout is assumed to be invariant at the time the precomputed "Params" object is constructed. +/// +/// Base pointer and tensor extents may be specified at the time the iterator is constructed. +/// Subsequently, they are assumed to be immutable. +/// +/// Adding a logical coordinate offset may be performed at the time the iterator is constructed. +/// Subsequent additions to logical coordinate offset may be performed but are relatively expensive. +/// +/// Visitation order is intended to first visit a "residual" tile that may be partially full in +/// both the advance dimension and the steady-state dimension. This is assumed to be the last +/// tile in the iteration sequence. Advancing an iterator that has just been constructed moves to +/// the first tile that is full in the advance dimension and recomputes predicates. Subsequent +/// accesses may be performed without updating internal predicates and are efficient in terms of +/// live register state and pointer arithmetic instructions. +/// +/// To be efficient, this assumes the iterator will be dereferenced and advanced at least once +/// outside any looping structure to minimize integer arithmetic. +/// +/// Acceses out of bounds are safe so long as `clear_mask()` is called prior to dereferencing +/// the iterator. +/// +/// +/// Example: +/// +/// An efficient pipeline structure may be constructed as follows: +/// +// template +// __global__ void kernel( +// typename Iterator::Params params, +// typename Iterator::Element *ptr, +// TensorCoord extent) { +// +// typename Iterator::Fragment fragment; +// +// TensorCoord threadblock_offset(0, 0); +// +// Iterator iter(params, ptr, extent, threadIdx.x, threadblock_offsets); +// +// +// fragment = *iter; // load "residue" tile first +// ++iter; // advance to first "steady state" tile and update internal masks +// +// +// #pragma unroll +// for (int i = Remaining - 1; i >= 0; --i) { +// +// f(fragment); +// +// if (!i) { +// iter.clear_mask(); // light-weight operation to clear masks - subsequent loads become NO-OPs. +// } +// +// fragment = *iter; // load tile during "steady state" phase +// ++iter; // advance to next tile - lightweight due to steady-state masks +// } +// } +// +// void host(TensorView view) { +// +// using Iterator = transform::threadblock::PredicatedTileIterator; +// +// typename Iterator::Params params(view.layout()); +// +// kernel(params, view.data()); +// } +/// +/// +template < + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + int AccessSize = ThreadMap::kElementsPerAccess, + bool Gather = false, + typename PermuteLayout = layout::NoPermute +> +class PredicatedTileIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + /// Type used for internal memory accesses + using AccessType = AlignedArray::value / 8)>; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = + PredicatedTileAccessIterator; + + static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename TileAccessIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + using Base = typename TileAccessIterator::Params::Base; + + friend PredicatedTileIterator; + + private: + /// Parameters object + typename TileAccessIterator::Params params_; + + public: + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) : params_(layout) { } + + // Default constructor + Params() = default; + + MCTLASS_HOST_DEVICE + Params(Base const &base) + : params_(base) {} + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + + public: + + /// Default constructor + PredicatedTileIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset, + /// Gather indices + int const *indices = nullptr) + : address_iterator_(params.params_, pointer, extent, thread_id, + threadblock_offset, indices) {} + + /// Construct a PredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator &operator++() { + if (kAdvanceRank) + address_iterator_.add_tile_offset({0, 1}); + else + address_iterator_.add_tile_offset({1, 0}); + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator operator++(int) { + PredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { address_iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { address_iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { address_iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { address_iterator_.get_mask(mask); } + + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + load_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + address_iterator_.set_iteration_index(idx); + char const *byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + + AccessType const *access_ptr = reinterpret_cast(byte_ptr); + + mctlass::arch::global_load( + frag_ptr[idx], access_ptr, address_iterator_.valid()); + + ++address_iterator_; + } + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_byte_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + store_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) { + address_iterator_.set_iteration_index(0); + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + char *byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType *access_ptr = reinterpret_cast(byte_ptr); + + if (address_iterator_.valid()) { + *access_ptr = frag_ptr[idx]; + } + ++address_iterator_; + } + } + } + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_byte_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize, + bool Gather, + typename PermuteLayout +> +class PredicatedTileIterator { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileIterator< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessSize, + Gather, + PermuteLayout + >; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + + friend PredicatedTileIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + // Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) + {} + + MCTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const &base) + : params_(base) {} + }; + + +private: + + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + + /// Default constructor + PredicatedTileIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const &threadblock_offset, ///< Initial offset of threadblock + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ): + iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column()), + indices) + { } + + /// Construct a PredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ): PredicatedTileIterator(params, pointer, extent, thread_id, make_Coord(0, 0)) { } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator operator++(int) { + PredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize, + bool Gather, + typename PermuteLayout +> +class PredicatedTileIterator { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileIterator< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize, + Gather, + PermuteLayout + >; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + + friend PredicatedTileIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) {} + + MCTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const &base) + : params_(base) {} + + }; + +private: + + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + + /// Default constructor + PredicatedTileIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const &threadblock_offset, ///< Initial offset of threadblock + int const *indices = nullptr ///< Gather indices + ): + iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row()), + indices + ) { } + + /// Construct a PredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ): PredicatedTileIterator(params, pointer, extent, thread_id, make_Coord(0, 0)) { } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator operator++(int) { + PredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for affine rank-2 data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIterator, AdvanceRank, + ThreadMap_, AccessSize, false> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRankN<2>; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + /// Type used for internal memory accesses + using AccessType = AlignedArray::value / 8)>; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = + PredicatedTileAccessIterator; + + static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename TileAccessIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + + friend PredicatedTileIterator; + + private: + /// Parameters object + typename TileAccessIterator::Params params_; + + public: + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) : params_(layout) { } + + /// Default constructor + Params() = default; + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + + public: + + /// Default constructor + PredicatedTileIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset, + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ) + : address_iterator_(params.params_, pointer, extent, thread_id, + threadblock_offset) {} + + /// Construct a PredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator &operator++() { + if (kAdvanceRank) + address_iterator_.add_tile_offset(make_Coord(0, 1)); + else + address_iterator_.add_tile_offset(make_Coord(1, 0)); + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator operator++(int) { + PredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { address_iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { address_iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { address_iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { address_iterator_.get_mask(mask); } + + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + load_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + address_iterator_.set_iteration_index(idx); + char const *byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + + AccessType const *access_ptr = reinterpret_cast(byte_ptr); + + mctlass::arch::global_load( + frag_ptr[idx], access_ptr, address_iterator_.valid()); + + ++address_iterator_; + } + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_byte_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + store_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) { + address_iterator_.set_iteration_index(0); + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + char *byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType *access_ptr = reinterpret_cast(byte_ptr); + + if (address_iterator_.valid()) { + *access_ptr = frag_ptr[idx]; + } + ++address_iterator_; + } + } + } + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_byte_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for affine rank 2 column-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize +> +class PredicatedTileIterator { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = PredicatedTileIterator< + layout::PitchLinearShape, + Element, + layout::AffineRankN<2>, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + AccessSize + >; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + + friend PredicatedTileIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default constructor + Params() = default; + + /// Construct the Params object given an AffineRankN<2> tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout): params_(layout::AffineRankN<2>(layout.stride(0), layout.stride(1))) + {} + }; + +private: + + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + +public: + + /// Default constructor + PredicatedTileIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const &threadblock_offset, ///< Initial offset of threadblock + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ): + iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column()) + ) { } + + /// Construct a PredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ): PredicatedTileIterator(params, pointer, extent, thread_id, make_Coord(0, 0)) { } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator operator++(int) { + PredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for affine rank 2 row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int AccessSize +> +class PredicatedTileIterator { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::AffineRank2RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + // Map to the underlying AffineRankN<2> layout + using UnderlyingIterator = PredicatedTileIterator< + layout::PitchLinearShape, + Element, + layout::AffineRankN<2>, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + AccessSize + >; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + + friend PredicatedTileIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default constructor + Params() = default; + + /// Construct the Params object given an AffineRankN<2> tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout): params_(layout::AffineRankN<2>(layout.stride(1), layout.stride(0))) {} + }; + + +private: + + // + // Data members + // + + /// Underlying AffineRankN<2> tile iterator + UnderlyingIterator iterator_; + +public: + + /// Default constructor + PredicatedTileIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const &threadblock_offset, ///< Initial offset of threadblock + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ): + iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row()) + ) { } + + /// Construct a PredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ): PredicatedTileIterator(params, pointer, extent, thread_id, make_Coord(0, 0)) { } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator operator++(int) { + PredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for interleaved data. It is mapped +/// to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// + +template +class PredicatedTileIterator, + AdvanceRank, ThreadMap_, AccessSize, false> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::ColumnMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileIterator< + layout::PitchLinearShape, + Element, layout::PitchLinear, (kAdvanceRank == 0 ? 0 : 1), ThreadMap, AccessSize>; + + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default constructor + Params() = default; + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + + MCTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const &base) + : params_(base) {} + + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + + /// Default constructor + PredicatedTileIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset, + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.row() * kInterleavedK, + extent.column() / kInterleavedK), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.row() * kInterleavedK, + threadblock_offset.column() / kInterleavedK)) {} + + /// Construct a PredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator operator++(int) { + PredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator for interleaved-32 data. It is +/// mapped to the congruous layout. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIterator, + AdvanceRank, ThreadMap_, AccessSize, false> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + static int const kInterleavedK = InterleavedK; + using Layout = layout::RowMajorInterleaved; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileIterator< + layout::PitchLinearShape, + Element, layout::PitchLinear, (kAdvanceRank == 0 ? 1 : 0), ThreadMap, AccessSize>; + + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + friend PredicatedTileIterator; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + /// Default constructor + Params() = default; + + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) + : params_(layout::PitchLinear(layout.stride(0))) {} + + MCTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const &base) + : params_(base) {} + }; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + + /// Default constructor + PredicatedTileIterator() = default; + + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset, + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ) + : iterator_(params.params_, pointer, + layout::PitchLinearCoord(extent.column() * kInterleavedK, + extent.row() / kInterleavedK), + thread_id, + layout::PitchLinearCoord( + threadblock_offset.column() * kInterleavedK, + threadblock_offset.row() / kInterleavedK)) {} + + /// Construct a PredicatedTileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIterator( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIterator(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator operator++(int) { + PredicatedTileIterator self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { iterator_.get_mask(mask); } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_iterator_2dthreadtile.h b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_iterator_2dthreadtile.h new file mode 100644 index 0000000..53eea95 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_iterator_2dthreadtile.h @@ -0,0 +1,787 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of tiles from pitch-linear rank=2 tensors. + + This iterator uses masks to guard out-of-bounds accesses and visits the last "residue" tile + first, with the objective of minimizing predicate mask updates during steady-state operation. + + A precomputed "Params" object minimizes the amount of state that must be stored in registers, + and integer addition is used to advance the pointer through memory. +*/ + +#pragma once + +#include "mctlass/transform/threadblock/predicated_tile_access_iterator_2dthreadtile.h" +#include "mctlass/transform/thread/transpose.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileIterator2dThreadTile +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +/// Regular tile iterator using a precomputed control structure to minimize register liveness +/// and integer arithmetic. +/// +/// Layout is assumed to be invariant at the time the precomputed "Params" object is constructed. +/// +/// Base pointer and tensor extents may be specified at the time the iterator is constructed. +/// Subsequently, they are assumed to be immutable. +/// +/// Adding a logical coordinate offset may be performed at the time the iterator is constructed. +/// Subsequent additions to logical coordinate offset may be performed but are relatively expensive. +/// +/// Vistitation order is intended to first visit a "residual" tile that may be partially full in +/// both the advance dimension and the steady-state dimension. This is assumed to be the last +/// tile in the iteration sequence. Advancing an iterator that has just been constructed moves to +/// the first tile that is full in the advance dimension and recomputes predicates. Subsequent +/// accesses may be performed without updating internal predicates and are efficient in terms of +/// live register state and pointer arithmetic instructions. +/// +/// To be efficient, this assumes the iteraor will be dereferenced and advanced at least once +/// outside any looping structure to minimize integer arithmetic. +/// +/// Acceses out of bounds are safe so long as `clear_mask()` is called prior to dereferencing +/// the iterator. +/// +/// +/// Example: +/// +/// An efficient pipeline structure may be constructed as follows: +/// +// template +// __global__ void kernel( +// typename Iterator::Params params, +// typename Iterator::Element *ptr, +// TensorCoord extent) { +// +// typename Iterator::Fragment fragment; +// +// TensorCoord threadblock_offset(0, 0); +// +// Iterator iter(params, ptr, extent, threadIdx.x, threadblock_offsets); +// +// +// fragment = *iter; // load "residue" tile first +// ++iter; // advance to first "steady state" tile and update internal masks +// +// +// #pragma unroll +// for (int i = Remaining - 1; i >= 0; --i) { +// +// f(fragment); +// +// if (!i) { +// iter.clear_mask(); // light-weight operation to clear masks - subsequent loads become NO-OPs. +// } +// +// fragment = *iter; // load tile during "steady state" phase +// ++iter; // advance to next tile - lightweight due to steady-state masks +// } +// } +// +// void host(TensorView view) { +// +// using Iterator = transform::threadblock::PredicatedTileIterator2dThreadTile; +// +// typename Iterator::Params params(view.layout()); +// +// kernel(params, view.data()); +// } +/// +/// +template < + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + bool Transpose = false +> +class PredicatedTileIterator2dThreadTile; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator2dThreadTile for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIterator2dThreadTile { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + /// Type used for internal memory accesses + /// extra set of parenthesis is needed for VS compiler + struct alignas((ThreadMap::kElementsPerAccess * sizeof_bits::value / + 8)) AccessType { + + Array storage; + + static int const kElements = ThreadMap::kElementsPerAccess; + }; + + /// Optinally this fragment can be 4x4 transposed + using Transform = thread::Transpose< ThreadMap::Iterations::kCount * ThreadMap::ThreadAccessShape::kCount , layout::PitchLinearShape<4,4>, Element>; + static bool const transpose = Transpose_; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = + PredicatedTileAccessIterator2dThreadTile; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename TileAccessIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + using Base = typename TileAccessIterator::Params::Base; + + friend PredicatedTileIterator2dThreadTile; + + private: + /// Parameters object + typename TileAccessIterator::Params params_; + + public: + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) : params_(layout) { } + + MCTLASS_HOST_DEVICE + Params() { } + + MCTLASS_HOST_DEVICE + Params(Base const &base) + : params_(base) {} + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIterator2dThreadTile( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset, + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ) + : address_iterator_(params.params_, pointer, extent, thread_id, + threadblock_offset) {} + + /// Construct a PredicatedTileIterator2dThreadTile with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIterator2dThreadTile( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIterator2dThreadTile(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator2dThreadTile &operator++() { + if (kAdvanceRank) + address_iterator_.add_tile_offset({0, 1}); + else + address_iterator_.add_tile_offset({1, 0}); + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator2dThreadTile operator++(int) { + PredicatedTileIterator2dThreadTile self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { address_iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { address_iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { address_iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { address_iterator_.get_mask(mask); } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + MCTLASS_PRAGMA_UNROLL + for (int ts = 0; ts < ThreadMap::ThreadAccessShape::kStrided; ts++){ + + int access_idx = ts + c * ThreadMap::ThreadAccessShape::kStrided + \ + s * ThreadMap::Iterations::kContiguous * ThreadMap::ThreadAccessShape::kStrided; + + address_iterator_.set_iteration_index(access_idx); + if (address_iterator_.valid()) { + + frag_ptr[access_idx] = + *(address_iterator_.get() + pointer_offset); + } + + ++address_iterator_; + } + } + } + + if (transpose) { + Transform t; + t.transform(frag, frag); + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + MCTLASS_PRAGMA_UNROLL + for (int ts = 0; ts < ThreadMap::ThreadAccessShape::kStrided; ts++){ + + int access_idx = ts + c * ThreadMap::ThreadAccessShape::kStrided + \ + s * ThreadMap::Iterations::kContiguous * ThreadMap::ThreadAccessShape::kStrided; + + address_iterator_.set_iteration_index(access_idx); + if (address_iterator_.valid()) { + *(address_iterator_.get() + pointer_offset) = frag_ptr[access_idx]; + } + ++address_iterator_; + } + } + } + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator2dThreadTile for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + bool Transpose_ +> +class PredicatedTileIterator2dThreadTile { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + static bool const Transpose = Transpose_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileIterator2dThreadTile< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + Transpose + >; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + + friend PredicatedTileIterator2dThreadTile; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) {} + + MCTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const &base) + : params_(base) {} + }; + + +private: + + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + + /// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIterator2dThreadTile( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const &threadblock_offset, ///< Initial offset of threadblock + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ): + iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column()) + ) { } + + /// Construct a PredicatedTileIterator2dThreadTile with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIterator2dThreadTile( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ): PredicatedTileIterator2dThreadTile(params, pointer, extent, thread_id, make_Coord(0, 0)) { } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator2dThreadTile &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator2dThreadTile operator++(int) { + PredicatedTileIterator2dThreadTile self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIterator2dThreadTile for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + bool Transpose_ +> +class PredicatedTileIterator2dThreadTile { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + static bool const Transpose = Transpose_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileIterator2dThreadTile< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + Transpose + >; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + + friend PredicatedTileIterator2dThreadTile; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) { } + + MCTLASS_HOST_DEVICE + Params(typename UnderlyingIterator::Params::Base const &base) + : params_(base) {} + }; + + +private: + + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + + /// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIterator2dThreadTile( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const &threadblock_offset, ///< Initial offset of threadblock + int const *indices = nullptr ///< gather/scatter indices, note no support for gather/scatter at this specialization + ): + iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row()) + ) { } + + /// Construct a PredicatedTileIterator2dThreadTile with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIterator2dThreadTile( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ): PredicatedTileIterator2dThreadTile(params, pointer, extent, thread_id, make_Coord(0, 0)) { } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator2dThreadTile &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIterator2dThreadTile operator++(int) { + PredicatedTileIterator2dThreadTile self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_iterator_triangular_matrix.h b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_iterator_triangular_matrix.h new file mode 100644 index 0000000..ab53f00 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_tile_iterator_triangular_matrix.h @@ -0,0 +1,818 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of tiles from pitch-linear rank=2 tensors. + + This iterator uses masks to guard out-of-bounds accesses and visits the last "residue" tile + first, with the objective of minimizing predicate mask updates during steady-state operation. + + A precomputed "Params" object minimizes the amount of state that must be stored in registers, + and integer addition is used to advance the pointer through memory. +*/ + +#pragma once + +#include "mctlass/arch/memory.h" +#include "mctlass/transform/threadblock/predicated_tile_access_iterator_triangular_matrix.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedTileIteratorTriangularMatrix +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +/// Regular tile iterator using a precomputed control structure to minimize register liveness +/// and integer arithmetic. +/// +/// Layout is assumed to be invariant at the time the precomputed "Params" object is constructed. +/// +/// Base pointer and tensor extents may be specified at the time the iterator is constructed. +/// Subsequently, they are assumed to be immutable. +/// +/// Adding a logical coordinate offset may be performed at the time the iterator is constructed. +/// Subsequent additions to logical coordinate offset may be performed but are relatively expensive. +/// +/// Vistitation order is intended to first visit a "residual" tile that may be partially full in +/// both the advance dimension and the steady-state dimension. This is assumed to be the last +/// tile in the iteration sequence. Advancing an iterator that has just been constructed moves to +/// the first tile that is full in the advance dimension and recomputes predicates. Subsequent +/// accesses may be performed without updating internal predicates and are efficient in terms of +/// live register state and pointer arithmetic instructions. +/// +/// To be efficient, this assumes the iteraor will be dereferenced and advanced at least once +/// outside any looping structure to minimize integer arithmetic. +/// +/// Acceses out of bounds are safe so long as `clear_mask()` is called prior to dereferencing +/// the iterator. +/// +/// +/// Example: +/// +/// An efficient pipeline structure may be constructed as follows: +/// +// template +// __global__ void kernel( +// typename Iterator::Params params, +// typename Iterator::Element *ptr, +// TensorCoord extent) { +// +// typename Iterator::Fragment fragment; +// +// TensorCoord threadblock_offset(0, 0); +// +// Iterator iter(params, ptr, extent, threadIdx.x, threadblock_offsets); +// +// +// fragment = *iter; // load "residue" tile first +// ++iter; // advance to first "steady state" tile and update internal masks +// +// +// #pragma unroll +// for (int i = Remaining - 1; i >= 0; --i) { +// +// f(fragment); +// +// if (!i) { +// iter.clear_mask(); // light-weight operation to clear masks - subsequent loads become NO-OPs. +// } +// +// fragment = *iter; // load tile during "steady state" phase +// ++iter; // advance to next tile - lightweight due to steady-state masks +// } +// } +// +// void host(TensorView view) { +// +// using Iterator = transform::threadblock::PredicatedTileIteratorTriangularMatrix; +// +// typename Iterator::Params params(view.layout()); +// +// kernel(params, view.data()); +// } +/// +/// +template < + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + SideMode kSideMode, + FillMode kFillMode, + DiagType kDiagType, + int AccessSize = ThreadMap::kElementsPerAccess +> +class PredicatedTileIteratorTriangularMatrix; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorTriangularMatrix for pitch-linear data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template +class PredicatedTileIteratorTriangularMatrix { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + /// Type used for internal memory accesses + using AccessType = AlignedArray::value / 8)>; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = + PredicatedTileAccessIteratorTriangularMatrix; + + static int const kAccessesPerVector = TileAccessIterator::kAccessesPerVector; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename TileAccessIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + public: + friend PredicatedTileIteratorTriangularMatrix; + + private: + /// Parameters object + typename TileAccessIterator::Params params_; + + public: + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout) : params_(layout) { } + + MCTLASS_HOST_DEVICE + Params() { } + }; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIteratorTriangularMatrix( + /// Precomputed parameters object + Params const ¶ms, + /// Pointer to start of tensor + Pointer pointer, + /// Extent of tensor + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : address_iterator_(params.params_, pointer, extent, thread_id, + threadblock_offset) {} + + /// Construct a PredicatedTileIteratorTriangularMatrix with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIteratorTriangularMatrix( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ) + : PredicatedTileIteratorTriangularMatrix(params, pointer, extent, thread_id, + make_Coord(0, 0)) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIteratorTriangularMatrix &operator++() { + if (kAdvanceRank) + address_iterator_.add_tile_offset({0, 1}); + else + address_iterator_.add_tile_offset({1, 0}); + + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIteratorTriangularMatrix operator++(int) { + PredicatedTileIteratorTriangularMatrix self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { address_iterator_.clear_mask(enable); } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { address_iterator_.enable_mask(); } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { address_iterator_.set_mask(mask); } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { address_iterator_.get_mask(mask); } + + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + load_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + address_iterator_.set_iteration_index(idx); + char const *byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + + AccessType const *access_ptr = reinterpret_cast(byte_ptr); + + mctlass::arch::global_load( + frag_ptr[idx], access_ptr, address_iterator_.valid()); + + ++address_iterator_; + } + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_byte_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + store_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) { + address_iterator_.set_iteration_index(0); + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + MCTLASS_PRAGMA_UNROLL + for (int v = 0; v < kAccessesPerVector; ++v) { + + int idx = v + kAccessesPerVector * (c + s * ThreadMap::Iterations::kContiguous); + + char *byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType *access_ptr = reinterpret_cast(byte_ptr); + + if (address_iterator_.valid()) { + *access_ptr = frag_ptr[idx]; + } + ++address_iterator_; + } + } + } + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_byte_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorTriangularMatrix for column-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + SideMode kSideMode, + FillMode kFillMode, + DiagType kDiagType, + int AccessSize +> +class PredicatedTileIteratorTriangularMatrix { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileIteratorTriangularMatrix< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap, + kSideMode, + kFillMode, + kDiagType, + AccessSize + >; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + + friend PredicatedTileIteratorTriangularMatrix; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) { + + } + }; + + +private: + + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + + /// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIteratorTriangularMatrix( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const &threadblock_offset ///< Initial offset of threadblock + ): + iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.row(), extent.column()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.row(), threadblock_offset.column()) + ) { } + + /// Construct a PredicatedTileIteratorTriangularMatrix with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIteratorTriangularMatrix( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ): PredicatedTileIteratorTriangularMatrix(params, pointer, extent, thread_id, make_Coord(0, 0)) { } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIteratorTriangularMatrix &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIteratorTriangularMatrix operator++(int) { + PredicatedTileIteratorTriangularMatrix self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedTileIteratorTriangularMatrix for row-major data. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept | +/// MaskedTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + SideMode kSideMode, + FillMode kFillMode, + DiagType kDiagType, + int AccessSize +> +class PredicatedTileIteratorTriangularMatrix { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using Pointer = Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedTileIteratorTriangularMatrix< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + kSideMode, + kFillMode, + kDiagType, + AccessSize + >; + + using AccessType = typename UnderlyingIterator::AccessType; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array; + + /// Predicate vector stores mask to guard accesses + using Mask = typename UnderlyingIterator::Mask; + + /// Parameters object is precomputed state and is host-constructible + class Params { + private: + + friend PredicatedTileIteratorTriangularMatrix; + + /// Parameters object + typename UnderlyingIterator::Params params_; + + public: + + MCTLASS_HOST_DEVICE + Params() { } + + /// Construct the Params object given a pitch-linear tensor's layout + MCTLASS_HOST_DEVICE + Params(Layout const &layout): params_(layout::PitchLinear(layout.stride(0))) { + + }; + }; + + +private: + + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + +public: + + /// Constructs a TileIterator from its precomputed state, threadblock offset, and thread ID + MCTLASS_HOST_DEVICE + PredicatedTileIteratorTriangularMatrix( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + TensorCoord const &threadblock_offset ///< Initial offset of threadblock + ): + iterator_( + params.params_, + pointer, + layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, + layout::PitchLinearCoord(threadblock_offset.column(), threadblock_offset.row()) + ) { } + + /// Construct a PredicatedTileIteratorTriangularMatrix with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedTileIteratorTriangularMatrix( + Params const ¶ms, ///< Precomputed parameters object + Pointer pointer, ///< Pointer to start of tensor + TensorCoord extent, ///< Extent of tensor + int thread_id ///< ID of each participating thread + ): PredicatedTileIteratorTriangularMatrix(params, pointer, extent, thread_id, make_Coord(0, 0)) { } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIteratorTriangularMatrix &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the iterator's + /// internal pointer is reverted to the first "steady state" tile. Subsequent calls + /// are lightweight and must only update the internal pointer. + MCTLASS_HOST_DEVICE + PredicatedTileIteratorTriangularMatrix operator++(int) { + PredicatedTileIteratorTriangularMatrix self(*this); + operator++(); + return self; + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void clear_mask(bool enable = true) { + iterator_.clear_mask(enable); + } + + /// Clears the predicate set efficiently + MCTLASS_HOST_DEVICE + void enable_mask() { + iterator_.enable_mask(); + } + + /// Sets the predicate mask, overriding value stored in predicate iterator + MCTLASS_HOST_DEVICE + void set_mask(Mask const &mask) { + iterator_.set_mask(mask); + } + + /// Gets the mask + MCTLASS_HOST_DEVICE + void get_mask(Mask &mask) { + iterator_.get_mask(mask); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, LongIndex byte_offset) { + iterator_.load_with_byte_offset(frag, byte_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, LongIndex byte_offset) { + iterator_.store_with_byte_offset(frag, byte_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/predicated_vector_access_iterator.h b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_vector_access_iterator.h new file mode 100644 index 0000000..1d400c2 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/predicated_vector_access_iterator.h @@ -0,0 +1,417 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates implementing computing the addresses of loading small + vectors from the global memory. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/coord.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// PredicatedVectorAccessIterator +/// +template < + /// Shape of the vector accessed by the entire threadblock + typename Shape, + /// Shape of the vector accessed by the warp + typename WarpShape, + /// Type of Element + typename Element, + /// Layout of the vector + typename Layout, + /// Number of elements for each access + int ElementsPerAccess, + /// Support residual tile + bool EnableResidualAccess = false +> +class PredicatedVectorAccessIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Vector access iterator specialized for vectors, e.g. scale and bias +/// Thread arrangements are for TensorOps +/// +template < + typename Shape_, + typename WarpShape_, + typename Element_, + int ElementsPerAccess, + bool EnableResidualAccess +> +class PredicatedVectorAccessIterator < + Shape_, + WarpShape_, + Element_, + layout::PitchLinear, + ElementsPerAccess, + EnableResidualAccess +> { + public: + + using Shape = Shape_; + using WarpShape = WarpShape_; + using Element = Element_; + using Layout = layout::PitchLinear; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + +// static int const kElementsPerAccess = 128 / sizeof_bits::value; + static int const kElementsPerAccess = ElementsPerAccess; + static int const kThreads = 32; + static int const kRowsPerIteration = 8; + static int const kThreadsPerRow = kThreads / kRowsPerIteration; + static int const kThreadsPerRowMask = 0x3; + static int const kIterations = WarpShape::kContiguous / (kThreadsPerRow * kElementsPerAccess); + static int const kWarpCountStrided = Shape::kStrided / WarpShape::kStrided; + + using AccessType = AlignedArray; + + private: + /// Internal pointer type permits fast address arithmetic + using BytePointer = char *; + + private: + // + // Data members + // + + /// Internal pointer to first access of tile + BytePointer pointer_; + + /// Extent of tensor + TensorCoord extent_; + + /// pointer offset of each thread + TensorCoord thread_offset_; + + /// iteration index + LongIndex iteration_; + + /// residual access + bool is_residual_; + + /// residual offset of each thread + TensorCoord residual_offset_; + + public: + /// Constructs a vector access iterator + MCTLASS_HOST_DEVICE + PredicatedVectorAccessIterator( + /// Pointer to the start of the vector + ConstPointer pointer, + /// Extent of vector + TensorCoord extent, + /// ID of each participating thread + int thread_id, + /// ID of each participating warp + int warp_id, + /// Initial offset of threadblock + TensorCoord const &threadblock_offset) + : pointer_(reinterpret_cast( + const_cast(pointer))), + extent_(extent), + is_residual_(false) { + + + int warp_offset = (warp_id / kWarpCountStrided) * WarpShape::kContiguous; + + // Per-thread offset in logical coordinates of tensor + + thread_offset_ = threadblock_offset + TensorCoord(warp_offset, 0) + + TensorCoord((thread_id & kThreadsPerRowMask) * kElementsPerAccess, 0); + + set_iteration_index(0); + + if(EnableResidualAccess) { + // compute residual offset + typename TensorCoord::Index residual_size = extent_.contiguous() % WarpShape::kContiguous; + if (residual_size) { + is_residual_ = true; + residual_offset_ = make_Coord(residual_size, 0); + } + } + } + + /// Construct a PredicatedVectorAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedVectorAccessIterator( + /// Pointer to start of vector + ConstPointer pointer, + /// Extent of vector + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + /// ID of each participating warp + int warp_id) + : PredicatedVectorAccessIterator(pointer, extent, thread_id, warp_id, + make_Coord(0, 0)) {} + + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iteration_ = index; + } + + /// Advances an iterator along logical dimensions of matrix in units of whole tiles + MCTLASS_DEVICE + void add_tile_offset( + TensorCoord const &tile_offset) { + + thread_offset_ = + thread_offset_ + + TensorCoord(WarpShape::kContiguous * tile_offset.contiguous(), 0); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + + return reinterpret_cast( + pointer_ + + ((thread_offset_.contiguous() + iteration_ * kThreadsPerRow * kElementsPerAccess) + * sizeof_bits::value / 8)); + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + PredicatedVectorAccessIterator &operator++() { + ++iteration_; + if(iteration_ >= kIterations) + iteration_ = 0; + + return *this; + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + void advance() { + if(EnableResidualAccess && is_residual_) { + is_residual_ = false; + thread_offset_ += residual_offset_; + } + else + add_tile_offset(TensorCoord(1, 0)); + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + PredicatedVectorAccessIterator operator++(int) { + PredicatedVectorAccessIterator self(*this); + operator++(); + return self; + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return ((thread_offset_.contiguous() + + iteration_ * kThreadsPerRow * kElementsPerAccess) < extent_.contiguous()); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Specialization of PredicatedVectorAccessIterator for row-major data. +/// +template < + typename Shape_, + typename WarpShape_, + typename Element_, + int ElementsPerAccess, + bool EnableResidualAccess +> +class PredicatedVectorAccessIterator< + Shape_, + WarpShape_, + Element_, + layout::RowMajor, + ElementsPerAccess, + EnableResidualAccess +> { + public: + + using Shape = Shape_; + using WarpShape = WarpShape_; + using Element = Element_; + using Layout = layout::RowMajor; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorView = TensorView; + using TensorCoord = typename Layout::TensorCoord; + + using ConstPointer = const Element *; + using NonConstPointer = typename platform::remove_const::type *; + + using UnderlyingIterator = PredicatedVectorAccessIterator< + layout::PitchLinearShape, + layout::PitchLinearShape, + Element, + layout::PitchLinear, + ElementsPerAccess, + EnableResidualAccess>; + + using AccessType = typename UnderlyingIterator::AccessType; + static int const kElementsPerAccess = UnderlyingIterator::kElementsPerAccess; + static int const kRowsPerIteration = UnderlyingIterator::kRowsPerIteration; + static int const kThreads = UnderlyingIterator::kThreads; + static int const kIterations = UnderlyingIterator::kIterations; + + private: + // + // Data members + // + + /// Underlying pitch-linear tile iterator + UnderlyingIterator iterator_; + + public: + /// Constructs a TileIterator from its precomputed state, threadblock offset, + /// and thread ID + MCTLASS_HOST_DEVICE + PredicatedVectorAccessIterator( + ///< Pointer to the start of the vector + ConstPointer pointer, + ///< Extent of tensor + TensorCoord extent, + ///< ID of each participating thread + int thread_id, + ///< ID of each participating warp + int warp_id, + ///< Initial offset of threadblock + TensorCoord const &threadblock_offset) + : iterator_(pointer, layout::PitchLinearCoord(extent.column(), extent.row()), + thread_id, warp_id, + layout::PitchLinearCoord(threadblock_offset.column(), + threadblock_offset.row())) {} + + /// Construct a PredicatedVectorAccessIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + PredicatedVectorAccessIterator( + ConstPointer pointer, ///< Pointer to the start of the vector + TensorCoord extent, ///< Extent of tensor + int thread_id, ///< ID of each participating thread + int warp_id ///< ID of each participating warp + ) + : PredicatedVectorAccessIterator(pointer, extent, thread_id, warp_id, + make_Coord(0, 0)) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Advances an iterator along logical dimensions of matrix in units of whole + /// tiles + MCTLASS_HOST_DEVICE + void add_tile_offset(TensorCoord const &tile_offset) { + iterator_.add_tile_offset({tile_offset.column(), tile_offset.row()}); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedVectorAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + /// + /// The first time this method is called, predicates are updated, and the + /// iterator's internal pointer is reverted to the first "steady state" tile. + /// Subsequent calls are lightweight and must only update the internal + /// pointer. + MCTLASS_HOST_DEVICE + PredicatedVectorAccessIterator operator++(int) { + PredicatedVectorAccessIterator self(*this); + operator++(); + return self; + } + + /// Increment and return an instance to self. + MCTLASS_HOST_DEVICE + void advance() { + iterator_.advance(); + } + + /// Returns whether access is valid or not + MCTLASS_HOST_DEVICE + bool valid() { + return iterator_.valid(); + } +}; + + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/regular_scale_bias_vector_access_iterator.h b/csrc/mctlass/include/mctlass/transform/threadblock/regular_scale_bias_vector_access_iterator.h new file mode 100644 index 0000000..367dc18 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/regular_scale_bias_vector_access_iterator.h @@ -0,0 +1,253 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +/*! \file + \brief Templates implementing computing the addresses of storing of small + scale and bias vectors in the shared memory. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// RegularScaleBiasVectorAccessIterator +/// +template +class RegularScaleBiasVectorAccessIterator; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for congruous arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularScaleBiasVectorAccessIterator { + public: + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + /// Element type per access + static int const kElementsPerAccess = 128 / sizeof_bits::value; + static int const kThreads = Shape::kContiguous / kElementsPerAccess; + using AccessType = Array; + + private: + // + // Data members + // + + /// Internal pointer + AccessType *pointer_; + + /// Internal byte offset + Index byte_offset_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularScaleBiasVectorAccessIterator( + TensorRef scale_bias_ref, ///< Pointer to the start of the scale and bias + ///< vector + int thread_id ///< ID of each participating thread + ) + : byte_offset_(0) { + // Per-thread offset in logical coordinates of tensor + int thread_offset = thread_id * kElementsPerAccess; + + // initialize pointer + pointer_ = + reinterpret_cast(scale_bias_ref.data() + thread_offset); + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_offset_ += pointer_offset * sizeof(Element); + } + + /// Returns a pointer + MCTLASS_DEVICE + AccessType *get() const { + + char *access_byte_ptr = + reinterpret_cast(pointer_); + + return reinterpret_cast(access_byte_ptr + byte_offset_); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularScaleBiasVectorAccessIterator &operator++() { return *this; } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularScaleBiasVectorAccessIterator operator++(int) { + RegularScaleBiasVectorAccessIterator prev(*this); + this->operator++(); + + return prev; + } + + /// Adds a tile offset in the unit of tile. + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + // Multiply by 2 because we store scale and bias belong to the same stage + // next to each other. + add_pointer_offset(coord.contiguous() * Shape::kContiguous * 2); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for row major layouts +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularScaleBiasVectorAccessIterator< + Shape_, Element_, + layout::RowMajor> { + public: + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + /// Underlying iterator type + using UnderlyingIterator = RegularScaleBiasVectorAccessIterator< + layout::PitchLinearShape, Element, + layout::PitchLinear>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularScaleBiasVectorAccessIterator( + TensorRef scale_bias_ref, ///< Pointer to the start of the scale and bias + ///< vector + int thread_id ///< ID of each participating thread + ) + : iterator_({scale_bias_ref.data(), scale_bias_ref.stride()}, thread_id) { + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularScaleBiasVectorAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularScaleBiasVectorAccessIterator operator++(int) { + RegularScaleBiasVectorAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator.h b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator.h new file mode 100644 index 0000000..af7bc0d --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator.h @@ -0,0 +1,58 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing the address computation of storing of tiles + from pitch-linear rank=2 tensors. +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template ::value* ThreadMap::kElementsPerAccess / 8> +class RegularTileAccessIterator; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear.h b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear.h new file mode 100644 index 0000000..f2fd5c6 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear.h @@ -0,0 +1,408 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing computing the addresses of storing of tiles + from pitch-linear rank=2 tensors. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" + +#include "mctlass/transform/threadblock/regular_tile_access_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for congruous arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::PitchLinear, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Element type per access + using AccessType = Array; + + private: + // + // Data members + // + + /// Stride value + StrideIndex stride_; + + /// Internal pointer to first access of tile + AccessType *pointer_; + + /// Internal byte offset + Index byte_offset_; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : stride_(ref.stride(0) / ThreadMap::kElementsPerAccess), + byte_offset_(0) { + + layout::PitchLinearCoord thread_offset_base = ThreadMap::initial_offset(thread_id); + + // initialize pointer + pointer_ = reinterpret_cast(ref.data() + ref.offset(thread_offset_base)); + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_offset_ += pointer_offset * sizeof(Element); + } + + /// Returns a pointer + MCTLASS_DEVICE + AccessType *get() const { + + AccessType *access_ptr = pointer_; + + int access_offset = iteration_strided_ * ThreadMap::Delta::kStrided * stride_ + + iteration_contiguous_ * ThreadMap::Delta::kContiguous / + ThreadMap::kElementsPerAccess; + + char *access_byte_ptr = + reinterpret_cast(access_ptr + access_offset); + + return reinterpret_cast(access_byte_ptr + byte_offset_); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iteration_contiguous_; + + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) + return *this; + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + iteration_contiguous_ = 0; + ++iteration_strided_; + + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + iteration_strided_ = 0; + + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + this->operator++(); + + return prev; + } + + /// Adds a tile offset in the unit of tile. + /// In GEMM/Conv implementation, this is used to move in the k dimension in the shared memory. + /// Below layouts are the shared memory layouts. Current SM50 SIMT kernels only use col major A and row major B. + /// For row major A operand, k dimension is contiguous dimension; + /// For col major A operand, k dimension is strided dimension; + /// For row major B operand, k dimension is strided dimension; + /// For col major B operand, k dimension is contiguous dimension. + /// Below two classes map col/row major to the pitch linear coordinates used + /// in this base class. + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + add_pointer_offset(coord.contiguous() * Shape::kContiguous + + coord.strided() * Shape::kStrided * stride_ * + ThreadMap::kElementsPerAccess); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for column major layouts +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::ColumnMajor, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for row major layouts +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::RowMajor, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear_direct_conv.h b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear_direct_conv.h new file mode 100644 index 0000000..140478b --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_pitch_linear_direct_conv.h @@ -0,0 +1,587 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing computing the addresses of storing of tiles + from pitch-linear rank=2 tensors. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" + +#include "mctlass/transform/threadblock/regular_tile_access_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + + +//////////////////////////////////////////////////////////////////////////////// + +template ::value* ThreadMap::kElementsPerAccess / 8 + > +class RegularTileAccessIteratorDirectConv; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for congruous arrangements for TensorOps with dynamic_iterations OFF +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIteratorDirectConv< + Shape_, Element_, + layout::PitchLinear, + AdvanceRank, ThreadMap_, false, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Element type per access + using AccessType = Array; + + private: + // + // Data members + // + + /// Stride value + StrideIndex stride_; + + /// Internal pointer to first access of tile + AccessType *pointer_; + + /// Internal byte offset + Index byte_offset_; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIteratorDirectConv(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : stride_(ref.stride(0) / ThreadMap::kElementsPerAccess), + byte_offset_(0) { + + layout::PitchLinearCoord thread_offset_base = ThreadMap::initial_offset(thread_id); + + // initialize pointer + pointer_ = reinterpret_cast(ref.data() + ref.offset(thread_offset_base)); + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_num(int num) { + //Do nothing + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_offset_ += pointer_offset * sizeof(Element); + } + + /// Returns a pointer + MCTLASS_DEVICE + AccessType *get() const { + + AccessType *access_ptr = pointer_; + + int access_offset = iteration_strided_ * ThreadMap::Delta::kStrided * stride_ + + iteration_contiguous_ * ThreadMap::Delta::kContiguous / + ThreadMap::kElementsPerAccess; + + char *access_byte_ptr = + reinterpret_cast(access_ptr + access_offset); + + return reinterpret_cast(access_byte_ptr + byte_offset_); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIteratorDirectConv &operator++() { + ++iteration_contiguous_; + + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) + return *this; + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + iteration_contiguous_ = 0; + ++iteration_strided_; + + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + iteration_strided_ = 0; + + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIteratorDirectConv operator++(int) { + RegularTileAccessIteratorDirectConv prev(*this); + this->operator++(); + + return prev; + } + + /// Adds a tile offset in the unit of tile. + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + add_pointer_offset(coord.contiguous() * Shape::kContiguous + + coord.strided() * ThreadMap::Iterations::kStrided * + ThreadMap::Delta::kStrided * stride_ * ThreadMap::kElementsPerAccess); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for congruous arrangements for TensorOps with dynamic_iterations ON +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIteratorDirectConv< + Shape_, Element_, + layout::PitchLinear, + AdvanceRank, ThreadMap_,true, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Element type per access + using AccessType = Array; + + private: + // + // Data members + // + + /// Stride value + StrideIndex stride_; + + /// Internal pointer to first access of tile + AccessType *pointer_; + + /// Internal byte offset + Index byte_offset_; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + /// Total iterattions in the strided dimension: Dynamic value + int total_iteration_strided_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIteratorDirectConv(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : stride_(ref.stride(0) / ThreadMap::kElementsPerAccess), + byte_offset_(0) { + + layout::PitchLinearCoord thread_offset_base = ThreadMap::initial_offset(thread_id); + + // initialize pointer + pointer_ = reinterpret_cast(ref.data() + ref.offset(thread_offset_base)); + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_num(int num) { + total_iteration_strided_ = num; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_offset_ += pointer_offset * sizeof(Element); + } + + /// Returns a pointer + MCTLASS_DEVICE + AccessType *get() const { + + AccessType *access_ptr = pointer_; + + int access_offset = iteration_strided_ * ThreadMap::Delta::kStrided * stride_ + + iteration_contiguous_ * ThreadMap::Delta::kContiguous / + ThreadMap::kElementsPerAccess; + + char *access_byte_ptr = + reinterpret_cast(access_ptr + access_offset); + + return reinterpret_cast(access_byte_ptr + byte_offset_); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIteratorDirectConv &operator++() { + ++iteration_contiguous_; + + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) + return *this; + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + iteration_contiguous_ = 0; + ++iteration_strided_; + + if (iteration_strided_ < total_iteration_strided_) { + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + iteration_strided_ = 0; + + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIteratorDirectConv operator++(int) { + RegularTileAccessIteratorDirectConv prev(*this); + this->operator++(); + + return prev; + } + + /// Adds a tile offset in the unit of tile. + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + add_pointer_offset(coord.contiguous() * Shape::kContiguous + + coord.strided() * total_iteration_strided_ * ThreadMap::Delta::kStrided * stride_ * + ThreadMap::kElementsPerAccess); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for column major layouts +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIteratorDirectConv< + Shape_, Element_, + layout::ColumnMajor, + AdvanceRank, ThreadMap_, Dynamic_iterations , Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIteratorDirectConv< + layout::PitchLinearShape, Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap_, + Dynamic_iterations>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIteratorDirectConv(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_num(int num) { + iterator_.set_iteration_num(num); + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIteratorDirectConv &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIteratorDirectConv operator++(int) { + RegularTileAccessIteratorDirectConv prev(*this); + ++iterator_; + + return prev; + } +}; + + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for row major layouts +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIteratorDirectConv< + Shape_, Element_, + layout::RowMajor, + AdvanceRank, ThreadMap_, Dynamic_iterations, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIteratorDirectConv< + layout::PitchLinearShape, Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap_, + Dynamic_iterations>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIteratorDirectConv(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_num(int num) { + iterator_.set_iteration_num(num); + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIteratorDirectConv &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIteratorDirectConv operator++(int) { + RegularTileAccessIteratorDirectConv prev(*this); + ++iterator_; + + return prev; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h new file mode 100644 index 0000000..463f040 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h @@ -0,0 +1,820 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing computing the addresses of storing of tiles + from pitch-linear rank=2 tensors. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for congruous arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::TensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element_))>, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = + layout::TensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element_))>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Internal details made public to facilitate introspection + struct Detail { + /// This iterator is specialized for an access size that is 128 bits in + /// length. + static int const kAccessSizeInBits = 128; + + static_assert(sizeof_bits::value * + ThreadMap::kElementsPerAccess == + kAccessSizeInBits, + "This iterator requires a policy whose access size is 128bs"); + + ///< Number of pointers + static int const kPointerCount = + (ThreadMap::Iterations::kStrided > 1 ? 2 : 1); + }; + + /// Element type per access + using AccessType = Array; + + private: + // + // Data members + // + + /// Stride value + StrideIndex stride_; + + /// Internal pointer to first access of tile + AccessType *pointer_[Detail::kPointerCount]; + + /// Internal byte offset + Index byte_offset_; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : stride_(ref.stride(0) / Layout::kElementsPerAccess), + byte_offset_(0) { + layout::PitchLinearCoord thread_offset_base = + ThreadMap::initial_offset(thread_id); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Detail::kPointerCount; ++i) { + // This is the offset of a thread within a threadblock tile for a specific + // pointer (units of elements) + layout::PitchLinearCoord thread_offset_in_threadblock_tile = + thread_offset_base + + layout::PitchLinearCoord{ + 0, ThreadMap::Detail::WarpThreadArrangement::kStrided * i}; + + // initialize pointer + pointer_[i] = reinterpret_cast( + ref.data() + ref.offset(thread_offset_in_threadblock_tile)); + } + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_offset_ += pointer_offset * sizeof(Element); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + AccessType *access_ptr = pointer_[iteration_strided_ & 1]; + int stride_idx = (iteration_strided_ & ~1); + + int access_offset = stride_idx * ThreadMap::Delta::kStrided * stride_ + + iteration_contiguous_ * ThreadMap::Delta::kContiguous / + ThreadMap::kElementsPerAccess; + + char *access_byte_ptr = + reinterpret_cast(access_ptr + access_offset); + return reinterpret_cast(access_byte_ptr + byte_offset_); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iteration_contiguous_; + + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) + return *this; + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + iteration_contiguous_ = 0; + ++iteration_strided_; + + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + + // Enter here only if (iteration_strided_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + iteration_strided_ = 0; + + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + this->operator++(); + + return prev; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + add_pointer_offset(coord.contiguous() * Shape::kContiguous + + coord.strided() * Shape::kStrided * stride_ * + Layout::kElementsPerAccess); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for column-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element_))>, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for column-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element_))>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element_))>, + (kAdvanceRank == 0 ? 0 : 1), ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for row-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::RowMajorTensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element_))>, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for row-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element_))>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element_))>, + (kAdvanceRank == 0 ? 1 : 0), ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for crosswise arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator::value, Crosswise>, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = + layout::TensorOpMultiplicandCrosswise::value, + Crosswise>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + static int const kCrosswise = Crosswise; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + static_assert(!(ThreadMap::Delta::kContiguous % kCrosswise), + "kCrosswise is the smallest unit in the contiguous dimension " + "for shared memory swizzling."); + + /// Internal details made public to facilitate introspection + struct Detail { + /// This iterator is specialized for an access size that is 128 bits in + /// length. + static int const kAccessSizeInBits = 128; + + static_assert(sizeof_bits::value * + ThreadMap::kElementsPerAccess == + kAccessSizeInBits, + "This iterator requires a policy whose access size is 128bs"); + + /// Number of pointers + /// + /// Note:TN kblock32 layouts only needs 1 pointer, but strangely + /// reducing pointer count hurts perfomrnace + static int const kPointerCount = + (ThreadMap::Iterations::kStrided > 1 ? 2 : 1); + }; + + /// Element type per access + using AccessType = Array; + + private: + // + // Data members + // + + /// Total number of sections. The memory is divided into stages. One stage + /// can store one tile. Stage is divided into sections. Interleaved layout + /// can have multiple sections in a stage. The rest layout only has one section + /// in a stage. + int sections_; + + /// Sections that a stage has + int sections_per_stage_; + + /// Stride value + StrideIndex stride_; + + /// Internal pointer to first access of tile + AccessType *pointer_[Detail::kPointerCount]; + + /// Internal byte offset + Index byte_offset_; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : sections_(ref.stride(0) / kCrosswise), + sections_per_stage_(Shape::kContiguous / kCrosswise), + // stride_ = kCrosswise x sections_ x kFactor + stride_(ref.stride(0) * Layout::kFactor / Layout::kElementsPerAccess), + byte_offset_(0) { + layout::PitchLinearCoord thread_offset_base = + ThreadMap::initial_offset(thread_id); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Detail::kPointerCount; ++i) { + // This is the offset of a thread within a threadblock tile for a specific + // pointer (units of elements) + layout::PitchLinearCoord thread_offset_in_threadblock_tile = + thread_offset_base + + layout::PitchLinearCoord{ + 0, ThreadMap::Detail::WarpThreadArrangement::kStrided * i}; + // initialize pointer + pointer_[i] = reinterpret_cast(ref.data()) + + ref.offset(thread_offset_in_threadblock_tile) / + Layout::kElementsPerAccess; + } + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_offset_ += pointer_offset * sizeof_bits::value / 8; + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + AccessType *access_ptr = pointer_[iteration_strided_ & 1]; + int stride_idx = (iteration_strided_ & ~1); + + int access_offset = + stride_idx * ThreadMap::Delta::kStrided * stride_ / Layout::kFactor + + // kCrosswise elements in the contiguous dimension would span to a + // shared memory cache line. + iteration_contiguous_ * (ThreadMap::Delta::kContiguous / kCrosswise) * + Layout::TileShape::kContiguous; + char *access_byte_ptr = + reinterpret_cast(access_ptr + access_offset); + return reinterpret_cast(access_byte_ptr + byte_offset_); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iteration_contiguous_; + + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) + return *this; + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + iteration_contiguous_ = 0; + ++iteration_strided_; + + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + + // Enter here only if (iteration_strided_ == ThreadMap::Iteration::kStrided) + // which means we enter the next section. + iteration_strided_ = 0; + + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + this->operator++(); + + return prev; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + add_pointer_offset(coord.contiguous() * sections_per_stage_ * stride_ * + ThreadMap::kElementsPerAccess / sections_ + + coord.strided() * Shape::kStrided * stride_ * + Layout::kElementsPerAccess / Layout::kFactor); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for column-major crosswise TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Crosswise>, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for column-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Crosswise>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCrosswise::value, + Crosswise>, + (kAdvanceRank == 0 ? 0 : 1), ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for row-major crosswise TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator::value, Crosswise>, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for row-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Crosswise>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCrosswise::value, + Crosswise>, + (kAdvanceRank == 0 ? 1 : 0), ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op_sm80.h b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op_sm80.h new file mode 100644 index 0000000..c674588 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op_sm80.h @@ -0,0 +1,1532 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing computing the addresses of storing of tiles + from pitch-linear rank=2 tensors. +*/ + +#pragma once + +#include "mctlass/array.h" +#include "mctlass/mctlass.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor_op_multiplicand_sm75.h" +#include "mctlass/layout/tensor_op_multiplicand_sm80.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for congruous arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::TensorOpMultiplicandCongruous64b, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorOpMultiplicandCongruous64b; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + static_assert(ThreadMap::kThreads / 32 > 1, + "This tile iterator requires at least two warps."); + + /// Internal details made public to facilitate introspection + struct Detail { + /// This iterator is specialized for an access size that is 128 bits in + /// length. + static int const kAccessSizeInBits = 64; + + static_assert(sizeof_bits::value * + ThreadMap::kElementsPerAccess == + kAccessSizeInBits, + "This iterator requires a policy whose access size is 64b"); + + ///< Number of pointers + static int const kPointerCount = 1; + }; + + /// Element type per access + using AccessType = Array; + + private: + // + // Data members + // + + /// Stride value + StrideIndex stride_; + + /// Internal pointer to first access of tile + AccessType *pointer_; + + /// Internal byte offset + Index byte_offset_; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + public: + + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): + stride_(ref.stride(0) / Layout::kElementsPerAccess), + byte_offset_(0) { + + layout::PitchLinearCoord thread_offset_base = ThreadMap::initial_offset(thread_id); + + // This is the offset of a thread within a threadblock tile for a specific + // pointer (units of elements) + layout::PitchLinearCoord thread_offset_in_threadblock_tile = thread_offset_base; + + // initialize pointer + pointer_ = reinterpret_cast(ref.data() + ref.offset(thread_offset_in_threadblock_tile)); + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + + byte_offset_ += pointer_offset * sizeof(Element); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + + AccessType *access_ptr = pointer_; + + int access_offset = iteration_strided_ * ThreadMap::Delta::kStrided * stride_ + + iteration_contiguous_ * ThreadMap::Delta::kContiguous / + ThreadMap::kElementsPerAccess; + + char *access_byte_ptr = + reinterpret_cast(access_ptr + access_offset); + + return reinterpret_cast(access_byte_ptr + byte_offset_); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iteration_contiguous_; + + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) + return *this; + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + iteration_contiguous_ = 0; + ++iteration_strided_; + + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + iteration_strided_ = 0; + + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + + RegularTileAccessIterator prev(*this); + + this->operator++(); + + return prev; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + + add_pointer_offset( + coord.contiguous() * Shape::kContiguous + + coord.strided() * Shape::kStrided * stride_ * Layout::kElementsPerAccess); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for column-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::ColumnMajorTensorOpMultiplicandCongruous64b, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for column-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajorTensorOpMultiplicandCongruous64b; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCongruous64b, + (kAdvanceRank == 0 ? 0 : 1), ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for row-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for row-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajorTensorOpMultiplicandCongruous64b; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCongruous64b, + (kAdvanceRank == 0 ? 1 : 0), ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +//////////////////////////////////////////////////////////////////////////////// +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for crosswise arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::TensorOpMultiplicand64bCrosswise, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorOpMultiplicand64bCrosswise; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + static_assert(ThreadMap::kThreads / 32 > 1, + "This tile iterator requires at least two warps."); + + /// Internal details made public to facilitate introspection + struct Detail { + /// This iterator is specialized for an access size that is 128 bits in + /// length. + static int const kAccessSizeInBits = 64; + + static_assert(sizeof_bits::value * + ThreadMap::kElementsPerAccess == + kAccessSizeInBits, + "This iterator requires a policy whose access size is 64b"); + + ///< Number of pointers - two pointers are needed if making more than 4 iterations along + ///< strided dimension + static int const kPointerCount = (ThreadMap::Iterations::kStrided > 4 ? 2 : 1); + }; + + /// Element type per access + using AccessType = Array; + + private: + // + // Data members + // + + /// Stride value + StrideIndex stride_; + + /// Internal pointer to first access of tile + AccessType *pointer_; + + /// Internal byte offset + Index byte_offset_[Detail::kPointerCount]; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + public: + + /// Construct a TileIterator with zero threadblock offset + MCTLASS_DEVICE + RegularTileAccessIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): + stride_(ref.stride(0) / ThreadMap::kElementsPerAccess) { + + layout::PitchLinearCoord thread_offset_base = ThreadMap::initial_offset(thread_id); + + // This is the offset of a thread within a threadblock tile for a specific + // pointer (units of elements) + layout::PitchLinearCoord thread_offset_in_threadblock_tile = thread_offset_base; + + // initialize pointer + pointer_ = reinterpret_cast(ref.data()); + + byte_offset_[0] = ref.offset(thread_offset_in_threadblock_tile) * sizeof(Element); + + if (Detail::kPointerCount == 2) { + byte_offset_[1] = byte_offset_[0] ^ 8; + } + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + + pointer_ += pointer_offset / ThreadMap::kElementsPerAccess; + } + + /// Returns a pointer + MCTLASS_DEVICE + AccessType *get() const { + + // Map the logical contiguous and strided access to the internal swizzled structure. + int uniform_offset = (iteration_strided_ & 0x3) * stride_ + (iteration_strided_ >> 3) * 16 + stride_ * ThreadMap::Delta::kContiguous * iteration_contiguous_; + + char *access_byte_ptr = reinterpret_cast(pointer_ + uniform_offset); + + int byte_offset; + + // This iterator may require two byte offsets if it must load more than 8 rows (or 2 iterations) + // in the strided dimension + if (Detail::kPointerCount == 2 && (iteration_strided_ & 0x4)) { + byte_offset = byte_offset_[1]; + } + else { + byte_offset = byte_offset_[0]; + } + + return reinterpret_cast(access_byte_ptr + byte_offset); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iteration_contiguous_; + + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) + return *this; + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + iteration_contiguous_ = 0; + ++iteration_strided_; + + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + iteration_strided_ = 0; + + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + + RegularTileAccessIterator prev(*this); + + this->operator++(); + + return prev; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + + add_pointer_offset(coord.strided() * Shape::kStrided + coord.contiguous() * Shape::kContiguous * stride_); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for column-major crosswise TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::ColumnMajorTensorOpMultiplicand64bCrosswise, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for column-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajorTensorOpMultiplicand64bCrosswise; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicand64bCrosswise, + (kAdvanceRank == 0 ? 0 : 1), ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for row-major crosswise TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for row-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajorTensorOpMultiplicand64bCrosswise; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicand64bCrosswise, + (kAdvanceRank == 0 ? 1 : 0), ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for congruous arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::TensorOpMultiplicandCongruous128b, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorOpMultiplicandCongruous128b; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + static_assert(ThreadMap::kThreads / 32 > 1, + "This tile iterator requires at least two warps."); + + /// Internal details made public to facilitate introspection + struct Detail { + /// This iterator is specialized for an access size that is 128 bits in + /// length. + static int const kAccessSizeInBits = 128; + + static_assert(sizeof_bits::value * + ThreadMap::kElementsPerAccess == + kAccessSizeInBits, + "This iterator requires a policy whose access size is 128b"); + + ///< Number of pointers + static int const kPointerCount = 1; + }; + + /// Element type per access + using AccessType = Array; + + private: + // + // Data members + // + + /// Stride value + StrideIndex stride_; + + /// Internal pointer to first access of tile + AccessType *pointer_; + + /// Internal byte offset + Index byte_offset_; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + public: + + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): + stride_(ref.stride(0) / Layout::kElementsPerAccess), + byte_offset_(0) { + + layout::PitchLinearCoord thread_offset_base = ThreadMap::initial_offset(thread_id); + + // This is the offset of a thread within a threadblock tile for a specific + // pointer (units of elements) + layout::PitchLinearCoord thread_offset_in_threadblock_tile = thread_offset_base; + + // initialize pointer + pointer_ = reinterpret_cast(ref.data() + ref.offset(thread_offset_in_threadblock_tile)); + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + + byte_offset_ += pointer_offset * sizeof(Element); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + + AccessType *access_ptr = pointer_; + + int access_offset = iteration_strided_ * ThreadMap::Delta::kStrided * stride_ + + iteration_contiguous_ * ThreadMap::Delta::kContiguous / + ThreadMap::kElementsPerAccess; + + char *access_byte_ptr = + reinterpret_cast(access_ptr + access_offset); + + return reinterpret_cast(access_byte_ptr + byte_offset_); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iteration_contiguous_; + + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) + return *this; + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + iteration_contiguous_ = 0; + ++iteration_strided_; + + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + iteration_strided_ = 0; + + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + + RegularTileAccessIterator prev(*this); + + this->operator++(); + + return prev; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + + add_pointer_offset( + coord.contiguous() * Shape::kContiguous + + coord.strided() * Shape::kStrided * stride_ * Layout::kElementsPerAccess); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for column-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::ColumnMajorTensorOpMultiplicandCongruous128b, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for column-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajorTensorOpMultiplicandCongruous128b; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCongruous128b, + (kAdvanceRank == 0 ? 0 : 1), ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for row-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for row-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajorTensorOpMultiplicandCongruous128b; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCongruous128b, + (kAdvanceRank == 0 ? 1 : 0), ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): + iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for congruous arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::TensorOpMultiplicandCrosswise128x4, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::TensorOpMultiplicandCrosswise128x4; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + static_assert(ThreadMap::kThreads / 32 > 1, + "This tile iterator requires at least two warps."); + + /// Internal details made public to facilitate introspection + struct Detail { + /// This iterator is specialized for an access size that is 128 bits in + /// length. + static int const kAccessSizeInBits = 128; + + static_assert(sizeof_bits::value * + ThreadMap::kElementsPerAccess == + kAccessSizeInBits, + "This iterator requires a policy whose access size is 128b"); + + ///< Number of pointers + static int const kPointerCount = 1; + }; + + + // static_assert(!(ThreadMap::Iterations::kStrided % 2), "This iterator requires at least two iterations along the strided dimension"); + + /// Element type per access + using AccessType = Array; + + private: + // + // Data members + // + + /// Stride value + StrideIndex stride_; + + /// Internal pointer to first access of tile + AccessType *pointer_; + + /// Internal byte offset + Index byte_offset_; + + /// Iteration in the contiguous dimension + int iteration_contiguous_; + + /// Iteration in the strided dimension + int iteration_strided_; + + public: + + /// Construct a TileIterator with zero threadblock offset + MCTLASS_DEVICE + RegularTileAccessIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): + stride_(ref.stride(0) / Layout::kElementsPerAccess), + byte_offset_(0) { + + layout::PitchLinearCoord thread_offset_base = ThreadMap::initial_offset(thread_id); + + // This is the offset of a thread within a threadblock tile for a specific + // pointer (units of elements) + layout::PitchLinearCoord thread_offset_in_threadblock_tile = thread_offset_base; + + // initialize pointer + pointer_ = reinterpret_cast(ref.data() + ref.offset(thread_offset_in_threadblock_tile)); + + set_iteration_index(0); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + + iteration_contiguous_ = index % ThreadMap::Iterations::kContiguous; + iteration_strided_ = index / ThreadMap::Iterations::kContiguous; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + + byte_offset_ += pointer_offset * sizeof(Element); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + + AccessType *access_ptr = pointer_; + + int offset_c = (iteration_contiguous_ * ThreadMap::Delta::kContiguous + (iteration_strided_ & 1) * 2); + int offset_s = (iteration_strided_ / 2) * 8; + + int access_offset = offset_c * stride_ + offset_s; + + char *access_byte_ptr = + reinterpret_cast(access_ptr + access_offset); + + return reinterpret_cast(access_byte_ptr + byte_offset_); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iteration_contiguous_; + + if (iteration_contiguous_ < ThreadMap::Iterations::kContiguous) + return *this; + + // Enter here only if (iteration_contiguous_ == + // ThreadMap::Iteration::kContiguous) + iteration_contiguous_ = 0; + ++iteration_strided_; + + if (iteration_strided_ < ThreadMap::Iterations::kStrided) { + return *this; + } + + // Enter here only if (iteration_stride_ == ThreadMap::Iteration::kStrided) + // which means we enter the next tile. + iteration_strided_ = 0; + + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + + RegularTileAccessIterator prev(*this); + + this->operator++(); + + return prev; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + + add_pointer_offset( + coord.contiguous() * Shape::kContiguous * stride_ + + coord.strided() * Shape::kStrided * Layout::kElementsPerAccess); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for column-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator< + Shape_, Element_, + layout::ColumnMajorTensorOpMultiplicandCrosswise128x4, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for column-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajorTensorOpMultiplicandCrosswise128x4; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCrosswise128x4, + (kAdvanceRank == 0 ? 0 : 1), ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for row-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileAccessIterator { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for row-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajorTensorOpMultiplicandCrosswise128x4; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileAccessIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCrosswise128x4, + (kAdvanceRank == 0 ? 1 : 0), ThreadMap_>; + + using AccessType = typename UnderlyingIterator::AccessType; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileAccessIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): + iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { iterator_.set_iteration_index(index); } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return reinterpret_cast(iterator_.get()); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileAccessIterator operator++(int) { + RegularTileAccessIterator prev(*this); + ++iterator_; + + return prev; + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator.h b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator.h new file mode 100644 index 0000000..66f8f02 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator.h @@ -0,0 +1,62 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing storing of tiles from pitch-linear rank=2 tensors. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +template < + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + int Alignment = sizeof_bits::value * ThreadMap::kElementsPerAccess / 8 +> +class RegularTileIterator; + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h new file mode 100644 index 0000000..2cd2e39 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_pitch_linear.h @@ -0,0 +1,552 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of tiles from pitch-linear rank=2 tensors. + + This iterator uses masks to guard out-of-bounds accesses and visits the last "residue" tile + first, with the objective of minimizing predicate mask updates during steady-state operation. + + A precomputed "Params" object minimizes the amount of state that must be stored in registers, + and integer addition is used to advance the pointer through memory. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" + +#include "regular_tile_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Regular tile iterator specialized for pitch-linear. This one is used by 2-stage SIMT kernels +/// and sparse tensor core meta data. +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator { +public: + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using Fragment = Array; + + using AccessType = AlignedArray; + + static_assert(kAdvanceRank == 0 || kAdvanceRank == 1, + "Advance rank may only be along the contiguous or strided dimensions."); + +private: + + // + // Types + // + + // + // Data members + // + + /// Pointer to memory + uint8_t *pointer_; + + /// Stride quantity + StrideIndex stride_; + + /// Amount to increment pointer along strided dimension + Index increment_strided_; + + /// Amount to advance pointer between tiles + Index increment_advance_; + +public: + + MCTLASS_DEVICE + RegularTileIterator(): pointer_(nullptr), increment_strided_(0), increment_advance_(0) { } + + MCTLASS_DEVICE + RegularTileIterator( + TensorRef const &ref, + int thread_idx + ): + pointer_(reinterpret_cast(ref.data()) + (ref.offset(ThreadMap::initial_offset(thread_idx)) * sizeof_bits::value / 8)) { + + stride_ = ref.stride()[0]; + increment_strided_ = (ref.stride()[0] * sizeof_bits::value) * ThreadMap::Delta::kStrided / 8; + + increment_advance_ = + (kAdvanceRank == 0 ? + Shape::kContiguous * sizeof_bits::value / 8 : + Shape::kStrided * (ref.stride()[0] * sizeof_bits::value / 8)); + } + + /// Loads a fragment + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + AccessType *frag_ptr = reinterpret_cast(&frag); + uint8_t const *byte_pointer = pointer_ + pointer_offset * sizeof_bits::value / 8; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + AccessType const *access_ptr = reinterpret_cast(byte_pointer); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int idx = c + s * ThreadMap::Iterations::kContiguous; + frag_ptr[idx] = access_ptr[c * ThreadMap::Delta::kContiguous / + ThreadMap::kElementsPerAccess]; + } + + if (s + 1 < ThreadMap::Iterations::kStrided) { + byte_pointer += increment_strided_; + } + } + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load(Fragment &frag, TensorCoord const & tile_offset) { + load_with_pointer_offset( + frag, + tile_offset.contiguous() * Shape::kContiguous / ThreadMap::kElementsPerAccess + + tile_offset.strided() * Shape::kStrided * stride_ + ); + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + uint8_t *byte_pointer = pointer_ + pointer_offset * sizeof_bits::value / 8; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + AccessType *access_ptr = reinterpret_cast(byte_pointer); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int idx = c + s * ThreadMap::Iterations::kContiguous; + access_ptr[c * ThreadMap::Delta::kContiguous / + ThreadMap::kElementsPerAccess] = frag_ptr[idx]; + } + + if (s + 1 < ThreadMap::Iterations::kStrided) { + byte_pointer += increment_strided_; + } + } + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store(Fragment const &frag, TensorCoord const & tile_offset) { + store_with_pointer_offset( + frag, + tile_offset.contiguous() * Shape::kContiguous + tile_offset.strided() * Shape::kStrided * stride_ + ); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } + + /// Advances the pointer + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + pointer_ += increment_advance_; + return *this; + } + + /// Advances the pointer + MCTLASS_HOST_DEVICE + RegularTileIterator &operator--() { + pointer_ -= increment_advance_; + return *this; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset; + } + + /// Adds a tile offset in the unit of tile. + /// In GEMM/Conv implementation, this is used to move in the k dimension in the shared memory. + /// Below layouts are the shared memory layouts. Current SM50 SIMT kernels only use col major A and row major B. + /// For row major A operand, k dimension is contiguous dimension; + /// For col major A operand, k dimension is strided dimension; + /// For row major B operand, k dimension is strided dimension; + /// For col major B operand, k dimension is contiguous dimension. + /// Below two classes map col/row major to the pitch linear coordinates used + /// in this base class. + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + int offset = sizeof_bits::value * + (coord.contiguous() * Shape::kContiguous + coord.strided() * Shape::kStrided * stride_) / 8; + add_pointer_offset(offset); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { +#if 0 + AccessType *access_ptr = pointer_[iteration_strided_ & 1]; + int stride_idx = (iteration_strided_ & ~1); + + int access_offset = stride_idx * ThreadMap::Delta::kStrided * stride_ + + iteration_contiguous_ * ThreadMap::Delta::kContiguous / + ThreadMap::kElementsPerAccess; + + char *access_byte_ptr = + reinterpret_cast(access_ptr + access_offset); + return reinterpret_cast(access_byte_ptr + byte_offset_); +#endif + return reinterpret_cast(pointer_); + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Regular tile iterator specialized for row major +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator { +public: + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using Fragment = Array; + + using Underlying = RegularTileIterator< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + kAlignment + >; + + using AccessType = typename Underlying::AccessType; + + static_assert(kAdvanceRank == 0 || kAdvanceRank == 1, + "Advance rank may only be along the row or column dimensions."); + +private: + + Underlying iterator_; + +public: + + MCTLASS_DEVICE + RegularTileIterator() { } + + MCTLASS_DEVICE + RegularTileIterator( + TensorRef const &ref, + int thread_idx + ): + iterator_({ref.data(), ref.stride()}, thread_idx) { + + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load(Fragment &frag, TensorCoord const & tile_offset) { + iterator_.load_with_pointer_offset(frag, {tile_offset.column(), tile_offset.row()}); + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load(Fragment &frag) { + iterator_.load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store(Fragment const &frag, TensorCoord const & tile_offset) { + iterator_.store_with_pointer_offset(frag, {tile_offset.column(), tile_offset.row()}); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + iterator_.store_with_pointer_offset(frag, 0); + } + + /// Advances the pointer + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances the pointer + MCTLASS_HOST_DEVICE + RegularTileIterator &operator--() { + --iterator_; + return *this; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return iterator_.get(); + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Regular tile iterator specialized for pitch-linear +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator { +public: + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajor; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using Fragment = Array; + + using Underlying = RegularTileIterator< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap + >; + + using AccessType = typename Underlying::AccessType; + + static_assert(kAdvanceRank == 0 || kAdvanceRank == 1, + "Advance rank may only be along the row or column dimensions."); + +private: + + Underlying iterator_; + +public: + + MCTLASS_DEVICE + RegularTileIterator() { } + + MCTLASS_DEVICE + RegularTileIterator( + TensorRef const &ref, + int thread_idx + ): + iterator_({ref.data(), ref.stride()}, thread_idx) { + + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load(Fragment &frag, TensorCoord const & tile_offset) { + iterator_.load_with_pointer_offset(frag, {tile_offset.row(), tile_offset.column()}); + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load(Fragment &frag) { + iterator_.load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store(Fragment const &frag, TensorCoord const & tile_offset) { + iterator_.store_with_pointer_offset(frag, {tile_offset.row(), tile_offset.column()}); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + iterator_.store_with_pointer_offset(frag, 0); + } + + /// Advances the pointer + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances the pointer + MCTLASS_HOST_DEVICE + RegularTileIterator &operator--() { + --iterator_; + return *this; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Overrides the internal iteration index + MCTLASS_HOST_DEVICE + void set_iteration_index(int index) { + } + + /// Returns a pointer + MCTLASS_HOST_DEVICE + AccessType *get() const { + return iterator_.get(); + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_pitch_linear_2dthreadtile.h b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_pitch_linear_2dthreadtile.h new file mode 100644 index 0000000..15ea280 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_pitch_linear_2dthreadtile.h @@ -0,0 +1,509 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of tiles from pitch-linear rank=2 tensors. + + This iterator uses masks to guard out-of-bounds accesses and visits the last "residue" tile + first, with the objective of minimizing predicate mask updates during steady-state operation. + + A precomputed "Params" object minimizes the amount of state that must be stored in registers, + and integer addition is used to advance the pointer through memory. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/pitch_linear.h" + +#include "regular_tile_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// +template < + typename Shape, + typename Element, + typename Layout, + int AdvanceRank, + typename ThreadMap, + int Alignment = sizeof_bits::value * ThreadMap::kElementsPerAccess / 8 +> +class RegularTileIterator2dThreadTile; + + +/// Regular tile iterator specialized for pitch-linear + 2d thread-tiled threadmapping +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator2dThreadTile { +public: + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::PitchLinear; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using Fragment = Array; + + static_assert(kAdvanceRank == 0 || kAdvanceRank == 1, + "Advance rank may only be along the contiguous or strided dimensions."); + +private: + + // + // Types + // + + using AccessType = AlignedArray; + + // + // Data members + // + + /// Pointer to memory + uint8_t *pointer_; + + /// Stride quantity + StrideIndex stride_; + + /// Amount to increment pointer along strided dimension + LongIndex increment_strided_; + + /// Amount to advance pointer between tiles + LongIndex increment_advance_; + +public: + + MCTLASS_DEVICE + RegularTileIterator2dThreadTile(): pointer_(nullptr), increment_strided_(0), increment_advance_(0) { } + + MCTLASS_DEVICE + RegularTileIterator2dThreadTile( + TensorRef const &ref, + int thread_idx, + int interleave + ){ + + TensorCoord t = ThreadMap::initial_offset(thread_idx); + long int offset = t[0] * interleave + t[1] * ref.stride()[0]/interleave; + pointer_ = reinterpret_cast(ref.data() + offset); + + stride_ = ref.stride()[0] / interleave; + increment_strided_ = (ref.stride()[0] * sizeof_bits::value / 8) * ThreadMap::Delta::kStrided / interleave; + + increment_advance_ = + (kAdvanceRank == 0 ? + Shape::kContiguous * sizeof_bits::value / 8 : + Shape::kStrided * (ref.stride()[0] * sizeof_bits::value / 8) / interleave); + } + + /// Loads a fragment + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + AccessType *frag_ptr = reinterpret_cast(&frag); + uint8_t const *byte_pointer = pointer_ + pointer_offset * sizeof_bits::value / 8; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + AccessType const *access_ptr = reinterpret_cast(byte_pointer); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int idx = c + s * ThreadMap::Iterations::kContiguous; + frag_ptr[idx] = access_ptr[c * ThreadMap::Delta::kContiguous / ThreadMap::ThreadAccessShape::kStrided]; + } + + if (s + 1 < ThreadMap::Iterations::kStrided) { + byte_pointer += increment_strided_; + } + } + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load(Fragment &frag, TensorCoord const & tile_offset) { + load_with_pointer_offset( + frag, + tile_offset.contiguous() * Shape::kContiguous / ThreadMap::kElementsPerAccess + + tile_offset.strided() * Shape::kStrided * stride_ + ); + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + uint8_t *byte_pointer = pointer_ + pointer_offset * sizeof_bits::value / 8; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + AccessType *access_ptr = reinterpret_cast(byte_pointer); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int idx = c + s * ThreadMap::Iterations::kContiguous; + access_ptr[c * ThreadMap::Delta::kContiguous / ThreadMap::ThreadAccessShape::kStrided] = frag_ptr[idx]; + } + + if (s + 1 < ThreadMap::Iterations::kStrided) { + byte_pointer += increment_strided_; + } + } + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store(Fragment const &frag, TensorCoord const & tile_offset) { + store_with_pointer_offset( + frag, + tile_offset.contiguous() * Shape::kContiguous + tile_offset.strided() * Shape::kStrided * stride_ + ); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } + + /// Advances the pointer + MCTLASS_HOST_DEVICE + RegularTileIterator2dThreadTile &operator++() { + pointer_ += increment_advance_; + return *this; + } + + /// Advances the pointer + MCTLASS_HOST_DEVICE + RegularTileIterator2dThreadTile &operator--() { + pointer_ -= increment_advance_; + return *this; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + pointer_ += pointer_offset; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + int offset = sizeof_bits::value * + (coord.contiguous() * Shape::kContiguous + coord.strided() * Shape::kStrided * stride_) / 8; + add_pointer_offset(offset); + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Regular tile iterator specialized for interleaved layout + 2d thread-tiled threadmapping +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator2dThreadTile, AdvanceRank, ThreadMap_, Alignment> { +public: + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajorInterleaved<4>; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using Fragment = Array; + + using Underlying = RegularTileIterator2dThreadTile< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap, + kAlignment + >; + + static_assert(kAdvanceRank == 0 || kAdvanceRank == 1, + "Advance rank may only be along the row or column dimensions."); + +private: + + Underlying iterator_; + +public: + + MCTLASS_DEVICE + RegularTileIterator2dThreadTile() { } + + MCTLASS_DEVICE + RegularTileIterator2dThreadTile( + TensorRef const &ref, + int thread_idx + ): + iterator_({ref.data(), ref.stride()}, thread_idx, 4) { + + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load(Fragment &frag, TensorCoord const & tile_offset) { + iterator_.load_with_pointer_offset(frag, {tile_offset.column(), tile_offset.row()}); + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load(Fragment &frag) { + iterator_.load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store(Fragment const &frag, TensorCoord const & tile_offset) { + iterator_.store_with_pointer_offset(frag, {tile_offset.column(), tile_offset.row()}); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + iterator_.store_with_pointer_offset(frag, 0); + } + + /// Advances the pointer + MCTLASS_HOST_DEVICE + RegularTileIterator2dThreadTile &operator++() { + ++iterator_; + return *this; + } + + /// Advances the pointer + MCTLASS_HOST_DEVICE + RegularTileIterator2dThreadTile &operator--() { + --iterator_; + return *this; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Regular tile iterator specialized for interleaved layout + 2d thread-tiled threadmapping +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator2dThreadTile, AdvanceRank, ThreadMap_, Alignment> { +public: + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajorInterleaved<4>; + static int const kAdvanceRank = AdvanceRank; + using ThreadMap = ThreadMap_; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using Fragment = Array; + using PitchLinearThreadMap = PitchLinearStripminedThreadMap< layout::PitchLinearShape, + ThreadMap::kThreads, ThreadMap::ThreadAccessShape::kCount >; + + + using Underlying = RegularTileIterator2dThreadTile< + layout::PitchLinearShape, + Element, + layout::PitchLinear, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap + >; + + static_assert(kAdvanceRank == 0 || kAdvanceRank == 1, + "Advance rank may only be along the row or column dimensions."); + +private: + + Underlying iterator_; + +public: + + MCTLASS_DEVICE + RegularTileIterator2dThreadTile() { } + + MCTLASS_DEVICE + RegularTileIterator2dThreadTile( + TensorRef const &ref, + int thread_idx + ): + iterator_({ref.data(), ref.stride()}, thread_idx, 4) { + + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load(Fragment &frag, TensorCoord const & tile_offset) { + iterator_.load_with_pointer_offset(frag, {tile_offset.row(), tile_offset.column()}); + } + + /// Loads a fragment + MCTLASS_HOST_DEVICE + void load(Fragment &frag) { + iterator_.load_with_pointer_offset(frag, 0); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store(Fragment const &frag, TensorCoord const & tile_offset) { + iterator_.store_with_pointer_offset(frag, {tile_offset.row(), tile_offset.column()}); + } + + /// Stores a fragment + MCTLASS_HOST_DEVICE + void store(Fragment const &frag) { + iterator_.store_with_pointer_offset(frag, 0); + } + + /// Advances the pointer + MCTLASS_HOST_DEVICE + RegularTileIterator2dThreadTile &operator++() { + ++iterator_; + return *this; + } + + /// Advances the pointer + MCTLASS_HOST_DEVICE + RegularTileIterator2dThreadTile &operator--() { + --iterator_; + return *this; + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_tensor_op.h b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_tensor_op.h new file mode 100644 index 0000000..055d8ee --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_tensor_op.h @@ -0,0 +1,1107 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing storing of tiles from pitch-linear rank=2 tensors. +*/ + +#pragma once + +#include "mctlass/transform/threadblock/regular_tile_iterator.h" +#include "mctlass/transform/threadblock/regular_tile_access_iterator_tensor_op.h" + +//////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for congruous arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileIterator< + Shape_, Element_, + layout::TensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element_))>, + AdvanceRank, ThreadMap_, Alignment> { + public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = + layout::TensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element))>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Internal details made public to facilitate introspection + struct Detail { + + /// This iterator is specialized for an access size that is 128 bits in length. + static int const kAccessSizeInBits = 128; + + static_assert( + sizeof_bits::value * ThreadMap::kElementsPerAccess == kAccessSizeInBits, + "This iterator requires a policy whose access size is 128bs"); + }; + +private: + + /// Element type per access + using AccessType = Array; + +public: + + /// Fragment object to be loaded or stored + using Fragment = Array; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = RegularTileAccessIterator; + +private: + + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + +public: + + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : address_iterator_(ref, thread_id) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + address_iterator_.add_tile_offset({0, 1}); + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + RegularTileIterator prev(*this); + this->operator++(); + + return prev; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + address_iterator_.add_tile_offset(coord); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + load_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_byte_offset(Fragment &frag, Index byte_offset) { + address_iterator_.set_iteration_index(0); + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + int access_idx = c + s * ThreadMap::Iterations::kContiguous; + + char const *byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType const *access_ptr = reinterpret_cast(byte_ptr); + + frag_ptr[access_idx] = *access_ptr; + ++address_iterator_; + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + store_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, Index byte_offset) { + address_iterator_.set_iteration_index(0); + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + int access_idx = c + s * ThreadMap::Iterations::kContiguous; + + char *byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType *access_ptr = reinterpret_cast(byte_ptr); + + *access_ptr = frag_ptr[access_idx]; + ++address_iterator_; + } + } + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_byte_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for column-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileIterator< + Shape_, Element_, + layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element_))>, + AdvanceRank, ThreadMap_, Alignment> { + public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for column-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element))>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element))>, + (kAdvanceRank == 0 ? 0 : 1), ThreadMap_>; + + public: + + /// Fragment object to be loaded or stored + using Fragment = Array; + +private: + + /// Underlying iterator + UnderlyingIterator iterator_; + +public: + + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): iterator_({ref.data(), ref.stride()}, thread_id) { + + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + RegularTileIterator prev(*this); + ++iterator_; + + return prev; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, + Index pointer_offset) { + + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for row-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileIterator< + Shape_, Element_, + layout::RowMajorTensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element_))>, + AdvanceRank, ThreadMap_, Alignment> { + public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for row-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajorTensorOpMultiplicandCongruous< + sizeof_bits::value, int(128 / sizeof(Element))>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCongruous::value, + int(128 / sizeof(Element))>, + (kAdvanceRank == 0 ? 1 : 0), ThreadMap_>; + + public: + + /// Fragment object to be loaded or stored + using Fragment = Array; + +private: + + /// Underlying iterator + UnderlyingIterator iterator_; + +public: + + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): iterator_({ref.data(), ref.stride()}, thread_id) { + + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + + RegularTileIterator prev(*this); + ++iterator_; + + return prev; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, + Index pointer_offset) { + + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for crosswise arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileIterator::value, Crosswise>, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = + layout::TensorOpMultiplicandCrosswise::value, + Crosswise>; + + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Internal details made public to facilitate introspection + struct Detail { + /// This iterator is specialized for an access size that is 128 bits in + /// length. + static int const kAccessSizeInBits = 128; + + static_assert(sizeof_bits::value * ThreadMap::kElementsPerAccess == + kAccessSizeInBits, + "This iterator requires a policy whose access size is 128bs"); + }; + + private: + /// Element type per access + using AccessType = Array; + + public: + /// Fragment object to be loaded or stored + using Fragment = + Array; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = RegularTileAccessIterator; + + private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : address_iterator_(ref, thread_id) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + address_iterator_.add_tile_offset({1, 0}); + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + RegularTileIterator prev(*this); + this->operator++(); + + return prev; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + address_iterator_.add_tile_offset(coord); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + address_iterator_.set_iteration_index(0); + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + int access_idx = c + s * ThreadMap::Iterations::kContiguous; + frag_ptr[access_idx] = *(address_iterator_.get() + pointer_offset); + ++address_iterator_; + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + store_with_byte_offset(frag, pointer_offset * sizeof_bits::value / 8); + } + + MCTLASS_DEVICE + void store_with_byte_offset(Fragment const &frag, Index byte_offset) { + address_iterator_.set_iteration_index(0); + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + int access_idx = c + s * ThreadMap::Iterations::kContiguous; + + char *byte_ptr = reinterpret_cast(address_iterator_.get()) + byte_offset; + AccessType *access_ptr = reinterpret_cast(byte_ptr); + + *access_ptr = frag_ptr[access_idx]; + ++address_iterator_; + } + } + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for column-major crosswise TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileIterator::value, Crosswise>, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for column-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Crosswise>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCrosswise::value, + Crosswise>, + (kAdvanceRank == 0 ? 0 : 1), ThreadMap_>; + + public: + /// Fragment object to be loaded or stored + using Fragment = Array; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + RegularTileIterator prev(*this); + ++iterator_; + + return prev; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for row-major crosswise TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileIterator::value, Crosswise>, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for row-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajorTensorOpMultiplicandCrosswise< + sizeof_bits::value, Crosswise>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileIterator< + layout::PitchLinearShape, Element, + layout::TensorOpMultiplicandCrosswise::value, + Crosswise>, + (kAdvanceRank == 0 ? 1 : 0), ThreadMap_>; + + public: + /// Fragment object to be loaded or stored + using Fragment = Array; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + RegularTileIterator prev(*this); + ++iterator_; + + return prev; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for k interleaved arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template +class RegularTileIterator< + Shape_, Element_, + layout::TensorOpMultiplicandRowMajorInterleaved::value, + InterleavedK>, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = + layout::TensorOpMultiplicandRowMajorInterleaved::value, + InterleavedK>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Internal details made public to facilitate introspection + struct Detail { + /// This iterator is specialized for an access size that is 128 bits in + /// length. + static int const kAccessSizeInBits = 128; + + static_assert(sizeof_bits::value * ThreadMap::kElementsPerAccess == + kAccessSizeInBits, + "This iterator requires a policy whose access size is 128bs"); + }; + + private: + + /// Element type per access + using AccessType = Array; + + public: + /// Fragment object to be loaded or stored + using Fragment = + Array; + + /// Underlying iterator to compute the addresses + using TileAccessIterator = RegularTileAccessIterator; + + private: + // + // Data members + // + + /// Data member to the tile access iterator + TileAccessIterator address_iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : address_iterator_(ref, thread_id) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + address_iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + address_iterator_.add_pointer_offset(Shape::kCount); + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + RegularTileIterator prev(*this); + this->operator++(); + + return prev; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + address_iterator_.add_pointer_offset(coord.contiguous() * Shape::kCount); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + address_iterator_.set_iteration_index(0); + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + int access_idx = c + s * ThreadMap::Iterations::kContiguous; + frag_ptr[access_idx] = *(address_iterator_.get() + pointer_offset); + ++address_iterator_; + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + AccessType const *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + int access_idx = c + s * ThreadMap::Iterations::kContiguous; + *(address_iterator_.get() + pointer_offset) = frag_ptr[access_idx]; + ++address_iterator_; + } + } + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } +}; + +//////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for k interleaved arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// + +template +class RegularTileIterator< + Shape_, Element_, + layout::TensorOpMultiplicandColumnMajorInterleaved::value, + InterleavedK>, + AdvanceRank, ThreadMap_, Alignment> { + + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = + layout::TensorOpMultiplicandColumnMajorInterleaved::value, + InterleavedK>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileIterator< + mctlass::MatrixShape, + Element, + layout::TensorOpMultiplicandRowMajorInterleaved::value, InterleavedK>, + (kAdvanceRank == 1 ? 0 : 1), + ThreadMap + >; + + public: + /// Fragment object to be loaded or stored + using Fragment = Array; + + private: + + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + RegularTileIterator prev(*this); + ++iterator_; + + return prev; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.strided(), coord.contiguous()}); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_tensor_op_sm70.h b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_tensor_op_sm70.h new file mode 100644 index 0000000..c56d472 --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/regular_tile_iterator_tensor_op_sm70.h @@ -0,0 +1,1460 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Templates implementing loading of tiles from pitch-linear rank=2 tensors. + + This iterator uses masks to guard out-of-bounds accesses and visits the last "residue" tile + first, with the objective of minimizing predicate mask updates during steady-state operation. + + A precomputed "Params" object minimizes the amount of state that must be stored in registers, + and integer addition is used to advance the pointer through memory. +*/ + +#pragma once + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/matrix_coord.h" +#include "mctlass/tensor_ref.h" +#include "mctlass/layout/pitch_linear.h" +#include "mctlass/layout/tensor_op_multiplicand_sm70.h" + +#include "mctlass/transform/threadblock/regular_tile_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tile iterator specialized for congruous arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator< + Shape_, + Element_, + layout::VoltaTensorOpMultiplicandCongruous::value>, + AdvanceRank, + ThreadMap_, + Alignment> { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::VoltaTensorOpMultiplicandCongruous::value>; + static int const kAdvanceRank = AdvanceRank; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Internal details made public to facilitate introspection + struct Detail { + + /// This iterator is specialized for an access size that is 128 bits in length. + static int const kAccessSizeInBits = 128; + + static_assert( + sizeof_bits::value * ThreadMap::kElementsPerAccess == kAccessSizeInBits, + "This iterator requires a policy whose access size is 128bs"); + + ///< Number of pointers + static int const kPointerCount = (ThreadMap::Iterations::kStrided > 1 ? 2 : 1); + }; + + +private: + + /// Element type per access + using AccessType = Array; + +public: + + /// Fragment object to be loaded or stored + using Fragment = Array; + +private: + + // + // Data members + // + + /// Stride value + StrideIndex stride_; + + /// Internal pointer to first access of tile + AccessType * pointer_[Detail::kPointerCount]; + + /// Internal byte offset + Index byte_offset_; + +public: + + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): stride_(ref.stride(0) / Layout::kElementsPerAccess), byte_offset_(0) { + + layout::PitchLinearCoord thread_offset_base = ThreadMap::initial_offset(thread_id); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Detail::kPointerCount; ++i) { + + // This is the offset of a thread within a threadblock tile for a specific pointer + // (units of elements) + layout::PitchLinearCoord thread_offset_in_threadblock_tile = + thread_offset_base + layout::PitchLinearCoord{0, ThreadMap::Detail::WarpThreadArrangement::kStrided * i}; + + // initialize pointer + pointer_[i] = reinterpret_cast(ref.data() + ref.offset(thread_offset_in_threadblock_tile)); + } + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + + byte_offset_ += pointer_offset * sizeof(Element); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + + add_pointer_offset((kAdvanceRank ? Shape::kStrided * stride_ * Layout::kElementsPerAccess : Shape::kContiguous)); + + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + + RegularTileIterator prev(*this); + this->operator++(); + + return prev; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + add_pointer_offset( + coord.contiguous() * Shape::kContiguous / ThreadMap::kElementsPerAccess + + coord.strided() * Shape::kStrided * stride_ * Layout::kElementsPerAccess + ); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + Index vec_pointer_offset = pointer_offset / ThreadMap::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + AccessType *access_ptr = pointer_[s & 1]; + int stride_idx = (s & ~1); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int access_offset = stride_idx * ThreadMap::Delta::kStrided * stride_ + + c * ThreadMap::Delta::kContiguous / ThreadMap::kElementsPerAccess + + vec_pointer_offset; + + int access_idx = c + s * ThreadMap::Iterations::kContiguous; + + char const *access_byte_ptr = reinterpret_cast(access_ptr + access_offset); + + frag_ptr[access_idx] = *reinterpret_cast(access_byte_ptr + byte_offset_); + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, + Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + Index vec_pointer_offset = pointer_offset / ThreadMap::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + AccessType *access_ptr = pointer_[s & 1]; + int stride_idx = (s & ~1); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int access_offset = stride_idx * ThreadMap::Delta::kStrided * stride_ + + c * ThreadMap::Delta::kContiguous / ThreadMap::kElementsPerAccess + + vec_pointer_offset; + + int access_idx = c + s * ThreadMap::Iterations::kContiguous; + + char *access_byte_ptr = reinterpret_cast(access_ptr + access_offset); + + *reinterpret_cast(access_byte_ptr + byte_offset_) = frag_ptr[access_idx]; + } + } + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +// Tile Iterator specialized for column-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator< + Shape_, + Element_, + layout::ColumnMajorVoltaTensorOpMultiplicandCongruous::value>, + AdvanceRank, + ThreadMap_, + Alignment> { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for column-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajorVoltaTensorOpMultiplicandCongruous::value>; + static int const kAdvanceRank = AdvanceRank; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileIterator< + layout::PitchLinearShape, + Element, + layout::VoltaTensorOpMultiplicandCongruous::value>, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap_>; + +public: + + /// Fragment object to be loaded or stored + using Fragment = Array; + +private: + + /// Underlying iterator + UnderlyingIterator iterator_; + +public: + + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): iterator_({ref.data(), ref.stride()}, thread_id) { + + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + + RegularTileIterator prev(*this); + ++iterator_; + + return prev; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, + Index pointer_offset) { + + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for row-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator< + Shape_, + Element_, + layout::RowMajorVoltaTensorOpMultiplicandCongruous::value>, + AdvanceRank, + ThreadMap_, + Alignment> { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for row-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajorVoltaTensorOpMultiplicandCongruous::value>; + static int const kAdvanceRank = AdvanceRank; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileIterator< + layout::PitchLinearShape, + Element, + layout::VoltaTensorOpMultiplicandCongruous::value>, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap_>; + +public: + + /// Fragment object to be loaded or stored + using Fragment = Array; + +private: + + /// Underlying iterator + UnderlyingIterator iterator_; + +public: + + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): iterator_({ref.data(), ref.stride()}, thread_id) { + + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + + RegularTileIterator prev(*this); + ++iterator_; + + return prev; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, + Index pointer_offset) { + + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; +/// Tile iterator specialized for congruous arrangements for TensorOps +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator< + Shape_, + Element_, + layout::VoltaTensorOpMultiplicandBCongruous::value>, + AdvanceRank, + ThreadMap_, + Alignment> { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::VoltaTensorOpMultiplicandBCongruous::value>; + static int const kAdvanceRank = AdvanceRank; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + using StrideIndex = typename Layout::Stride::Index; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Internal details made public to facilitate introspection + struct Detail { + + /// This iterator is specialized for an access size that is 128 bits in length. + static int const kAccessSizeInBits = 128; + + static_assert( + sizeof_bits::value * ThreadMap::kElementsPerAccess == kAccessSizeInBits, + "This iterator requires a policy whose access size is 128bs"); + + ///< Number of pointers + static int const kPointerCount = (ThreadMap::Iterations::kStrided > 1 ? 2 : 1); + }; + + +private: + + /// Element type per access + using AccessType = Array; + +public: + + /// Fragment object to be loaded or stored + using Fragment = Array; + +private: + + // + // Data members + // + + /// Stride value + StrideIndex stride_; + + /// Internal pointer to first access of tile + AccessType * pointer_[Detail::kPointerCount]; + + /// Internal byte offset + Index byte_offset_; + +public: + + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): stride_(ref.stride(0) / Layout::kElementsPerAccess), byte_offset_(0) { + + layout::PitchLinearCoord thread_offset_base = ThreadMap::initial_offset(thread_id); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Detail::kPointerCount; ++i) { + + // This is the offset of a thread within a threadblock tile for a specific pointer + // (units of elements) + layout::PitchLinearCoord thread_offset_in_threadblock_tile = + thread_offset_base + layout::PitchLinearCoord{0, ThreadMap::Detail::WarpThreadArrangement::kStrided * i}; + + // initialize pointer + pointer_[i] = reinterpret_cast(ref.data() + ref.offset(thread_offset_in_threadblock_tile)); + } + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + + byte_offset_ += pointer_offset * sizeof(Element); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + + add_pointer_offset((kAdvanceRank ? Shape::kStrided * stride_ * Layout::kElementsPerAccess : Shape::kContiguous)); + + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + + RegularTileIterator prev(*this); + this->operator++(); + + return prev; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + add_pointer_offset( + coord.contiguous() * Shape::kContiguous / ThreadMap::kElementsPerAccess + + coord.strided() * Shape::kStrided * stride_ * Layout::kElementsPerAccess + ); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + Index vec_pointer_offset = pointer_offset / ThreadMap::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + AccessType *access_ptr = pointer_[s & 1]; + int stride_idx = (s & ~1); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int access_offset = stride_idx * ThreadMap::Delta::kStrided * stride_ + + c * ThreadMap::Delta::kContiguous / ThreadMap::kElementsPerAccess + + vec_pointer_offset; + + int access_idx = c + s * ThreadMap::Iterations::kContiguous; + + char const *access_byte_ptr = reinterpret_cast(access_ptr + access_offset); + + frag_ptr[access_idx] = *reinterpret_cast(access_byte_ptr + byte_offset_); + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, + Index pointer_offset) { + + AccessType const *frag_ptr = reinterpret_cast(&frag); + + Index vec_pointer_offset = pointer_offset / ThreadMap::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + AccessType *access_ptr = pointer_[s & 1]; + int stride_idx = (s & ~1); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + int access_offset = stride_idx * ThreadMap::Delta::kStrided * stride_ + + c * ThreadMap::Delta::kContiguous / ThreadMap::kElementsPerAccess + + vec_pointer_offset; + + int access_idx = c + s * ThreadMap::Iterations::kContiguous; + + char *access_byte_ptr = reinterpret_cast(access_ptr + access_offset); + + *reinterpret_cast(access_byte_ptr + byte_offset_) = frag_ptr[access_idx]; + } + } + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for column-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator< + Shape_, + Element_, + layout::ColumnMajorVoltaTensorOpMultiplicandBCongruous::value>, + AdvanceRank, + ThreadMap_, + Alignment> { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for column-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajorVoltaTensorOpMultiplicandBCongruous::value>; + static int const kAdvanceRank = AdvanceRank; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileIterator< + layout::PitchLinearShape, + Element, + layout::VoltaTensorOpMultiplicandBCongruous::value>, + (kAdvanceRank == 0 ? 0 : 1), + ThreadMap_>; + +public: + + /// Fragment object to be loaded or stored + using Fragment = Array; + +private: + + /// Underlying iterator + UnderlyingIterator iterator_; + +public: + + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): iterator_({ref.data(), ref.stride()}, thread_id) { + + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + + RegularTileIterator prev(*this); + ++iterator_; + + return prev; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, + Index pointer_offset) { + + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for row-major congruous TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator< + Shape_, + Element_, + layout::RowMajorVoltaTensorOpMultiplicandBCongruous::value>, + AdvanceRank, + ThreadMap_, + Alignment> { +public: + + static_assert(AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for row-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajorVoltaTensorOpMultiplicandBCongruous::value>; + static int const kAdvanceRank = AdvanceRank; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileIterator< + layout::PitchLinearShape, + Element, + layout::VoltaTensorOpMultiplicandBCongruous::value>, + (kAdvanceRank == 0 ? 1 : 0), + ThreadMap_>; + +public: + + /// Fragment object to be loaded or stored + using Fragment = Array; + +private: + + /// Underlying iterator + UnderlyingIterator iterator_; + +public: + + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator( + TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ): iterator_({ref.data(), ref.stride()}, thread_id) { + + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + + RegularTileIterator prev(*this); + ++iterator_; + + return prev; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + load_with_pointer_offset(frag, 0); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset( + Fragment const &frag, + Index pointer_offset) { + + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { + store_with_pointer_offset(frag, 0); + } +}; + + +/// Tile iterator specialized for crosswise arrangements for TensorOps. +/// +/// Volta TN SMEM layout is a little diffrent: +/// Crosseised elements will be stored in a line, while contiguous elements +/// sre stored in line-by-line. +/// Padding is used to reduce SMEM bank conflicts. +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator< + Shape_, Element_, + layout::VoltaTensorOpMultiplicandCrosswise::value, + Shape_::kContiguous>, + AdvanceRank, ThreadMap_, Alignment> { + + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for pitch-linear iterator may along advance along the " + "contiguous(rank=0) or strided(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = + layout::VoltaTensorOpMultiplicandCrosswise::value, + Shape::kContiguous>; + static int const kAdvanceRank = AdvanceRank; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Internal details made public to facilitate introspection + struct Detail { + + ///< Number of pointers + static int const kPointerCount = (ThreadMap::Iterations::kStrided > 1 ? 2 : 1); + + /// Iterations for the kElementsPerAccess of ThreadMap + static int const kIterarionsPerAccess = + ThreadMap::kElementsPerAccess / Layout::kElementsPerAccess; + + /// Contiguous elements per line + static int const kContiguousElementsPerLine = 4; + }; + + private: + /// Element type per access + using AccessType = Array; + + public: + /// Fragment object to be loaded or stored + using Fragment = + Array; + + private: + // + // Data members + // + + /// The crosswised elements will be stored in a line. + /// line_size is size of crosswised dimension plus padding. + /// in units of AccessType + Index line_size; + + /// Internal pointer to first access of tile + AccessType *pointer_[Detail::kPointerCount]; + + /// Internal byte offset + Index byte_offset_; + + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : line_size(ref.stride(0) * Detail::kContiguousElementsPerLine / Layout::kElementsPerAccess), + byte_offset_(0) { + + layout::PitchLinearCoord thread_offset_base = + ThreadMap::initial_offset(thread_id); + + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < Detail::kPointerCount; ++i) { + // This is the offset of a thread within a threadblock tile for a specific + // pointer (units of elements) + layout::PitchLinearCoord thread_offset_in_threadblock_tile = + thread_offset_base + + layout::PitchLinearCoord{ + 0, ThreadMap::Detail::WarpThreadArrangement::kStrided * i}; + + // initialize pointer + pointer_[i] = reinterpret_cast( + ref.data() + ref.offset(thread_offset_in_threadblock_tile)); + } + } + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + byte_offset_ += pointer_offset * sizeof(Element); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + // (Shape::kContiguous/Layout::kElementsPerAccess)* + // line_size * Layout::kElementsPerAccess + add_pointer_offset(Shape::kContiguous * line_size); + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + RegularTileIterator prev(*this); + this->operator++(); + + return prev; + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + add_pointer_offset((coord.contiguous() * (Shape::kContiguous / Layout::kElementsPerAccess) * + line_size + coord.strided() * Shape::kStrided) * + Layout::kElementsPerAccess); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + AccessType *frag_ptr = reinterpret_cast(&frag); + + Index vec_pointer_offset = pointer_offset / Layout::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + AccessType *access_ptr = pointer_[(s & 1) ^ (s / 2)]; + + access_ptr += 16 * (s / 2); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + + MCTLASS_PRAGMA_UNROLL + for(int i = 0; i < Detail::kIterarionsPerAccess; ++i) { + + int access_offset = + c * ThreadMap::Delta::kContiguous / Detail::kContiguousElementsPerLine * line_size + + vec_pointer_offset + i * line_size; + + int access_idx = (c + s * ThreadMap::Iterations::kContiguous) * + Detail::kIterarionsPerAccess + i; + + char const *access_byte_ptr = reinterpret_cast(access_ptr + access_offset); + + frag_ptr[access_idx] = *reinterpret_cast( + access_byte_ptr + byte_offset_); + } + } + } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + AccessType const *frag_ptr = reinterpret_cast(&frag); + + Index vec_pointer_offset = pointer_offset / Layout::kElementsPerAccess; + + MCTLASS_PRAGMA_UNROLL + for (int s = 0; s < ThreadMap::Iterations::kStrided; ++s) { + + AccessType *access_ptr = pointer_[(s & 1) ^ ((s >> 1) & 1)]; + + access_ptr += 16 * (s / 2) + vec_pointer_offset; + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < ThreadMap::Iterations::kContiguous; ++c) { + MCTLASS_PRAGMA_UNROLL + for(int i = 0; i < Detail::kIterarionsPerAccess; ++i) { + + int access_offset = + c * ThreadMap::Delta::kContiguous / Detail::kContiguousElementsPerLine * line_size + i * line_size; + + int access_idx = (c + s * ThreadMap::Iterations::kContiguous) * + Detail::kIterarionsPerAccess + i; + + char *access_byte_ptr = reinterpret_cast(access_ptr + access_offset); + + *reinterpret_cast(access_byte_ptr + byte_offset_) = + frag_ptr[access_idx]; + } + } + } + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for column-major crosswise TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator::value, Shape_::kRow>, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for column-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::ColumnMajorVoltaTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kRow>; + static int const kAdvanceRank = AdvanceRank; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileIterator< + layout::PitchLinearShape, Element, + layout::VoltaTensorOpMultiplicandCrosswise::value, + Shape::kRow>, + (kAdvanceRank == 0 ? 0 : 1), ThreadMap_>; + + public: + /// Fragment object to be loaded or stored + using Fragment = Array; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.row(), coord.column()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + RegularTileIterator prev(*this); + ++iterator_; + + return prev; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +/// Tile Iterator specialized for row-major crosswise TensorOp formats. +/// +/// +/// Satisfies: ForwardTileIteratorConcept | +/// ReadableContiguousTileIteratorConcept | +/// WriteableContiguousTileIteratorConcept +/// +template < + typename Shape_, + typename Element_, + int AdvanceRank, + typename ThreadMap_, + int Alignment +> +class RegularTileIterator::value, Shape_::kColumn>, + AdvanceRank, ThreadMap_, Alignment> { + public: + static_assert( + AdvanceRank == 0 || AdvanceRank == 1, + "Specialization for row-major iterator may along advance along the " + "columns(rank=0) or rows(rank=1) dimension."); + + using Shape = Shape_; + using Element = Element_; + using Layout = layout::RowMajorVoltaTensorOpMultiplicandCrosswise< + sizeof_bits::value, Shape::kColumn>; + static int const kAdvanceRank = AdvanceRank; + static int const kAlignment = Alignment; + + using Index = typename Layout::Index; + using LongIndex = typename Layout::LongIndex; + + using TensorRef = TensorRef; + using TensorCoord = typename Layout::TensorCoord; + + using ThreadMap = ThreadMap_; + + /// Underlying iterator type + using UnderlyingIterator = RegularTileIterator< + layout::PitchLinearShape, Element, + layout::VoltaTensorOpMultiplicandCrosswise::value, + Shape::kColumn>, + (kAdvanceRank == 0 ? 1 : 0), ThreadMap_>; + + public: + /// Fragment object to be loaded or stored + using Fragment = Array; + + private: + /// Underlying iterator + UnderlyingIterator iterator_; + + public: + /// Construct a TileIterator with zero threadblock offset + MCTLASS_HOST_DEVICE + RegularTileIterator(TensorRef ref, ///< Pointer to start of tensor + int thread_id ///< ID of each participating thread + ) + : iterator_({ref.data(), ref.stride()}, thread_id) {} + + /// Adds a pointer offset in units of Element + MCTLASS_HOST_DEVICE + void add_pointer_offset(LongIndex pointer_offset) { + iterator_.add_pointer_offset(pointer_offset); + } + + /// Adds a tile offset + MCTLASS_DEVICE + void add_tile_offset(TensorCoord const &coord) { + iterator_.add_tile_offset({coord.column(), coord.row()}); + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator &operator++() { + ++iterator_; + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + RegularTileIterator operator++(int) { + RegularTileIterator prev(*this); + ++iterator_; + + return prev; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + iterator_.load_with_pointer_offset(frag, pointer_offset); + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { load_with_pointer_offset(frag, 0); } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store_with_pointer_offset(Fragment const &frag, Index pointer_offset) { + iterator_.store_with_pointer_offset(frag, pointer_offset); + } + + /// Store a fragment to memory + MCTLASS_DEVICE + void store(Fragment const &frag) { store_with_pointer_offset(frag, 0); } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass diff --git a/csrc/mctlass/include/mctlass/transform/threadblock/vector_iterator.h b/csrc/mctlass/include/mctlass/transform/threadblock/vector_iterator.h new file mode 100644 index 0000000..66cb3ca --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/threadblock/vector_iterator.h @@ -0,0 +1,149 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Template wraps the vector access iterator concept to load whole vector from tensors in + memory. This is typically used for per-channel scale and bias in convolution kernels. +*/ + +#pragma once + +#include "mctlass/transform/threadblock/predicated_vector_access_iterator.h" + +///////////////////////////////////////////////////////////////////////////////////////////////// + +namespace mctlass { +namespace transform { +namespace threadblock { + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +class VectorIterator { +public: + using VectorAccessIterator = VectorAccessIterator_; + + using Shape = typename VectorAccessIterator::Shape; + using Element = typename VectorAccessIterator::Element; + using Layout = typename VectorAccessIterator::Layout; + using TensorCoord = typename Layout::TensorCoord; + using AccessType = typename VectorAccessIterator::AccessType; + using TensorRef = typename VectorAccessIterator::TensorRef; + using Index = typename VectorAccessIterator::Index; + using LongIndex = typename VectorAccessIterator::LongIndex; + + static int const kElementsPerAccess = VectorAccessIterator::kElementsPerAccess; + static int const kRowsPerIteration = VectorAccessIterator::kRowsPerIteration; + static int const kThreads = VectorAccessIterator::kThreads; + static int const kIterations = VectorAccessIterator::kIterations; + + /// Fragment object to be loaded or stored + using Fragment = mctlass::Array< + Element, kElementsPerAccess * kIterations>; + +private: + + /// Internal state + VectorAccessIterator vector_access_iterator_; + +public: + + /// Constructor + MCTLASS_HOST_DEVICE + VectorIterator( + Element const *ptr, + TensorCoord extent, + int thread_idx, + int warp_idx, + MatrixCoord const &threadblock_offset = MatrixCoord() + ): + vector_access_iterator_(ptr, extent, thread_idx, warp_idx, threadblock_offset) { } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + VectorIterator &operator++() { + vector_access_iterator_.advance(); + return *this; + } + + /// Advances to the next tile in memory. + MCTLASS_HOST_DEVICE + VectorIterator operator++(int) { + VectorIterator self(*this); + operator++(); + return self; + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load_with_pointer_offset(Fragment &frag, Index pointer_offset) { + + frag.clear(); + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int c = 0; c < kIterations; ++c) { + + mctlass::arch::global_load< + AccessType, + sizeof(AccessType) + >( + frag_ptr[c], + vector_access_iterator_.get() + pointer_offset, + vector_access_iterator_.valid() + ); + + ++vector_access_iterator_; + } +// } + } + + /// Loads a fragment from memory + MCTLASS_DEVICE + void load(Fragment &frag) { + vector_access_iterator_.set_iteration_index(0); + load_with_pointer_offset(frag, 0); + } + + MCTLASS_DEVICE + void advance() { + vector_access_iterator_.advance(); + } + +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace threadblock +} // namespace transform +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// + diff --git a/csrc/mctlass/include/mctlass/transform/warp/vector_fragment_iterator.h b/csrc/mctlass/include/mctlass/transform/warp/vector_fragment_iterator.h new file mode 100644 index 0000000..fba456d --- /dev/null +++ b/csrc/mctlass/include/mctlass/transform/warp/vector_fragment_iterator.h @@ -0,0 +1,283 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + + +/*! \file + \brief This defines a "fragment" iterator for visiting the fragments of a warp vector + that participate in one warp-level mma operation. + + Typically, this is used to access the scale/bias fragement of a warp-level mma operation. + The scale/bias vector is then partitioned into smaller fragments that can be fed into + next warp-level mma operation. + + This iterator is necessary to accomplish warp-level mma fusion where the scale/bias vector is + applied to the multiplicand for the next mma. + +*/ + +#pragma once + +#include "mctlass/mctlass.h" + +#include "mctlass/array.h" +#include "mctlass/matrix_shape.h" +#include "mctlass/layout/matrix.h" +#include "mctlass/layout/tensor.h" +#include "mctlass/numeric_conversion.h" + +namespace mctlass { +namespace transform { +namespace warp { + + +//////////////////////////////////////////////////////////////////////////////// + +template < + /// Size of the input fragment tile shape (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Layout of operand in memory + typename Layout_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + //// Number of elements per access when loading fragment + int ElementsPerAccess> +class VectorFragmentIterator; + + +// Partial specialization for PitchLinear layout tile + +template < + /// Size of the input fragment vector shape (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + //// Number of elements per access when loading fragment + int ElementsPerAccess> +class VectorFragmentIterator { + public: + + /// Size of the input threadblock tile shape (concept: MatrixShape) + using Shape = Shape_; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::PitchLinear; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Number of participating threads + static int const kThreads = 32; + + static int const kElementsPerAccess = ElementsPerAccess; + static int const kRowsPerIteration = 8; + static int const kColumnsPerAccess = 8; + static int const kElementsPerIteration = kRowsPerIteration * InstructionShape::kK / kThreads; + static int const kAccessPerIteration = kElementsPerIteration / kElementsPerAccess; + + /// Number of iterations + using Iterations = MatrixShape; + +public: + + // + // Derived quantities + // + // All fragments have kElementsPerAccess scale followed by bias + + /// Fragment object holding a thread's part of a tile + /// This is the fragment size produced by one iteration of the iterator. + using Fragment = Array; + + /// Input threadblock fragment tile + using ThreadblockFragment = Array; + +private: + + /// Internal access type + using AccessType = Array; + +private: + // + // Data members + // + + /// Input threadblock fragment tile + AccessType const *iterator_; + + /// Internal index + int index_; + +public: + /// Constructs an iterator + MCTLASS_HOST_DEVICE + VectorFragmentIterator(ThreadblockFragment const &threadblock_frag) + : iterator_(reinterpret_cast(&threadblock_frag)), + index_(0) {} + + /// Add offset + MCTLASS_HOST_DEVICE + void add_offset(int index_offset) { + index_ += index_offset; + + if(index_ >= Iterations::kColumn) + index_ = 0; + } + + /// Increments + MCTLASS_HOST_DEVICE + VectorFragmentIterator &operator++() { + add_offset(1); + return *this; + } + + MCTLASS_HOST_DEVICE + void set_index(int idx) { + index_ = idx; + } + + /// Loads a fragment from the referenced part of the accumulator tile + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + + AccessType *frag_ptr = reinterpret_cast(&frag); + + MCTLASS_PRAGMA_UNROLL + for (int r = 0; r < Iterations::kRow; r++) { + MCTLASS_PRAGMA_UNROLL + for (int i = 0; i < kAccessPerIteration; i++) { + + frag_ptr[i * Iterations::kRow + r].clear(); + frag_ptr[i * Iterations::kRow + r] = iterator_[index_ * kAccessPerIteration + i]; + } + } + } + +}; + +// Partial specialization for Row-Major layout tile + +template < + /// Size of the input fragment tile shape (concept: MatrixShape) + typename Shape_, + /// Element type + typename Element_, + /// Shape of one matrix product operation (concept: MatrixShape) + typename InstructionShape_, + //// Number of elements per access when loading fragment + int ElementsPerAccess> +class VectorFragmentIterator { + public: + + /// Size of the input threadblock tile shape (concept: MatrixShape) + using Shape = Shape_; + + /// Element type + using Element = Element_; + + /// Layout of source tile + using Layout = mctlass::layout::RowMajor; + + /// Shape of one matrix product operation (concept: MatrixShape) + using InstructionShape = InstructionShape_; + + /// Underlying iterator + using Base = VectorFragmentIterator< + layout::PitchLinearShape, Element, + layout::PitchLinear, InstructionShape, ElementsPerAccess>; + + + public: + + // + // Derived quantities + // + /// Fragment object holding a thread's part of a tile + /// This is the fragment size produced by one iteration of the iterator. + using Fragment = typename Base::Fragment; + + /// Input threadblock fragment tile + using ThreadblockFragment = typename Base::ThreadblockFragment; + + private: + /// Underlying iterator + Base iterator_; + +public: + /// Constructs an iterator + MCTLASS_HOST_DEVICE + VectorFragmentIterator(ThreadblockFragment const &threadblock_frag) + : iterator_(threadblock_frag) {} + + /// Add offset + MCTLASS_HOST_DEVICE + void add_offset(int index_offset) { + iterator_.add_offset(index_offset); + } + + /// Increments + MCTLASS_HOST_DEVICE + VectorFragmentIterator &operator++() { + add_offset(1); + return *this; + } + + MCTLASS_HOST_DEVICE + void set_index(int idx) { + iterator_.set_index(idx); + } + + /// Loads a fragment from the referenced part of the accumulator tile + MCTLASS_HOST_DEVICE + void load(Fragment &frag) const { + iterator_.load(frag); + } + +}; + + +//////////////////////////////////////////////////////////////////////////////// + +} // namespace warp +} // namespace conv +} // namespace mctlass + +//////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/uint128.h b/csrc/mctlass/include/mctlass/uint128.h new file mode 100644 index 0000000..71a8e34 --- /dev/null +++ b/csrc/mctlass/include/mctlass/uint128.h @@ -0,0 +1,266 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! + \file + \brief Defines an unsigned 128b integer with several operators to support 64-bit integer division. +*/ + +#pragma once + +#if defined(__MACACC_RTC__) +#include +#else +#include +#include +#include +#include +#include +#endif + +#include "mctlass/mctlass.h" +#include "mctlass/numeric_types.h" + +/// Optionally enable GCC's built-in type +#if (defined(__x86_64) || defined (__aarch64__)) && !defined(__MACA_ARCH__) && defined(__GNUC__) +#define MCTLASS_UINT128_NATIVE +#elif defined(_MSC_VER) && defined(_M_AMD64) && !defined(__MACA_ARCH__) +#define MCTLASS_INT128_ARITHMETIC +#include +#if _MSC_VER >= 1920 +#define MCTLASS_INT128_ARITHMETIC_DIV +#include +#endif +#endif + +namespace mctlass { + +///! Unsigned 128b integer type +struct uint128_t { + + /// Size of one part of the uint's storage in bits + static constexpr int kPartSize = sizeof_bits::value; + + struct hilo { + uint64_t lo; + uint64_t hi; + + hilo() = default; + + MCTLASS_HOST_DEVICE hilo(uint64_t lo_, uint64_t hi_):lo(lo_), hi(hi_) {} + }; + + // Use a union to store either low and high parts or, if present, a built-in 128b integer type. + union { + struct hilo hilo_; + + #if defined(MCTLASS_UINT128_NATIVE) + unsigned __int128 native; + #endif // defined(MCTLASS_UINT128_NATIVE) + }; + + // + // Methods + // + + /// Default ctor + uint128_t() = default; + + /// Constructor from uint64 + MCTLASS_HOST_DEVICE + uint128_t(uint64_t lo_): hilo_(lo_, 0) { } + + /// Constructor from two 64b unsigned integers + MCTLASS_HOST_DEVICE + uint128_t(uint64_t lo_, uint64_t hi_): hilo_(lo_, hi_) { + + } + + /// Optional constructor from native value + #if defined(MCTLASS_UINT128_NATIVE) + uint128_t(unsigned __int128 value): native(value) { } + #endif + + /// Lossily cast to uint64 + MCTLASS_HOST_DEVICE + explicit operator uint64_t() const { + return hilo_.lo; + } + + MCTLASS_HOST_DEVICE + static void exception() { +#if defined(__MACA_ARCH__) +#if 0 + asm volatile (" brkpt;\n"); +#endif + printf("uint128.h L127 this function not work correctly now.\n"); +#else + // throw std::runtime_error("Not yet implemented."); + abort(); +#endif + } + + /// Add + MCTLASS_HOST_DEVICE + uint128_t operator+(uint128_t const &rhs) const { + uint128_t y; +#if defined(MCTLASS_UINT128_NATIVE) + y.native = native + rhs.native; +#else + y.hilo_.lo = hilo_.lo + rhs.hilo_.lo; + y.hilo_.hi = hilo_.hi + rhs.hilo_.hi + (!y.hilo_.lo && (rhs.hilo_.lo)); +#endif + return y; + } + + /// Subtract + MCTLASS_HOST_DEVICE + uint128_t operator-(uint128_t const &rhs) const { + uint128_t y; +#if defined(MCTLASS_UINT128_NATIVE) + y.native = native - rhs.native; +#else + y.hilo_.lo = hilo_.lo - rhs.hilo_.lo; + y.hilo_.hi = hilo_.hi - rhs.hilo_.hi - (rhs.hilo_.lo && y.hilo_.lo > hilo_.lo); +#endif + return y; + } + + /// Multiply by unsigned 64b integer yielding 128b integer + MCTLASS_HOST_DEVICE + uint128_t operator*(uint64_t const &rhs) const { + uint128_t y{}; +#if defined(MCTLASS_UINT128_NATIVE) + y.native = native * rhs; +#elif defined(MCTLASS_INT128_ARITHMETIC) + // Multiply by the low part + y.hilo_.lo = _umul128(hilo_.lo, rhs, &y.hilo_.hi); + + // Add the high part and ignore the overflow + uint64_t overflow; + y.hilo_.hi += _umul128(hilo_.hi, rhs, &overflow); +#else + MCTLASS_UNUSED(rhs); + exception(); +#endif + return y; + } + + /// Divide 128b operation by 64b operation yielding a 64b quotient + MCTLASS_HOST_DEVICE + uint64_t operator/(uint64_t const &divisor) const { + uint64_t quotient = 0; +#if defined(MCTLASS_UINT128_NATIVE) + quotient = uint64_t(native / divisor); +#elif defined(MCTLASS_INT128_ARITHMETIC_DIV) + // implemented using MSVC's arithmetic intrinsics + uint64_t remainder = 0; + quotient = _udiv128(hilo_.hi, hilo_.lo, divisor, &remainder); +#else + MCTLASS_UNUSED(divisor); + exception(); +#endif + return quotient; + } + + /// Divide 128b operation by 64b operation yielding a 64b quotient + MCTLASS_HOST_DEVICE + uint64_t operator%(uint64_t const &divisor) const { + uint64_t remainder = 0; +#if defined(MCTLASS_UINT128_NATIVE) + remainder = uint64_t(native % divisor); +#elif defined(MCTLASS_INT128_ARITHMETIC_DIV) + // implemented using MSVC's arithmetic intrinsics + (void)_udiv128(hilo_.hi, hilo_.lo, divisor, &remainder); +#else + MCTLASS_UNUSED(divisor); + exception(); +#endif + return remainder; + } + + /// Computes the quotient and remainder in a single method. + MCTLASS_HOST_DEVICE + uint64_t divmod(uint64_t &remainder, uint64_t divisor) const { + uint64_t quotient = 0; +#if defined(MCTLASS_UINT128_NATIVE) + quotient = uint64_t(native / divisor); + remainder = uint64_t(native % divisor); +#elif defined(MCTLASS_INT128_ARITHMETIC_DIV) + // implemented using MSVC's arithmetic intrinsics + quotient = _udiv128(hilo_.hi, hilo_.lo, divisor, &remainder); +#else + MCTLASS_UNUSED(remainder); + MCTLASS_UNUSED(divisor); + exception(); +#endif + return quotient; + } + + /// Left-shifts a 128b unsigned integer + MCTLASS_HOST_DEVICE + uint128_t operator<<(int sh) const { + if (sh == 0) { + return *this; + } + else if (sh >= kPartSize) { + return uint128_t(0, hilo_.lo << (sh - kPartSize)); + } + else { + return uint128_t( + (hilo_.lo << sh), + (hilo_.hi << sh) | uint64_t(hilo_.lo >> (kPartSize - sh)) + ); + } + } + + /// Right-shifts a 128b unsigned integer + MCTLASS_HOST_DEVICE + uint128_t operator>>(int sh) const { + if (sh == 0) { + return *this; + } + else if (sh >= kPartSize) { + return uint128_t((hilo_.hi >> (sh - kPartSize)), 0); + } + else { + return uint128_t( + (hilo_.lo >> sh) | (hilo_.hi << (kPartSize - sh)), + (hilo_.hi >> sh) + ); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/mctlass/include/mctlass/wmma_array.h b/csrc/mctlass/include/mctlass/wmma_array.h new file mode 100644 index 0000000..8fdfbe1 --- /dev/null +++ b/csrc/mctlass/include/mctlass/wmma_array.h @@ -0,0 +1,133 @@ +/*************************************************************************************************** + * Copyright (c) 2017 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ +/*! \file + \brief Statically sized array of elements that accommodates all MCTLASS-supported numeric types + and is safe to use in a union. +*/ + +#pragma once +// maca wmma head file +#include "mctlass/arch/wmma.h" + +#if defined(MCTLASS_ARCH_WMMA_ENABLED) + +#include "mctlass/mctlass.h" +#include "mctlass/array.h" +#include "mctlass/functional.h" + +namespace mctlass { + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +/// Wmma array type (WmmaFragmentArray holds elements of of type mxmaca::wmma::fragment) +template < + /// Element type + typename T, + /// Number of elements in the array + int N, + /// Whether the element type of T is half_t or __half + bool IsHalfType = (platform::is_same::value || + platform::is_same::value) +> +class WmmaFragmentArray: public Array { +public: + + /// Efficient clear method (override Array::clear()) + MCTLASS_HOST_DEVICE + void clear() + { + for(int i = 0; i < Array::kElements; i++) + { + mxmaca::wmma::fill_fragment((*this)[i], (typename T::element_type)0); + } + } + + MCTLASS_HOST_DEVICE + WmmaFragmentArray& operator+=(const WmmaFragmentArray& rhs) + { + using element_type = typename T::element_type; + plus add; + + for (int i = 0; i < Array::kElements; i++) + { + (*this)[i] = add((*this)[i], rhs[i]); + } + + return *this; + } +}; + +/// Partial specialization for the case in which T::element_type is +/// half_t or __half. This is needed because the cast (typename T::element_type)0 +/// in the primary template flags as an error when __CUDA_NO_HALF_CONVERSIONS__ +/// is set. +template < + /// Element type + typename T, + /// Number of elements in the array + int N +> +class WmmaFragmentArray: public Array { +public: + + /// Efficient clear method (override Array::clear()) + MCTLASS_HOST_DEVICE + void clear() + { + for(int i = 0; i < Array::kElements; i++) + { + mxmaca::wmma::fill_fragment((*this)[i], __float2half(0.f)); + } + } + + MCTLASS_HOST_DEVICE + WmmaFragmentArray& operator+=(const WmmaFragmentArray& rhs) + { + using element_type = typename T::element_type; + plus add; + + for (int i = 0; i < Array::kElements; i++) + { + (*this)[i] = add((*this)[i], rhs[i]); + } + + return *this; + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace mctlass + +/////////////////////////////////////////////////////////////////////////////////////////////////// + +#endif // if defined(MCTLASS_ARCH_WMMA_ENABLED) + diff --git a/csrc/utils/arch.h b/csrc/utils/arch.h new file mode 100644 index 0000000..6ecfef6 --- /dev/null +++ b/csrc/utils/arch.h @@ -0,0 +1,6 @@ +#pragma once + +enum class Arch { + xcore1000 = 1000, + xcore1500 = 1500, +}; \ No newline at end of file diff --git a/csrc/utils/block_info.h b/csrc/utils/block_info.h index e55d8c1..658ac47 100644 --- a/csrc/utils/block_info.h +++ b/csrc/utils/block_info.h @@ -1,5 +1,3 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - /****************************************************************************** * Copyright (c) 2023, Tri Dao. ******************************************************************************/ @@ -23,6 +21,7 @@ struct BlockInfo { // Otherwise it's cu_seqlens_k[bidb], i.e., we use cu_seqlens_k to store the sequence lengths of K. , seqlen_k_cache((!Varlen || params.cu_seqlens_k == nullptr ? params.seqlen_k : (params.is_seqlens_k_cumulative ? params.cu_seqlens_k[bidb + 1] - sum_s_k : params.cu_seqlens_k[bidb])) - leftpad_k) , actual_seqlen_k(params.seqused_k ? params.seqused_k[bidb] - leftpad_k : seqlen_k_cache + (params.knew_ptr == nullptr ? 0 : params.seqlen_knew)) + , tot_seqlen_k((params.cp_tot_seqused_k == nullptr || params.cp_world_size == 1) ? actual_seqlen_k : params.cp_tot_seqused_k[bidb]) { } @@ -43,6 +42,7 @@ struct BlockInfo { // We have to have seqlen_k_cache declared before actual_seqlen_k, otherwise actual_seqlen_k is set to 0. const int seqlen_k_cache; const int actual_seqlen_k; + const int tot_seqlen_k; }; //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/csrc/utils/hdim_switch.h b/csrc/utils/hdim_switch.h deleted file mode 100644 index 3d9192b..0000000 --- a/csrc/utils/hdim_switch.h +++ /dev/null @@ -1,94 +0,0 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - -#pragma once - -#include "flash_headdim.h" - -constexpr static int Headdim32 = FlashHeaddim<32>::Headdim; -constexpr static int Headdim64 = FlashHeaddim<64>::Headdim; -constexpr static int Headdim96 = FlashHeaddim<96>::Headdim; -constexpr static int Headdim128 = FlashHeaddim<128>::Headdim; -constexpr static int Headdim160 = FlashHeaddim<160>::Headdim; -constexpr static int Headdim192 = FlashHeaddim<192>::Headdim; -constexpr static int Headdim224 = FlashHeaddim<224>::Headdim; -constexpr static int Headdim256 = FlashHeaddim<256>::Headdim; -constexpr static int Headdim512 = FlashHeaddim<512>::Headdim; - -#ifndef MCFLASHINFER - -#define FWD_HEADDIM_SWITCH(HEADDIM, ...) \ - [&] { \ - if (HEADDIM <= 32) { \ - constexpr static int kHeadDim = Headdim32; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 64) { \ - constexpr static int kHeadDim = Headdim64; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 96) { \ - constexpr static int kHeadDim = Headdim96; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 128) { \ - constexpr static int kHeadDim = Headdim128; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 160) { \ - constexpr static int kHeadDim = Headdim160; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 192) { \ - constexpr static int kHeadDim = Headdim192; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 224) { \ - constexpr static int kHeadDim = Headdim224; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 256) { \ - constexpr static int kHeadDim = Headdim256; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 512) { \ - constexpr static int kHeadDim = Headdim512; \ - return __VA_ARGS__(); \ - } \ - }() - -#define BWD_HEADDIM_SWITCH(HEADDIM, ...) \ - [&] { \ - if (HEADDIM <= 32) { \ - constexpr static int kHeadDim = Headdim32; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 64) { \ - constexpr static int kHeadDim = Headdim64; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 96) { \ - constexpr static int kHeadDim = Headdim96; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 128) { \ - constexpr static int kHeadDim = Headdim128; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 160) { \ - constexpr static int kHeadDim = Headdim160; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 192) { \ - constexpr static int kHeadDim = Headdim192; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 224) { \ - constexpr static int kHeadDim = Headdim224; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 256) { \ - constexpr static int kHeadDim = Headdim256; \ - return __VA_ARGS__(); \ - } \ - }() - -#else -#define FWD_HEADDIM_SWITCH(HEADDIM, ...) \ - [&] { \ - if (HEADDIM <= 64) { \ - constexpr static int kHeadDim = Headdim64; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 128) { \ - constexpr static int kHeadDim = Headdim128; \ - return __VA_ARGS__(); \ - } else if (HEADDIM <= 256) { \ - constexpr static int kHeadDim = Headdim256; \ - return __VA_ARGS__(); \ - } \ - }() -#endif diff --git a/csrc/utils/host_utils.h b/csrc/utils/host_utils.h new file mode 100644 index 0000000..aaaa2a3 --- /dev/null +++ b/csrc/utils/host_utils.h @@ -0,0 +1,22 @@ +#pragma once + + +namespace flash { + +inline __host__ mcDeviceProp_t mcGetCurrentDeviceProperties() { + int deviceId{}; + mcGetDevice(&deviceId); + mcDeviceProp_t dprops; + mcGetDeviceProperties(&dprops, deviceId); + return dprops; +} + +inline __host__ int mcGetCurrentDeviceArch() { + int deviceId{}; + mcGetDevice(&deviceId); + mcDeviceProp_t dprops; + mcGetDeviceProperties(&dprops, deviceId); + return dprops.major * 100 + dprops.minor; +} + +} \ No newline at end of file diff --git a/csrc/utils/logger.cpp b/csrc/utils/logger.cpp new file mode 100644 index 0000000..fe7e5d6 --- /dev/null +++ b/csrc/utils/logger.cpp @@ -0,0 +1,170 @@ +/** + * @file logger.cpp + * @brief + * @version 0.1 + * @date 2021-09-02 + * + * @copyright Copyright (c) 2021 + * + */ +#include "logger.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace mcFlashAttn { +namespace utils { + +const char *Logger::LogLevelMap[] = + {"\033[0;31mERR\033[0m", "WARN", "INFO", "DBG"}; + +std::string basename(const std::string &in) { + auto loc = in.find_last_of("/"); + if (loc == std::string::npos) { + return in; + } + return in.substr(loc + 1); +} + +Logger &Logger::GetLogger() { + static Logger s_log; + return s_log; +} + +Logger::Logger() { + // log switch + char *p = getenv(MHA_LOG_ENABLE); + if(p != nullptr) { + external_en_ = true; + env_logext_enable_set_ = true; + } + + // log level + char *env_log_level = getenv(MHA_LOG_LEVEL); + if (env_log_level != nullptr) { + int level = strtol(env_log_level, nullptr, 10); + SetLogLevel(level); + env_log_level_set_ = true; + } + + // output direction + char *env_output = getenv(MHA_LOG_OUTPUT); + if (env_output != nullptr) { + if (strcmp(env_output, "stdout") == 0) { + log_output_ = LOG_STDOUT; + } else if (strcmp(env_output, "file") == 0) { + log_output_ = LOG_FILE; + } else if (strcmp(env_output, "syslog") == 0) { + log_output_ = LOG_SYSLOG; + } + env_log_output_set_ = true; + } + + if (external_en_) { + switch (log_output_) { + case LOG_STDOUT: + of_ = &std::cout; + break; + default: + case LOG_FILE: + logfile_.open(filename_.c_str()); + of_ = &logfile_; + break; + } + } +} + +void Logger::SetLogLevel(int level) { + log_level_ = std::min(level, static_cast(LOG_INFO)); +} + +uint32_t Logger::TimeStr(char *buf, uint32_t bufsize) { + time_t t; + struct tm *timeinfo; + std::time(&t); + timeinfo = std::localtime(&t); + return std::strftime(buf, bufsize, ":%D %s: ", timeinfo); +} + +uint32_t Logger::FmtLogHeader(char *buf, uint32_t bufsize, int level) { + uint32_t off = 0; + off += snprintf(buf + off, bufsize - off, "[MHA][%s] ", LogLevelMap[level]); + off += TimeStr(buf + off, bufsize - off); + return off; +} + +void Logger::IntervalCheck() { + ++counter_; + if (counter_ >= kCheckInterval) { + counter_ = 0; + } +} + +/* + @brief: log simple +*/ + +void Logger::LogSimple(const char *fmt, ...) { + std::lock_guard lock(mutex_); + if (log_output_ != LOG_SYSLOG) { + char message[4096]; + uint32_t off = 0; + + va_list args; + va_start(args, fmt); + off += std::vsnprintf(message + off, sizeof(message) - off, fmt, args); + va_end(args); + + *of_ << message; + } + + IntervalCheck(); +} + +/* + @brief: log external +*/ + +void Logger::LogExternal(int level, const char *fmt, ...) { + std::lock_guard lock(mutex_); + if (log_output_ != LOG_SYSLOG) { + char message[4096]; + uint32_t off = 0; + + off += snprintf(message, sizeof(message), "[MHA][%s] ", LogLevelMap[level]); + + va_list args; + va_start(args, fmt); + off += std::vsnprintf(message + off, sizeof(message) - off, fmt, args); + va_end(args); + + *of_ << message; + } +} + +void Logger::Log(int level, const char *fmt, ...) { + std::lock_guard lock(mutex_); + if (log_output_ != LOG_SYSLOG) { + char message[4096]; + uint32_t off = 0; + + off += FmtLogHeader(message, sizeof(message), level); + + va_list args; + va_start(args, fmt); + off += std::vsnprintf(message + off, sizeof(message) - off, fmt, args); + va_end(args); + *of_ << message; + } + + IntervalCheck(); +} + +} // namespace utils +} // namespace mcFlashAttn diff --git a/csrc/utils/logger.h b/csrc/utils/logger.h new file mode 100644 index 0000000..3e37f1c --- /dev/null +++ b/csrc/utils/logger.h @@ -0,0 +1,263 @@ +/** + * @file logger.h + * @brief + * @version 0.1 + * @date 2021-09-02 + * + * @copyright Copyright (c) 2021 + * + */ +#ifndef UTILS_LOGGER_H_ +#define UTILS_LOGGER_H_ + +#include +#include +#include +#include +#include +#include +#include +#include + +/** + * @brief log init + * usually do not use it. it will cover settings in environment + * Only for tests you want run locally + */ +#define LOG_INIT(enable, level, output, filename) \ + mcFlashAttn::utils::Logger::GetLogger().Init(level, output, filename) + +/** + * @brief log formating + * LOG_F(mcFlashAttn::utils::LOG_INFO, "%s\n", "enjoy"); + * use quick api like LOG_INFO("%s\n", dflakjf); + */ +#define LOG_F(level, fmt, ...) \ + do { \ + if (mcFlashAttn::utils::Logger::GetLogger().ExternalEn() && \ + mcFlashAttn::utils::Logger::Filter(level)) { \ + mcFlashAttn::utils::Logger::GetLogger().Log(level, \ + "%s:%d |\t" fmt, \ + mcFlashAttn::utils::basename(__FILE__).c_str(), \ + __LINE__, \ + ##__VA_ARGS__); \ + } \ + } while (false) + +#define LOG_ERR(fmt, ...) LOG_F(mcFlashAttn::utils::LOG_ERROR, fmt, ##__VA_ARGS__) +#define LOG_WARN(fmt, ...) LOG_F(mcFlashAttn::utils::LOG_WARNING, fmt, ##__VA_ARGS__) +#define LOG_INFO(fmt, ...) LOG_F(mcFlashAttn::utils::LOG_INFO, fmt, ##__VA_ARGS__) +#define LOG_DEBUG(fmt, ...) LOG_F(mcFlashAttn::utils::LOG_DEBUG, fmt, ##__VA_ARGS__) + + +/** + * @brief log format simple + * for logs you don't want header and timestamp + * For example log in a for loop + */ +#define LOG_FS(level, fmt, ...) \ + do { \ + if (mcFlashAttn::utils::Logger::GetLogger().ExternalEn() && \ + mcFlashAttn::utils::Logger::Filter(mcFlashAttn::utils::LOG_##level)) { \ + mcFlashAttn::utils::Logger::GetLogger().LogSimple(fmt, ##__VA_ARGS__); \ + } \ + } while (false) + +#define LOG_FS_ERR(fmt, ...) LOG_FS(ERROR, fmt, ##__VA_ARGS__) +#define LOG_FS_WARN(fmt, ...) LOG_FS(WARNING, fmt, ##__VA_ARGS__) +#define LOG_FS_INFO(fmt, ...) LOG_FS(INFO, fmt, ##__VA_ARGS__) +#define LOG_FS_DEBUG(fmt, ...) LOG_FS(DEBUG, fmt, ##__VA_ARGS__) + +/** + * @brief log external + * + * log external's format is more simple than internal's. Only three levels err, warn and info + * are available! + */ +#define LOG_EXT(level, fmt, ...) \ + do { \ + if (mcFlashAttn::utils::Logger::GetLogger().ExternalEn() && \ + mcFlashAttn::utils::Logger::Filter(level)) { \ + mcFlashAttn::utils::Logger::GetLogger().LogExternal(level, fmt, ##__VA_ARGS__); \ + } \ + } while (false) + +#define LOG_EXT_ERR(fmt, ...) LOG_EXT(mcFlashAttn::utils::LOG_ERROR, fmt, ##__VA_ARGS__) +#define LOG_EXT_WARN(fmt, ...) LOG_EXT(mcFlashAttn::utils::LOG_WARNING, fmt, ##__VA_ARGS__) +#define LOG_EXT_INFO(fmt, ...) LOG_EXT(mcFlashAttn::utils::LOG_INFO, fmt, ##__VA_ARGS__) + +/** + * @brief log stream + * Attention: + * LOG_S is not recommand to use, when internal log is disabled, StreamLogger function + * operator << still run, though nothing will print because of the Filter function + * in that situation. + * Examples: + * LOG_S(INFO) << a << b << "\n"; + * LOG_S(WARNING) << a << b << "\n"; + * + * ... + */ +#define LOG_S(level) mcFlashAttn::utils::StreamLogger(mcFlashAttn::utils::LOG_##level, __FILE__, __LINE__) + +/** + * @brief print shape log + */ +#define LOG_FS_SHAPE(fmt, ...) \ + do { \ + mcFlashAttn::utils::Logger::GetLogger().LogSimple(fmt, ##__VA_ARGS__); \ + } while (false) + +#define LOG_SHAPE(fmt, ...) LOG_FS_SHAPE(fmt, ##__VA_ARGS__) + +namespace mcFlashAttn { +namespace utils { + +enum LogLevel : int32_t { + LOG_ERROR = 0, + LOG_WARNING, + LOG_INFO, + LOG_DEBUG, +}; + +enum LogOutput : int32_t { + LOG_STDOUT = 0, + LOG_FILE, + LOG_SYSLOG, +}; + +std::string basename(const std::string &in); + +// control env variables +#define MHA_LOG_ENABLE "MHA_LOG_ENABLE" +#define MHA_LOG_LEVEL "MHA_LOG_LEVEL" +#define MHA_LOG_OUTPUT "MHA_LOG_OUTPUT" + +class Logger { + public: + Logger(); + void Init(int level, int output, const std::string &filename) { + if (!env_logext_enable_set_) { // env variable has higher priority + external_en_ = true; + } + if (!env_log_level_set_) { // env variable has higher priority + SetLogLevel(level); + } + if (!env_log_output_set_) { // env variable has higher priority + log_output_ = output; + } + + if (logfile_.is_open()) { + logfile_.close(); + } + filename_ = filename; + + if (external_en_) { + switch (log_output_) { + case LOG_STDOUT: + of_ = &std::cout; + break; + default: + case LOG_FILE: + logfile_.open(filename_.c_str()); + of_ = &logfile_; + break; + } + } + } + + static Logger &GetLogger(); + static bool Filter(int level) { return level <= GetLogger().log_level_; } + + bool ExternalEn() const { return external_en_; } + void SetLogLevel(int level); + + void Log(int level, const char *fmt, ...); + /** + * @brief For usage without header "info [time] xxx" + * usally if you want log something in a for loop without header + * and crlf, use this api. + * @param fmt + * @param ... + */ + void LogSimple(const char *fmt, ...); + + /** + * @brief For usage without header "info [time] xxx" + */ + void LogExternal(int level, const char *fmt, ...); + + private: + using LoggerClock = std::chrono::high_resolution_clock; + static uint64_t MicroSeconds() { + std::chrono::duration ts = + std::chrono::duration_cast( + LoggerClock::now().time_since_epoch()); + return ts.count(); + } + void IntervalCheck(); + + static uint32_t TimeStr(char *buf, uint32_t bufsize); + static uint32_t FmtLogHeader(char *buf, uint32_t bufsize, int level); + + int log_level_ = LOG_INFO; + uint32_t log_output_ = LOG_STDOUT; + + std::string ident_ = "mcFlashAttn"; + std::ofstream logfile_; + std::ostream *of_ = &std::cout; + std::mutex mutex_; + std::string filename_ = "mcFlashAttn.txt"; + uint32_t counter_ = 0; + bool env_logext_enable_set_{false}; + bool env_log_level_set_{false}; + bool env_log_output_set_{false}; + bool external_en_ = false; + + static constexpr int kCheckInterval = 1000; + static const char *LogLevelMap[]; +}; + +class StreamLogger { + public: + StreamLogger(int level, const char *filename, int linenum) + : level_(level), filename_(filename), linenum_(linenum) {} + ~StreamLogger() { Flush(); } + void Flush() { + if (Filter()) { + Logger::GetLogger().Log(level_, + "%s:%d |\t%s", + basename(filename_).c_str(), + linenum_, + ss_.str().c_str()); + ss_.clear(); + } + } + + template + StreamLogger &operator<<(const T &v) { + if (Filter()) { + ss_ << v; + } + return *this; + } + + StreamLogger &operator<<(std::ostream &(*f)(std::ostream &)) { + if (Filter()) { + f(ss_); + } + return *this; + } + + private: + bool Filter() const { return Logger::Filter(level_) && Logger::GetLogger().ExternalEn(); } + int level_; + std::string filename_; + int linenum_; + std::stringstream ss_; +}; + +} // namespace utils +} // namespace mcFlashAttn + +#endif // UTILS_LOGGER_H_ diff --git a/csrc/utils/philox.cuh b/csrc/utils/philox.cuh deleted file mode 100644 index 5884fee..0000000 --- a/csrc/utils/philox.cuh +++ /dev/null @@ -1,54 +0,0 @@ -// Adapted from Dao-AILab/flash-attention (https://github.com/Dao-AILab/flash-attention/tree/v2.6.3) - -// Pytorch also has an implementation of Philox RNG: https://github.com/pytorch/pytorch/blob/8ca3c881db3e3510fcb7725389f6a0633c9b992c/torch/csrc/jit/tensorexpr/cuda_random.h -#pragma once -// Philox CUDA. - -namespace flash { - -struct ull2 { - unsigned long long x; - unsigned long long y; -}; - -__forceinline__ __device__ uint2 mulhilo32(const unsigned int a, const unsigned int b) { - uint2 *res; - unsigned long long tmp; - // asm ("mul.wide.u32 %0, %1, %2;\n\t" - // : "=l"(tmp) - // : "r"(a), "r"(b)); - tmp = (uint64_t)a * (uint64_t)b; - res = (uint2*)(&tmp); - return *res; -} - -__forceinline__ __device__ uint4 philox_single_round(const uint4 ctr, const uint2 key) { - constexpr unsigned long kPhiloxSA = 0xD2511F53; - constexpr unsigned long kPhiloxSB = 0xCD9E8D57; - uint2 res0 = mulhilo32(kPhiloxSA, ctr.x); - uint2 res1 = mulhilo32(kPhiloxSB, ctr.z); - uint4 ret = {res1.y ^ ctr.y ^ key.x, res1.x, res0.y ^ ctr.w ^ key.y, res0.x}; - return ret; -} - -__forceinline__ __device__ uint4 philox(unsigned long long seed, - unsigned long long subsequence, - unsigned long long offset) { - constexpr unsigned long kPhilox10A = 0x9E3779B9; - constexpr unsigned long kPhilox10B = 0xBB67AE85; - uint2 key = reinterpret_cast(seed); - uint4 counter; - ull2 *tmp = reinterpret_cast(&counter); - tmp->x = offset; - tmp->y = subsequence; - #pragma unroll - for (int i = 0; i < 6; i++) { - counter = philox_single_round(counter, key); - key.x += (kPhilox10A); - key.y += (kPhilox10B); - } - uint4 output = philox_single_round(counter, key); - return output; -} - -} // namespace flash diff --git a/csrc/utils/print_parameter.cpp b/csrc/utils/print_parameter.cpp new file mode 100644 index 0000000..b91b981 --- /dev/null +++ b/csrc/utils/print_parameter.cpp @@ -0,0 +1,218 @@ +#include +#include +#include +#include +#include "flash_mla.h" +#include "logger.h" +#include "process_str.h" +#include "print_parameter.h" + +std::stringstream process_params(mcFlashAttn::Flash_fwd_mla_params params, bool Is_causal, const std::string& debug_flag) { + /* + + ==============================Parts that require special handling（cu_seqlens）============================== + + */ + std::stringstream cu_seqlen_q; + std::stringstream cu_seqlen_k; + if (params.cu_seqlens_q != nullptr) { + std::vector host_cuseq_q(params.b + 1); + cudaMemcpy(host_cuseq_q.data(), params.cu_seqlens_q, (sizeof(int) * (params.b + 1)), cudaMemcpyDeviceToHost); + for(int i = 0; i < (params.b + 1); ++i) { + if(cu_seqlen_q.str().size() == 0) + { + cu_seqlen_q << "["; + } else if(cu_seqlen_q.str().size() > 1) { + cu_seqlen_q << "-"; + } + cu_seqlen_q << std::to_string(host_cuseq_q[i]); + } + cu_seqlen_q << "]"; + } else { + cu_seqlen_q << "[nil]"; + } + + if (params.cu_seqlens_k != nullptr) { + // when debug_flag is kvcache, cu_seqlens_k_size == batch_size + // when debug_flag is fwd & bwd, cu_seqlens_k_size == batch_size + 1 + const int seq_k_size = debug_flag == "kvcache" ? params.b : (params.b + 1);//? + std::vector host_cuseq_k(seq_k_size); + cudaMemcpy(host_cuseq_k.data(), params.cu_seqlens_k, (sizeof(int) * (seq_k_size)), cudaMemcpyDeviceToHost); + for(int i = 0; i < seq_k_size; ++i) { + if(cu_seqlen_k.str().size() == 0) + { + cu_seqlen_k << "["; + } else if(cu_seqlen_k.str().size() > 1) { + cu_seqlen_k << "-"; + } + cu_seqlen_k << std::to_string(host_cuseq_k[i]); + } + cu_seqlen_k << "]"; + } else { + cu_seqlen_k << "[nil]"; + } + + /* + + ==============================The part where the Bool_switch is printed============================== + + */ + + bool Split = params.num_splits > 1; + + std::vector bool_info{ + /*================ + The unique part of the fwd and shared parts. + ================*/ + std::to_string(params.is_bf16), + std::to_string(Is_causal), + std::to_string(params.is_seqlens_k_cumulative), + std::to_string(params.unpadded_lse), + std::to_string(Split), + std::to_string(params.is_sparse_attn), + }; + /* + + ==============================The part where the Dim_info is printed============================== + + */ + std::vector dim_info{ + std::to_string(params.b), + std::to_string(params.seqlen_q), + std::to_string(params.seqlen_k), + std::to_string(params.ngroups), + std::to_string(params.topk), + std::to_string(params.h), + std::to_string(params.h_k), + std::to_string(params.d), + std::to_string(params.d_value), + std::to_string(params.d_value_rounded), + std::to_string(params.num_splits), + std::to_string(params.page_block_size), + std::to_string(params.scale_softmax), + std::to_string(params.seqlen_knew), + std::to_string(params.seqlen_q_rounded), + std::to_string(params.seqlen_k_rounded), + std::to_string(params.d_rounded), + std::to_string(params.rotary_dim), + /*================ + cu_seqlen + ================*/ + cu_seqlen_q.str(), + cu_seqlen_k.str(), + }; + + return concat_total_strs("MLA", debug_flag, bool_info, dim_info); +} + +void shape_print(mcFlashAttn::Flash_fwd_mla_params params, bool Is_causal, const std::string& debug_flag) { + auto total_strs = process_params(params, Is_causal, debug_flag); + /* + Bool Switch: + is_bf16, is_causal, is_seqlens_k_cumulative, unpadded_lse, Split, is_sparse_attn + Dim_info: + b, seqlen_q, seqlen_k, ngroups, topk, h, h_k, d, d_value, d_value_rounded, num_splits, page_block_size + scale_softmax, seqlen_knew, seqlen_q_rounded, seqlen_k_rounded, d_rounded, rotary_dim, + cu_seqlens: + cu_seqlens_q, cu_seqlens_k, + */ + LOG_SHAPE("%s\n", total_strs.str().c_str()); +} + +void shape_print(SparsePrefillParams params, bool Is_causal, const std::string& debug_flag) { + /* + Bool Switch: + is_causal + Dim_info: + s_q, s_kv, h_q, h_kv, d_qk, d_v, topk + sm_scale, sm_scale_div_log2 + */ + std::vector bool_info{ + std::to_string(Is_causal), + }; + /* + + ==============================The part where the Dim_info is printed============================== + + */ + std::vector dim_info{ + std::to_string(params.s_q), + std::to_string(params.s_kv), + std::to_string(params.h_q), + std::to_string(params.h_kv), + std::to_string(params.d_qk), + std::to_string(params.d_v), + std::to_string(params.topk), + std::to_string(params.sm_scale), + std::to_string(params.sm_scale_div_log2), + }; + auto total_strs = concat_total_strs("MLA", debug_flag, bool_info, dim_info); + LOG_SHAPE("%s\n", total_strs.str().c_str()); +} + + +void debug_print(mcFlashAttn::Flash_fwd_mla_params params, const std::string& debug_flag) { + printf("==============%s-debug parameters recored start...\n", debug_flag.c_str()); + + printf("----rng_state_seed=%d\n",params.rng_state_seed); + printf("----rng_state_offset=%d\n",params.rng_state_offset); + printf("----p_ptr=%p\n", params.p_ptr); + printf("----rp_dropout=%f\n",params.rp_dropout); + + printf("----rotary_cos_ptr=%p\n", params.rotary_cos_ptr); + printf("----rotary_sin_ptr=%p\n", params.rotary_sin_ptr); + printf("----cache_batch_idx=%p\n", params.cache_batch_idx); + printf("----block_table=%p\n", params.block_table); + printf("----block_table_batch_stride=%ld\n", params.block_table_batch_stride); + printf("----page_block_size=%d\n", params.page_block_size); + printf("----knew_ptr=%p\n",params.knew_ptr); + printf("----vnew_ptr=%p\n",params.vnew_ptr); + printf("----oaccum_ptr=%p\n", params.oaccum_ptr); + printf("----num_splits=%d\n",params.num_splits); + printf("----softmax_lse_ptr=%p\n", params.softmax_lse_ptr); + printf("----softmax_lseaccum_ptr=%p\n", params.softmax_lseaccum_ptr); + + printf("----softmax_lse_ptr=%p\n", params.softmax_lse_ptr); + printf("----seqused_k=%p\n",params.seqused_k); + printf("----q_ptr=%p\n", params.q_ptr); + printf("----k_ptr=%p\n", params.k_ptr); + printf("----v_ptr=%p\n", params.v_ptr); + printf("----o_ptr=%p\n", params.o_ptr); + printf("----scale_softmax=%f\n",params.scale_softmax); + printf("----scale_softmax_log2=%f\n",params.scale_softmax_log2); + printf("----o_batch_stride=%ld\n", params.o_batch_stride); + printf("----o_row_stride=%ld\n",params.o_row_stride); + printf("----o_head_stride=%ld\n",params.o_head_stride); + printf("----q_batch_stride=%ld\n",params.q_batch_stride); + printf("----k_batch_stride=%ld\n",params.k_batch_stride); + printf("----v_batch_stride=%ld\n",params.v_batch_stride); + printf("----q_row_stride=%ld\n",params.q_row_stride); + printf("----k_row_stride=%ld\n",params.k_row_stride); + printf("----v_row_stride=%ld\n",params.v_row_stride); + printf("----q_head_stride=%ld\n",params.q_head_stride); + printf("----k_head_stride=%ld\n",params.k_head_stride); + printf("----v_head_stride=%ld\n",params.v_head_stride); + printf("----unpadded lse=%d\n", params.unpadded_lse); + + printf("----d_value=%d\n", params.d_value); + printf("----d_value_rounded=%d\n", params.d_value_rounded); + printf("----num_sm_parts=%d\n", params.num_sm_parts); + printf("==============%s-debug parameters recored end...\n", debug_flag.c_str()); +} + +void debug_print(SparsePrefillParams params, const std::string& debug_flag) { + printf("==============%s-debug parameters recored start...\n", debug_flag.c_str()); + + printf("----s_q=%d\n",params.s_q); + printf("----s_kv=%d\n",params.s_kv); + printf("----h_q=%d\n",params.h_q); + printf("----h_kv=%d\n",params.h_kv); + printf("----d_qk=%d\n",params.d_qk); + printf("----d_v=%d\n",params.d_v); + printf("----topk=%d\n",params.topk); + + printf("----sm_scale=%f\n",params.sm_scale); + printf("----sm_scale_div_log2=%f\n",params.sm_scale_div_log2); + + printf("==============%s-debug parameters recored end...\n", debug_flag.c_str()); +} diff --git a/csrc/utils/print_parameter.h b/csrc/utils/print_parameter.h new file mode 100644 index 0000000..a1567a0 --- /dev/null +++ b/csrc/utils/print_parameter.h @@ -0,0 +1,15 @@ +#pragma once + +#include +#include +#include +#include "flash_mla.h" + + +void shape_print(mcFlashAttn::Flash_fwd_mla_params params, bool Is_causal, const std::string& debug_flag); + +void shape_print(SparsePrefillParams params, bool Is_causal, const std::string& debug_flag); + +void debug_print(mcFlashAttn::Flash_fwd_mla_params params, const std::string& debug_flag); + +void debug_print(SparsePrefillParams params, const std::string& debug_flag); \ No newline at end of file diff --git a/csrc/utils/process_str.cpp b/csrc/utils/process_str.cpp new file mode 100644 index 0000000..6bac0f6 --- /dev/null +++ b/csrc/utils/process_str.cpp @@ -0,0 +1,31 @@ +#include "process_str.h" + +template +void concat_str(std::stringstream &total_string, std::vector info){ + for(int i=0; i < info.size(); i++) { + if(i != 0) { + total_string << "," << info[i]; + } + else { + total_string << info[i]; + } + } +} + +std::stringstream concat_total_strs(const std::string& prefix_flag, const std::string& debug_flag, std::vector bool_info, std::vector shape_info) { + std::stringstream total_strs; + total_strs << "["; + total_strs << prefix_flag; + total_strs << "] "; + + total_strs << debug_flag; + total_strs << "_Bool_Switch:"; + concat_str(total_strs, bool_info); + total_strs << " "; + + total_strs << debug_flag; + total_strs << "_Shape_Info:"; + concat_str(total_strs, shape_info); + + return total_strs; +} \ No newline at end of file diff --git a/csrc/utils/process_str.h b/csrc/utils/process_str.h new file mode 100644 index 0000000..d433986 --- /dev/null +++ b/csrc/utils/process_str.h @@ -0,0 +1,10 @@ +#pragma once + +#include +#include +#include + +template +void concat_str(std::stringstream &total_string, std::vector info); + +std::stringstream concat_total_strs(const std::string& prefix_flag, const std::string& debug_flag, std::vector bool_info, std::vector shape_info); \ No newline at end of file diff --git a/csrc/utils/static_switch.h b/csrc/utils/static_switch.h index f1350e4..f78c55e 100644 --- a/csrc/utils/static_switch.h +++ b/csrc/utils/static_switch.h @@ -2,9 +2,11 @@ // Inspired by // https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h // and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h - +#include "arch.h" #pragma once +#define CHECK_MSG(x, ...) do { if((x) == false) {throw std::invalid_argument(__VA_ARGS__);} }while(0) + /// CONST_PRECOND && COND #define BOOL_SWITCH_AND_CONST_PRECOND(CONST_PRECOND, COND, CONST_NAME, ...) \ [&] { \ @@ -121,6 +123,23 @@ } \ }() +#define FLASH_ASSERT(cond) \ + do { \ + if (not (cond)) { \ + fprintf(stderr, "Assertion failed (%s:%d): %s\n", __FILE__, __LINE__, #cond); \ + exit(1); \ + } \ + } while(0) + + +#define FLASH_DEVICE_ASSERT(cond) \ + do { \ + if (not (cond)) { \ + printf("Assertion failed (%s:%d): %s\n", __FILE__, __LINE__, #cond); \ + asm("trap 0x1;"); \ + } \ + } while(0) + #define NUMSPLITS_SWITCH(NUMSPLITS, ...) \ [&] { \ if (NUMSPLITS <= 2) { \ @@ -144,19 +163,70 @@ } else if (NUMSPLITS <= 128) { \ constexpr static int kLogMaxSplits = 7; \ return __VA_ARGS__(); \ - } \ + }else if (NUMSPLITS <= 256) { \ + constexpr static int kLogMaxSplits = 8; \ + return __VA_ARGS__(); \ + } \ }() -#define ROWNUM_SWITCH(COND, CONST_NAME, ...) \ - [&] { \ - if (COND) \ - { \ - constexpr static int CONST_NAME = 2; \ - return __VA_ARGS__(); \ - } \ - else \ - { \ - constexpr static int CONST_NAME = 1; \ - return __VA_ARGS__(); \ - } \ +#define COMBINE_BLOCKM_SWITCH(BATCH, HEADQ, SEQLENQ, CONST_NAME, ...) \ + [&] { \ + const int tot_block_num = BATCH * HEADQ * SEQLENQ; \ + const int totl_qo_num = SEQLENQ * HEADQ; \ + if (totl_qo_num == 2) { \ + constexpr static int CONST_NAME = 2; \ + return __VA_ARGS__(); \ + } \ + else if (totl_qo_num < 16 && totl_qo_num % 4 == 0) { \ + constexpr static int CONST_NAME = 4; \ + return __VA_ARGS__(); \ + } else { \ + if (tot_block_num >= 1024 && totl_qo_num % 16 == 0) { \ + constexpr static int CONST_NAME = 16; \ + return __VA_ARGS__(); \ + } else if (totl_qo_num % 4 == 0) { \ + constexpr static int CONST_NAME = 4; \ + return __VA_ARGS__(); \ + } \ + else { \ + constexpr static int CONST_NAME = 1; \ + return __VA_ARGS__(); \ + } \ + } \ }() + +#if defined(XCORE1000) + #define ARCH_SWITCH ARCH_SWITCH_XCORE1000 +#elif defined(XCORE1500) + #define ARCH_SWITCH ARCH_SWITCH_XCORE1500 +#else + #define ARCH_SWITCH ARCH_SWITCH_ALL +#endif + +#define ARCH_SWITCH_ALL(ARCH, ARCH_NAME, ...) \ + [&] { \ + if (ARCH == 1000) { \ + constexpr static Arch ARCH_NAME = Arch::xcore1000; \ + return __VA_ARGS__(); \ + } else if (ARCH == 1500) { \ + constexpr static Arch ARCH_NAME = Arch::xcore1500; \ + return __VA_ARGS__(); \ + } else { \ + CHECK_MSG(false, "This arch xcore" + std::to_string(ARCH) + \ + " is not supported, please check your arch!"); \ + } \ + }() + +#define ARCH_SWITCH_XCORE1000(ARCH, ARCH_NAME, ...) \ + [&] { \ + constexpr static Arch ARCH_NAME = Arch::xcore1000; \ + return __VA_ARGS__(); \ + }() + +#define ARCH_SWITCH_XCORE1500(ARCH, ARCH_NAME, ...) \ + [&] { \ + constexpr static Arch ARCH_NAME = Arch::xcore1500; \ + return __VA_ARGS__(); \ + }() + + diff --git a/csrc/utils/utils.h b/csrc/utils/utils.h index 1685188..fdce045 100644 --- a/csrc/utils/utils.h +++ b/csrc/utils/utils.h @@ -26,125 +26,17 @@ #include #include #include - -//////////////////////////////////////////////////////////////////////////////////////////////////// +#define SET_BIT(var, pos) ((var) |= (1U << (pos))) +#define CLEAR_BIT(var, pos) ((var) &= ~(1U << (pos))) +#define TOGGLE_BIT(var, pos) ((var) ^= (1U << (pos))) +#define CHECK_BIT(var, pos) (((var) >> (pos)) & 1U) +#define WRITE_BIT(var, pos, val) \ + ((val) ? SET_BIT(var, pos) : CLEAR_BIT(var, pos)) namespace flash { //////////////////////////////////////////////////////////////////////////////////////////////////// -__forceinline__ __device__ dim3 get_bidInfo(const int& blockType) { - - int m_block = blockIdx.y; - int bidb = blockIdx.z; - int bidh = blockIdx.x; - - if (blockType == 0) { - int m_block = blockIdx.x; - int bidb = blockIdx.z; - int bidh = blockIdx.y; - return dim3(m_block, bidb, bidh); - } - if (blockType == 1) { - int m_block = blockIdx.x; - int bidb = blockIdx.y; - int bidh = blockIdx.z; - return dim3(m_block, bidb, bidh); - } - - if (blockType == 2) { - int m_block = blockIdx.y; - int bidb = blockIdx.z; - int bidh = blockIdx.x; - return dim3(m_block, bidb, bidh); - } - - if (blockType == 3) { - int m_block = blockIdx.y; - int bidb = blockIdx.x; - int bidh = blockIdx.z; - return dim3(m_block, bidb, bidh); - } - - if (blockType == 4) { - int m_block = blockIdx.z; - int bidb = blockIdx.x; - int bidh = blockIdx.y; - return dim3(m_block, bidb, bidh); - } - - if (blockType == 5) { - int m_block = blockIdx.z; - int bidb = blockIdx.y; - int bidh = blockIdx.x; - return dim3(m_block, bidb, bidh); - } - - return dim3(0, 0, 0); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -__forceinline__ __device__ dim3 get_bidInfo(const int& blockType, const int& h, int& n_split_idx) { - - if constexpr(Split) { - int m_block = blockIdx.y; - int bidb = blockIdx.x / h; - int bidh = blockIdx.x - bidb * h; - n_split_idx = blockIdx.z; - if (blockType == 0) { - m_block = blockIdx.z; - n_split_idx = blockIdx.x; - bidb = blockIdx.y / h; - bidh = blockIdx.y - bidb * h; - return dim3(m_block, bidb, bidh); - } - if (blockType == 1) { - m_block = blockIdx.y; - n_split_idx = blockIdx.x; - bidb = blockIdx.z / h; - bidh = blockIdx.z - bidb * h; - return dim3(m_block, bidb, bidh); - } - if (blockType == 2) { - m_block = blockIdx.z; - n_split_idx = blockIdx.y; - bidb = blockIdx.x / h; - bidh = blockIdx.x - bidb * h; - return dim3(m_block, bidb, bidh); - } - if (blockType == 3) { - m_block = blockIdx.x; - n_split_idx = blockIdx.y; - bidb = blockIdx.z / h; - bidh = blockIdx.z - bidb * h; - return dim3(m_block, bidb, bidh); - } - if (blockType == 4) { - m_block = blockIdx.y; - n_split_idx = blockIdx.z; - bidb = blockIdx.x / h; - bidh = blockIdx.x - bidb * h; - return dim3(m_block, bidb, bidh); - } - if (blockType == 5) { - m_block = blockIdx.x; - n_split_idx = blockIdx.z; - bidb = blockIdx.y / h; - bidh = blockIdx.y - bidb * h; - return dim3(m_block, bidb, bidh); - } - return dim3(m_block, bidb, bidh); - } - n_split_idx = 0; - return get_bidInfo(blockType); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -//////////////////////////////////////////////////////////////////////////////////////////////////// - template struct MaxOp { __device__ __forceinline__ T operator()(T const & x, T const & y) { return x > y ? x : y; } @@ -167,7 +59,7 @@ __device__ __forceinline__ T operator()(T const & x, T const & y) { return x + y template struct Allreduce { - static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4); + static_assert(THREADS == 64 || THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4); template static __device__ __forceinline__ T run(T x, Operator &op) { constexpr int OFFSET = THREADS / 2; @@ -179,7 +71,18 @@ struct Allreduce { //////////////////////////////////////////////////////////////////////////////////////////////////// template<> -struct Allreduce<64> { +struct Allreduce<2> { +template +static __device__ __forceinline__ T run(T x, Operator &op) { + x = op(x, __shfl_xor_sync(uint64_t(-1), x, 1)); + return x; +} +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// reduce val(tidx) val(tidx+16) val(tidx+32) val(tidx+48) +struct Partialreduce { template static __device__ __forceinline__ T run(T x, Operator &op) { #if 0 @@ -202,39 +105,63 @@ struct Allreduce<64> { //////////////////////////////////////////////////////////////////////////////////////////////////// -template<> -struct Allreduce<2> { -template -static __device__ __forceinline__ T run(T x, Operator &op) { - x = op(x, __shfl_xor_sync(uint64_t(-1), x, 1)); - return x; -} -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - template -__forceinline__ __device__ void gemm_opt(Tensor0 &acc, Tensor1 &tCrA, Tensor2 &tCrB, Tensor3 const& tCsA, +__forceinline__ __device__ void gemm(Tensor0 &acc, Tensor1 &tCrA, Tensor2 &tCrB, Tensor3 const& tCsA, Tensor4 const& tCsB, TiledMma tiled_mma, TiledCopyA smem_tiled_copy_A, TiledCopyB smem_tiled_copy_B, ThrCopyA smem_thr_copy_A, ThrCopyB smem_thr_copy_B) { CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc)); // MMA_M CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc)); // MMA_N - CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K + CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA); CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // M Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB); CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // N - if (!A_in_regs) { cute::copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{})); } - if (!B_in_regs) { cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{})); } + if constexpr (!A_in_regs) { cute::copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{})); } + if constexpr (!B_in_regs) { cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{})); } #pragma unroll for (int i = 0; i < size<2>(tCrA); ++i) { if (i < size<2>(tCrA) - 1) { - if (!A_in_regs) { cute::copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1)); } - if (!B_in_regs) { cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1)); } + if constexpr (!A_in_regs) { cute::copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1)); } + if constexpr (!B_in_regs) { cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1)); } + } + cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc); + } +} + +template +__forceinline__ __device__ void gemm_prefetch_lds(Tensor0 &acc, Tensor1 &tCrA, Tensor2 &tCrB, Tensor3 const& tCsA, + Tensor4 const& tCsB, TiledMma tiled_mma, + TiledCopyA smem_tiled_copy_A, TiledCopyB smem_tiled_copy_B, + ThrCopyA smem_thr_copy_A, ThrCopyB smem_thr_copy_B) { + CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc)); // MMA_M + CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc)); // MMA_N + CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K + Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA); + CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // M + Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB); + CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // N + static_assert(decltype(size<2>(tCrA))::value >= prefetch_lds_num); + + #pragma unroll + for (int i = 0; i < prefetch_lds_num + 1; ++i) { + if constexpr (!A_in_regs) { cute::copy(smem_tiled_copy_A, tCsA(_, _, i), tCrA_copy_view(_, _, i)); } + if constexpr (!B_in_regs) { cute::copy(smem_tiled_copy_B, tCsB(_, _, i), tCrB_copy_view(_, _, i)); } + } + // do first gemm outside the loop, for compiler obey the sequence + cute::gemm(tiled_mma, tCrA(_, _, _0{}), tCrB(_, _, _0{}), acc); + + #pragma unroll + for (int i = 1; i < size<2>(tCrA); ++i) { + if (i + prefetch_lds_num < size<2>(tCrA)) { + if constexpr (!A_in_regs) { cute::copy(smem_tiled_copy_A, tCsA(_, _, i + prefetch_lds_num), tCrA_copy_view(_, _, i + prefetch_lds_num)); } + if constexpr (!B_in_regs) { cute::copy(smem_tiled_copy_B, tCsB(_, _, i + prefetch_lds_num), tCrB_copy_view(_, _, i + prefetch_lds_num)); } } cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc); } @@ -252,6 +179,7 @@ __forceinline__ __device__ void gemm_rs(Tensor0 &acc, Tensor1 &tCrA, Tensor2 &tC CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB); CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view)); // N + cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{})); #pragma unroll for (int i = 0; i < size<2>(tCrA); ++i) { @@ -278,7 +206,7 @@ __forceinline__ __device__ void gemm_rr(Tensor0 &acc, Tensor1 &tCrA, Tensor2 &tC // mc: Convert acc_layout from (MMA=4, MMA_M, MMA_N) to (nrow=(1, MMA_M), ncol=(4, MMA_N)) template __forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) { - static_assert(decltype(size<0>(acc_layout))::value == 4); + // static_assert(decltype(size<0>(acc_layout))::value == 4); static_assert(decltype(rank(acc_layout))::value == 3); //auto l = logical_divide(acc_layout, Shape<_2>{}); // ((2, 2), MMA_M, MMA_N) //return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l))); @@ -334,19 +262,22 @@ __forceinline__ __device__ void cp_async_wait() { __builtin_mxc_arrive_gvmcnt(N); } +// barrier_ex(2) == barrier_inst() +template __forceinline__ __device__ void sync_threads() { __builtin_mxc_arrive_bsmcnt(0); - __builtin_mxc_barrier_inst(); + __builtin_mxc_barrier_ex(N); } +template __forceinline__ __device__ void barrier() { - __builtin_mxc_barrier_inst(); + __builtin_mxc_barrier_ex(N); } -template +template __forceinline__ __device__ void barrier_gvm() { __builtin_mxc_arrive_gvmcnt(N); - __builtin_mxc_barrier_inst(); + __builtin_mxc_barrier_ex(M); } //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -442,66 +373,6 @@ __forceinline__ __device__ void copy_reg_to_global(Tensor cons //////////////////////////////////////////////////////////////////////////////////////////////////// -template -__forceinline__ __device__ void copy_reg_to_global4x4fp32(Tensor const &S, - Tensor &D, const int &d, const int max_MN=0) { - CUTE_STATIC_ASSERT_V(rank(S) == Int<2>{}); - CUTE_STATIC_ASSERT_V(rank(D) == Int<2>{}); - CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); - CUTE_STATIC_ASSERT_V(size<0>(S) == _16{}); - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_K - typedef __NATIVE_VECTOR__(4, int) VecType; - constexpr int kAtomLayoutMO = Kernel_traits::kAtomLayoutMO; - const int tidx = threadIdx.x; - const int warp_idx = tidx / 64; - const int lane_idx = tidx % 64; - int row_idx = warp_idx % kAtomLayoutMO * 16 + lane_idx % 16; - #pragma unroll - for (int k = 0; k < size<1>(S); ++k) { - #pragma unroll - for (int m = 0; m < 4; ++m) { - int col_idx = lane_idx / 16 * 16 + k * 128 + m * 4; - auto D_ptr = (VecType *)(reinterpret_cast(&D(4*m, k))); - auto S_ptr = (VecType const *)(reinterpret_cast(&S(4*m, k))); - bool col_mask = Is_even_K || col_idx < d; - bool row_mask = Is_even_MN || row_idx < max_MN; - __builtin_mxc_stg_b128_predicator(D_ptr, 0, S_ptr[0], true, false, false, col_mask && row_mask, 1, MACA_ICMP_EQ); - } - } -} - - -template -__forceinline__ __device__ void copy_w_min_idx(Tensor const &S, - Tensor &D, Tensor const &identity_MN, - const int &d, const int max_MN=0, const int min_MN=0) { - CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); - CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); - CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA - CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M - CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K - Tensor reg = make_fragment_like(S); - typedef __NATIVE_VECTOR__(4, int) VecType; - #pragma unroll - for (int m = 0; m < size<1>(S); ++m) { - #pragma unroll - for (int k = 0; k < size<2>(S); ++k) { - bool row_mask = get<0>(identity_MN(0, m, 0)) >= min_MN && get<0>(identity_MN(0, m, 0)) < max_MN; - bool col_mask = Is_even_K || get<1>(identity_MN(0, 0, k)) < d; - auto src_ptr = (VecType *)(S(_, m, k).data().ptr_); - auto reg_ptr = (VecType *)(reg(_, m, k).data()); - reg_ptr[0] = __builtin_mxc_ldg_b128_predicator(src_ptr, 0, false, true, false, false, - col_mask && row_mask, 1, MACA_ICMP_EQ); - - auto dst_ptr = (VecType *)(D(_, m, k).data().ptr_); - __builtin_mxc_stg_b128_predicator(dst_ptr, 0, reg_ptr[0], true, false, false, col_mask && row_mask, 1, MACA_ICMP_EQ); - } - } -} - template __forceinline__ __device__ void swap(T &a, T &b) { T tmp = a; @@ -542,6 +413,48 @@ __forceinline__ __device__ void copy_b128(Tensor const &S, } } +// for tensor shape is (cols=8, m, k). +template +__forceinline__ __device__ void copy_b128_bsm_async(Tensor0 const &S, + Tensor1 &&D, + Tensor2 const &identity_MN, + const int d, + const int max_MN=0) { + CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); + CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); + CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K + + typedef __NATIVE_VECTOR__(4, int) VecType; + #pragma unroll + for (int m = 0; m < size<1>(S); ++m) { + #pragma unroll + for (int k = 0; k < size<2>(S); ++k) { + auto src_ptr = (VecType *)(S(_, m, k).data().get()); // gmem pointer + auto dst_ptr = (VecType *)(D(_, m, k).data().get()); // smem pointer + bool col_mask = Is_even_K || get<1>(identity_MN(0, 0, k)) < d; + bool row_mask = Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN; + if constexpr (Is_even_K && Is_even_MN) { + __builtin_mxc_ldg_b128_bsm(dst_ptr, src_ptr, 0, -1, true, true, false, true); + } else { + __builtin_mxc_ldg_b128_bsm_predicator( + dst_ptr, // shared memory pointer + src_ptr, // global memory pointer + 0, // Immediate value,use the default value 0. + true, // bool + true, // bool + false, // bool + true, // bool,If it is true, the compiler will not insert arrive. + col_mask && row_mask, + 1, + MACA_ICMP_EQ + ); + } + } + } +} + // for tensor shape is (cols=4, m, k). template @@ -603,28 +516,6 @@ __forceinline__ __device__ void swap_fragment(Tensor &S) { //////////////////////////////////////////////////////////////////////////////////////////////////// -template -__forceinline__ __device__ void apply_softcap(Tensor &tensor, const float softcap){ - // #pragma unroll - // for (int i = 0; i < size(tensor); ++i) { - // tensor(i) = mctlass::fast_tanh(tensor(i) * softcap); - // } - static_assert(decltype(size(tensor))::value % 2 == 0); - typedef __NATIVE_VECTOR__(2, float) Float2; - Float2 scale_vec = {softcap, softcap}; - Float2 beta_vec = {0.0f, 0.0f}; - #pragma unroll - for (int i = 0; i < size(tensor); i += 2) { - // tensor(i) = mctlass::fast_tanh(tensor(i) * softcap); - Float2 x_vec = {tensor(i), tensor(i + 1)}; - x_vec = __builtin_mxc_pk_fma_f32(x_vec, scale_vec, beta_vec); - tensor(i) = mctlass::fast_tanh(x_vec[0]); - tensor(i + 1) = mctlass::fast_tanh(x_vec[1]); - } -} - - - // resolves offset of a slice of a paged kv copy from gmem. // assumes that the tensor has already been positioned at the correct head. __forceinline__ __device__ @@ -637,7 +528,16 @@ int64_t resolve_thread_kv_page_slice_offset(const int page_block_size, const int + col_offset; } +// when prefetch ldg page_idx, use the follow function +__forceinline__ __device__ +int64_t resolve_thread_kv_page_slice_offset(const int page_block_size, const int page_idx, const int page_offset, const int page_stride, const int row_stride, const int col_offset) { + return ((int64_t) page_idx) * ((int64_t) page_stride) + + page_offset * ((int64_t) row_stride) + + col_offset; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// template @@ -687,7 +587,55 @@ __forceinline__ __device__ void copy_b128_page_one(Tensor cons } } -//////////////////////////////////////////////////////////////////////////////////////////////////// +template +__forceinline__ __device__ void copy_b64_page_one(Tensor const &S_base, + Tensor &S, + Tensor &D, + Tensor const &identity_MN, + const int d, + const int n_block, + const int *block_table, + const int page_stride, + const int row_stride, + const int page_block_size, + const int page_idx, + const int page_offset, + const int max_MN=0) { + CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); + CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); + CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kNThreads = Kernel_traits::kNThreads; + constexpr int kElementPerThread = 4; + constexpr int kGmemThreadsPerRow = Kernel_traits::kBlockKSmem / kElementPerThread; + constexpr int kGmemRowsPerThread = 1; + // load 1x4 per thread + int tidx = threadIdx.x; + + typedef __NATIVE_VECTOR__(2, int) VecType; + #pragma unroll + for (int m = 0; m < size<1>(S); ++m) { + bool row_mask = Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN; + const int col_offset = tidx % kGmemThreadsPerRow * kElementPerThread; + const int64_t global_kv_page_offset = flash::resolve_thread_kv_page_slice_offset(page_block_size, page_idx, page_offset, page_stride, row_stride, col_offset); + #pragma unroll + for (int k = 0; k < size<2>(S); ++k) { + auto src_ptr = (VecType *)(S_base.data().get() + global_kv_page_offset + get<2>(S.stride()) * k); + auto dst_ptr = (VecType *)(D(_, m, k).data()); // rf + bool col_mask = Is_even_K || get<1>(identity_MN(0, 0, k)) < d; + if constexpr (Is_even_MN && Is_even_K) { + *dst_ptr = __builtin_mxc_ldg_b64(src_ptr, 0, -1, true, true, false, false); + } else { + *dst_ptr = __builtin_mxc_ldg_b64_predicator(src_ptr, 0, true, true, false, false, + row_mask && col_mask, 1, MACA_ICMP_EQ); + } + } + } +} + template @@ -737,6 +685,173 @@ __forceinline__ __device__ void copy_b64_page_one(Tensor const } } +// when prefetch ldg page_idx and ldgbsm, use the follow function +template +__forceinline__ __device__ void copy_b128_page_bsm_async(Tensor const &S_base, + Tensor &S, + Tensor &&D, + Tensor const &identity_MN, + const int d, + const int n_block, + const int *block_table, + const int page_stride, + const int row_stride, + const int page_block_size, + const uint32_t *page_idx, + const uint32_t *page_offset, + const int swz_offset, + const int max_MN=0) { + CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); + CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); + CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kNThreads = Kernel_traits::kNThreads; + constexpr int kGmemThreadsPerRow = Kernel_traits::kBlockKSmem / 8; + constexpr int kGmemRowsPerThread = 1; + + // load 1x8 per thread + int tidx = threadIdx.x; + + typedef __NATIVE_VECTOR__(4, int) VecType; + #pragma unroll + for (int m = 0; m < size<1>(S); ++m) { + bool row_mask = Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN; + const int row_offset = tidx / kGmemThreadsPerRow * kGmemRowsPerThread + kNThreads / kGmemThreadsPerRow * m + n_block * kBlockN; + const int col_offset = tidx % kGmemThreadsPerRow * 8; + const int64_t global_kv_page_offset = flash::resolve_thread_kv_page_slice_offset(page_block_size, page_idx[m], page_offset[m], page_stride, row_stride, col_offset); + #pragma unroll + for (int k = 0; k < size<2>(S); ++k) { + auto src_ptr = (VecType *)(S_base.data().get() + swz_offset + global_kv_page_offset + get<2>(S.stride()) * k); + auto dst_ptr = (VecType *)(D(_, m, k).data().get()); // rf + bool col_mask = Is_even_K || get<1>(identity_MN(0, 0, k)) < d; + if constexpr (Is_even_K && Is_even_MN) { + __builtin_mxc_ldg_b128_bsm(dst_ptr, src_ptr, 0, -1, true, true, false, true); + } else { + __builtin_mxc_ldg_b128_bsm_predicator( + dst_ptr, // shared memory pointer + src_ptr, // global memory pointer + 0, // Immediate value,use the default value 0. + true, // bool + true, // bool + false, // bool + true, // bool,If it is true, the compiler will not insert arrive. + col_mask && row_mask, + 1, + MACA_ICMP_EQ + ); + } + } + } +} + +template +__forceinline__ __device__ void copy_page(Tensor &S, + uint32_t *page_idx, + uint32_t *page_offset, + const int n_block, + const int *block_table, + const int page_block_size) { + CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kNThreads = Kernel_traits::kNThreads; + constexpr int kGmemThreadsPerRow = Kernel_traits::kBlockKSmem / 8; + constexpr int kGmemRowsPerThread = 1; + // load 1x8 per thread + int tidx = threadIdx.x; + const int log2_page_size = __builtin_ctz(page_block_size); + + #pragma unroll + for (int m = 0; m < size<1>(S); ++m) { + const int row_offset = tidx / kGmemThreadsPerRow * kGmemRowsPerThread + kNThreads / kGmemThreadsPerRow * m + n_block * kBlockN; + int virtual_page_idx = row_offset >> log2_page_size; + page_offset[m] = row_offset - virtual_page_idx * page_block_size; + page_idx[m] = block_table[virtual_page_idx]; + } +} + +template +__forceinline__ __device__ void copy_b32(Tensor const &S, + Tensor &D, + Tensor const &identity_MN, + const int d, + const int max_MN=0) { + CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); + CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); + CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K + + typedef __NATIVE_VECTOR__(1, int) VecType; + #pragma unroll + for (int m = 0; m < size<1>(S); ++m) { + bool row_mask = Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN; + #pragma unroll + for (int k = 0; k < size<2>(S); ++k) { + auto src_ptr = (VecType *)(S(_, m, k).data().get()); // gmem + auto dst_ptr = (VecType *)(D(_, m, k).data()); // rf + bool col_mask = Is_even_K || get<1>(identity_MN(0, 0, k)) < d; + if constexpr (Is_even_MN && Is_even_K) { + *dst_ptr = __builtin_mxc_ldg_b32(src_ptr, 0, -1, true, true, false, false); + } else { + *dst_ptr = __builtin_mxc_ldg_b32_predicator(src_ptr, 0, true, true, false, false, + row_mask && col_mask, 1, MACA_ICMP_EQ); + } + } + } +} +template +__forceinline__ __device__ void copy_b32_page_one(Tensor const &S_base, + Tensor &S, + Tensor &D, + Tensor const &identity_MN, + const int d, + const int n_block, + const int *block_table, + const int page_stride, + const int row_stride, + const int page_block_size, + const int page_idx, + const int page_offset, + const int max_MN=0) { + CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); + CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); + CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kNThreads = Kernel_traits::kNThreads; + constexpr int kElementPerThread = Kernel_traits::kGmemElemsPerLoadB32; + constexpr int kGmemThreadsPerRow = Kernel_traits::kBlockKSmem / kElementPerThread; + constexpr int kGmemRowsPerThread = 1; + // load 32 bit per thread, 1x2 fp16 + int tidx = threadIdx.x; + + typedef __NATIVE_VECTOR__(1, int) VecType; + #pragma unroll + for (int m = 0; m < size<1>(S); ++m) { + bool row_mask = Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN; + const int col_offset = tidx % kGmemThreadsPerRow * kElementPerThread; + const int64_t global_kv_page_offset = flash::resolve_thread_kv_page_slice_offset(page_block_size, page_idx, page_offset, page_stride, row_stride, col_offset); + #pragma unroll + for (int k = 0; k < size<2>(S); ++k) { + auto src_ptr = (VecType *)(S_base.data().get() + global_kv_page_offset + get<2>(S.stride()) * k); + auto dst_ptr = (VecType *)(D(_, m, k).data()); // rf + bool col_mask = Is_even_K || get<1>(identity_MN(0, 0, k)) < d; + if constexpr (Is_even_MN && Is_even_K) { + *dst_ptr = __builtin_mxc_ldg_b32(src_ptr, 0, -1, true, true, false, false); + } else { + *dst_ptr = __builtin_mxc_ldg_b32_predicator(src_ptr, 0, true, true, false, false, + row_mask && col_mask, 1, MACA_ICMP_EQ); + } + } + } +} + template __forceinline__ __device__ void copy_b32_page_one(Tensor const &S_base, @@ -757,10 +872,10 @@ __forceinline__ __device__ void copy_b32_page_one(Tensor const CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K constexpr int kBlockN = Kernel_traits::kBlockN; constexpr int kNThreads = Kernel_traits::kNThreads; - constexpr int kElementPerThread = 2; + constexpr int kElementPerThread = Kernel_traits::kGmemElemsPerLoadB32; constexpr int kGmemThreadsPerRow = Kernel_traits::kBlockKSmem / kElementPerThread; constexpr int kGmemRowsPerThread = 1; - // load 1x2 per thread + // load 32 bit per thread, 1x2 fp16 int tidx = threadIdx.x; typedef __NATIVE_VECTOR__(1, int) VecType; @@ -785,6 +900,47 @@ __forceinline__ __device__ void copy_b32_page_one(Tensor const } } +template +__forceinline__ __device__ void copy_b32_sparse(Tensor const &S_base, + Tensor &S, + Tensor &D, + Tensor const &identity_MN, + const int d, + const int n_block, + const int row_stride, + const int topk_sparse_idx, + const int max_MN=0) { + CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{}); + CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{}); + CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA + CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M + CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kNThreads = Kernel_traits::kNThreads; + constexpr int kElementPerThread = 2; + constexpr int kGmemThreadsPerRow = Kernel_traits::kBlockKSmem / kElementPerThread; + constexpr int kGmemRowsPerThread = 1; + // load 1x2 per thread + int tidx = threadIdx.x; + + typedef __NATIVE_VECTOR__(1, int) VecType; + #pragma unroll + for (int m = 0; m < size<1>(S); ++m) { + bool row_mask = Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN; + const int col_offset = tidx % kGmemThreadsPerRow * kElementPerThread; + const int64_t global_kv_sparse_offset = topk_sparse_idx * row_stride + col_offset; + #pragma unroll + for (int k = 0; k < size<2>(S); ++k) { + auto src_ptr = (VecType *)(S_base.data().get() + global_kv_sparse_offset + get<2>(S.stride()) * k); + auto dst_ptr = (VecType *)(D(_, m, k).data()); // rf + bool col_mask = Is_even_K || get<1>(identity_MN(0, 0, k)) < d; + *dst_ptr = __builtin_mxc_ldg_b32_predicator(src_ptr, 0, true, true, false, false, + row_mask && col_mask, 1, MACA_ICMP_EQ); + } + } +} + template __forceinline__ __device__ void concat(Tensor0 &lhs, Tensor1 &rhs, Tensor2 &out) { CUTE_STATIC_ASSERT_V(rank(lhs) == Int<3>{}); @@ -882,41 +1038,6 @@ __forceinline__ __device__ decltype(auto) permute_4x4_b16(Tensor return t; } -//////////////////////////////////////////////////////////////////////////////////////////////////// - -__forceinline__ __host__ mcDeviceProp_t mcGetCurrentDeviceProperties() { - int deviceId{}; - mcGetDevice(&deviceId); - mcDeviceProp_t dprops; - mcGetDeviceProperties(&dprops, deviceId); - return dprops; -} - -// solving the bank conflict of swizzle<2,3,3>, match with offset_swz233 -template -__forceinline__ __device__ void swap_swz233(Tensor &t) { - if (__lane_id() >= 32) { - flash::swap_fragment(t); - } -} - -// solving the bank conflict of swizzle<3,3,4>, match with offset_swz334. todo: need to combine with SWIZZLE_STORE_QDO -template -__forceinline__ __device__ void swap_swz334(Tensor &t) { - if (__lane_id() / 8 % 2 == 1) { - flash::swap_fragment(t); - } -} - -#define UNPACK_GRID(bidx, bidh, bidb, type) \ - if (type == 0) {bidx = blockIdx.x; bidh = blockIdx.y; bidb = blockIdx.z;} \ - else if (type == 1) {bidx = blockIdx.x; bidh = blockIdx.z; bidb = blockIdx.y;} \ - else if (type == 2) {bidx = blockIdx.y; bidh = blockIdx.x; bidb = blockIdx.z;} \ - else if (type == 3) {bidx = blockIdx.y; bidh = blockIdx.z; bidb = blockIdx.x;} \ - else if (type == 4) {bidx = blockIdx.z; bidh = blockIdx.y; bidb = blockIdx.x;} \ - else if (type == 5) {bidx = blockIdx.z; bidh = blockIdx.x; bidb = blockIdx.y;} \ - else {bidx = blockIdx.x; bidh = blockIdx.y; bidb = blockIdx.z;} \ - #define CHECK_MSG(x, ...) do { if((x) == false) {throw std::invalid_argument(__VA_ARGS__);} }while(0) #define CUDA_CHECK(expr) {auto x = (expr); CHECK_MSG(x == cudaSuccess, #expr + std::string(" check failed!"));} #define CUDA_KERNEL_LAUNCH_CHECK() CUDA_CHECK(cudaGetLastError()) diff --git a/flash_mla/__init__.py b/flash_mla/__init__.py index e37cd19..b83b936 100644 --- a/flash_mla/__init__.py +++ b/flash_mla/__init__.py @@ -3,5 +3,6 @@ __version__ = "1.0.1" from flash_mla.flash_mla_interface import( get_mla_metadata, - flash_mla_with_kvcache + flash_mla_with_kvcache, + flash_mla_sparse_fwd ) diff --git a/flash_mla/flash_mla_interface.py b/flash_mla/flash_mla_interface.py index 616469a..777dc90 100644 --- a/flash_mla/flash_mla_interface.py +++ b/flash_mla/flash_mla_interface.py @@ -1,28 +1,147 @@ # Adapted from deepseek-ai/FlashMLA(https://github.com/deepseek-ai/FlashMLA) -from typing import Optional, Tuple - +from typing import Optional, Tuple, Callable, Dict import torch -# import flash_mla_cuda import flash_mla_cuda as flash_mla - +from functools import wraps +import warnings +import numpy as np def get_mla_metadata( cache_seqlens: torch.Tensor, - num_heads_per_head_k: int, - num_heads_k: int, + # num_q_tokens_per_head_k: int, + # num_heads_k: int, + # num_heads_q: Optional[int] = None, + # is_fp8_kvcache: bool = False, + # topk: Optional[int] = None + *args, + **kwargs ) -> Tuple[torch.Tensor, torch.Tensor]: """ Arguments: cache_seqlens: (batch_size), dtype torch.int32. - num_heads_per_head_k: Equals to seq_len_q * num_heads_q // num_heads_k. - num_heads_k: num_heads_k. + num_q_tokens_per_head_k: Equals to num_q_tokens_per_q_seq * num_heads_q // num_heads_k. + num_heads_k: The number of k heads. + num_heads_q: The number of q heads. This argument is optional when sparse attention is not enabled + is_fp8_kvcache: Whether the k_cache and v_cache are in fp8 format. + topk: If not None, sparse attention will be enabled, and only tokens in the `indices` array passed to `flash_mla_with_kvcache_sm90` will be attended to. - Return: + Returns: tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32. num_splits: (batch_size + 1), dtype torch.int32. """ - return flash_mla.get_mla_metadata(cache_seqlens, num_heads_per_head_k, num_heads_k) + + if "num_heads_per_head_k" in kwargs: + """ + Handle Deprecation calls (using num_heads_per_head_k) + Arguments: + cache_seqlens: (batch_size), dtype torch.int32. + num_heads_per_head_k: Equals to seq_len_q * num_heads_q // num_heads_k. + num_heads_k: num_heads_k. + """ + warnings.warn( + "Parameter 'num_heads_per_head_k' is deprecated. Please use 'num_q_tokens_per_head_k' instead.", + DeprecationWarning, + stacklevel=2, + ) + num_heads_per_head_k = kwargs.pop("num_heads_per_head_k") + + if "num_heads_k" in kwargs: + num_heads_k = kwargs.pop("num_heads_k") + elif len(args) >= 1: + num_heads_k = args[0] + args = args[1:] + else: + raise TypeError( + "Legacy call missing required 'num_heads_k' (position 2 or keyword argument)" + ) + + if "num_heads_q" in kwargs or "is_fp8_kvcache" in kwargs or "topk" in kwargs: + raise TypeError( + "The legacy call does not support the parameters: (num_heads_q, is_fp8_kvcache, topk). \ + If you want to use them, please replace the parameter name 'num_heads_per_head_k' with 'num_q_tokens_per_head_k'" + ) + + if len(args) > 0 or len(kwargs) > 0: + extra_args = list(args) + list(kwargs.keys()) + raise TypeError( + f"Legacy calls do not support extra position parameters: {extra_args}. Legacy parameters only support (cache_seqlens, num_heads_per_head_k, num_heads_k) \ + \nIf you want to use (num_heads_q, is_fp8_kvcache, topk), please replace the parameter name 'num_heads_per_head_k' with 'num_q_tokens_per_head_k'" + ) + + return flash_mla.get_mla_metadata( + cache_seqlens, num_heads_per_head_k, num_heads_k, None, False, None + ) + else: + """ + Handle v32 calls (using num_q_tokens_per_head_k) + Arguments: + cache_seqlens: (batch_size), dtype torch.int32. + num_q_tokens_per_head_k: Equals to num_q_tokens_per_q_seq * num_heads_q // num_heads_k. + num_heads_k: The number of k heads. + num_heads_q: The number of q heads. This argument is optional when sparse attention is not enabled + is_fp8_kvcache: Whether the k_cache and v_cache are in fp8 format. + topk: If not None, sparse attention will be enabled, and only tokens in the `indices` array passed to `flash_mla_with_kvcache_sm90` will be attended to. + """ + if len(args) > 5: + raise TypeError( + f"get_mla_metadata() takes 6 positional arguments but {len(args)+1} were given" + ) + + if "num_q_tokens_per_head_k" in kwargs: + num_q_tokens_per_head_k = kwargs.pop("num_q_tokens_per_head_k") + elif len(args) >= 1: + num_q_tokens_per_head_k = args[0] + args = args[1:] + else: + raise TypeError( + "get_mla_metadata() missing required 'num_q_tokens_per_head_k' (position 1 or keyword argument)" + ) + + if len(args) >= 1 and "num_heads_k" not in kwargs: + num_heads_k = args[0] + args = args[1:] + elif "num_heads_k" in kwargs: + num_heads_k = kwargs.pop("num_heads_k") + elif len(args) <= 2: + raise TypeError( + "get_mla_metadata() missing required 'num_heads_k' (position 2 or keyword argument)" + ) + + num_heads_q: Optional[int] = None + is_fp8_kvcache: bool = False + topk: Optional[int] = None + if len(args) >= 1 and "num_heads_q" not in kwargs: + num_heads_q = args[0] + args = args[1:] + elif "num_heads_q" in kwargs: + num_heads_q = kwargs.pop("num_heads_q") + + if len(args) >= 1 and "is_fp8_kvcache" not in kwargs: + is_fp8_kvcache = args[0] + args = args[1:] + elif "is_fp8_kvcache" in kwargs: + is_fp8_kvcache = kwargs.pop("is_fp8_kvcache") + + if len(args) >= 1 and "topk" not in kwargs: + topk = args[0] + args = args[1:] + elif "topk" in kwargs: + topk = kwargs.pop("topk") + + if len(kwargs) > 0: + raise TypeError( + f"Unrecognized keyword arguments: {list(kwargs.keys())}. Supported parameters: cache_seqlens, num_q_tokens_per_head_k, num_heads_k, num_heads_q, is_fp8_kvcache, topk" + ) + + return flash_mla.get_mla_metadata( + cache_seqlens, + num_q_tokens_per_head_k, + num_heads_k, + num_heads_q, + is_fp8_kvcache, + topk, + ) def flash_mla_with_kvcache( @@ -35,6 +154,14 @@ def flash_mla_with_kvcache( num_splits: torch.Tensor, softmax_scale: Optional[float] = None, causal: bool = False, + is_fp8_kvcache: bool = False, + indices: Optional[torch.Tensor] = None, + indices_all_valid_per_q: Optional[torch.tensor] = None, + descale_q: Optional[torch.Tensor] = None, + descale_k: Optional[torch.Tensor] = None, + cp_world_size=1, + cp_rank=0, + cp_tot_seqlen_k=None, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Arguments: @@ -47,13 +174,29 @@ def flash_mla_with_kvcache( num_splits: (batch_size + 1), torch.int32, return by get_mla_metadata. softmax_scale: float. The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim). causal: bool. Whether to apply causal attention mask. - + is_fp8_kvcache: bool. Whether the k_cache and v_cache are in fp8 format. Do not support fp8 kvcahce right now. + indices: (batch_size, seq_len_q, topk), torch.int32. If not None, sparse attention will be enabled, and only tokens in the `indices` array will be attended to. Invalid indices should be set to -1 . + indices_all_valid_per_q: [batch_size, seq_len_q, 1], bool. A scalar boolean tensor indicating wheter all entries in indices for every query token are valid. + - descale_q: (batch_size), + torch.float32. Descaling factors for Q, used for fp8 quantization. + - descale_k: (batch_size), + torch.float32. Descaling factors for K, used for fp8 quantization. Return: out: (batch_size, seq_len_q, num_heads_q, head_dim_v). softmax_lse: (batch_size, num_heads_q, seq_len_q), torch.float32. """ if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) + if indices is not None: + assert causal == False, "causal must be `false` if sparse attention is enabled." + assert (descale_q is None) == (descale_k is None), ( + "descale_q and descale_k should be both None or both not None" + ) + #bf16/fp16 dense/sparse MLA + if indices_all_valid_per_q is None: + batch_size = q.shape[0] + seqlen_q = q.shape[1] + indices_all_valid_per_q = torch.full((batch_size, seqlen_q, 1), False, dtype=torch.bool, device=q.device) out, softmax_lse = flash_mla.fwd_kvcache_mla( q, k_cache, @@ -65,5 +208,44 @@ def flash_mla_with_kvcache( causal, tile_scheduler_metadata, num_splits, + is_fp8_kvcache, + indices, + indices_all_valid_per_q, + cp_world_size, + cp_rank, + cp_tot_seqlen_k, ) return out, softmax_lse + +def flash_mla_sparse_fwd( + q: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + sm_scale: float, + d_v: int = 512, + indices_all_valid_per_q: Optional[torch.tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Sparse attention prefill kernel + + Args: + q: [s_q, h_q, d_qk], bfloat16 + kv: [s_kv, h_kv, d_qk], bfloat16 + indices: [s_q, h_kv, topk], int32. Invalid indices should be set to -1 or numbers >= s_kv + sm_scale: float + indices_all_valid_per_q: [s_q, 1], bool. A scalar boolean tensor indicating wheter all entries for every query token in indices are valid. + d_v: The dimension of value vectors. Can only be 512 + + Returns: + (output, max_logits, lse) + About the definition of output, max_logits and lse, please refer to README.md + - output: [s_q, h_q, d_v], bfloat16 + - max_logits: [s_q, h_q], float + - lse: [s_q, h_q], float, 2-based log-sum-exp + """ + if indices_all_valid_per_q is None: + indices_all_valid_per_q = torch.full((q.shape[0], 1), False, dtype=torch.bool, device=q.device) + results = flash_mla.sparse_prefill_fwd( + q, kv, indices, sm_scale, d_v, indices_all_valid_per_q + ) + return results diff --git a/setup.py b/setup.py index deb687e..c1124b7 100644 --- a/setup.py +++ b/setup.py @@ -26,6 +26,193 @@ from torch.utils.cpp_extension import ( CUDA_HOME, ) +# HACK: we monkey patch pytorch's _write_ninja_file to pass +# "--offload-arch=xcore1000" to files ending in '_xcore1000.cu', +# "--offload-arch=xcore1500" to files ending in '_xcore1500.cu' +from torch.utils.cpp_extension import ( + IS_HIP_EXTENSION, + COMMON_HIP_FLAGS, + SUBPROCESS_DECODE_ARGS, + IS_WINDOWS, + get_cxx_compiler, + _join_rocm_home, + _join_cuda_home, + _is_cuda_file, + _maybe_write, +) +def _write_ninja_file(path, + cflags, + post_cflags, + cuda_cflags, + cuda_post_cflags, + cuda_dlink_post_cflags, + sources, + objects, + ldflags, + library_target, + with_cuda, + **kwargs, # kwargs (ignored) to absorb new flags in torch.utils.cpp_extension + ) -> None: + r"""Write a ninja file that does the desired compiling and linking. + + `path`: Where to write this file + `cflags`: list of flags to pass to $cxx. Can be None. + `post_cflags`: list of flags to append to the $cxx invocation. Can be None. + `cuda_cflags`: list of flags to pass to $nvcc. Can be None. + `cuda_postflags`: list of flags to append to the $nvcc invocation. Can be None. + `sources`: list of paths to source files + `objects`: list of desired paths to objects, one per source. + `ldflags`: list of flags to pass to linker. Can be None. + `library_target`: Name of the output library. Can be None; in that case, + we do no linking. + `with_cuda`: If we should be compiling with CUDA. + """ + def sanitize_flags(flags): + if flags is None: + return [] + else: + return [flag.strip() for flag in flags] + + cflags = sanitize_flags(cflags) + post_cflags = sanitize_flags(post_cflags) + cuda_cflags = sanitize_flags(cuda_cflags) + cuda_post_cflags = sanitize_flags(cuda_post_cflags) + cuda_dlink_post_cflags = sanitize_flags(cuda_dlink_post_cflags) + ldflags = sanitize_flags(ldflags) + + # Sanity checks... + assert len(sources) == len(objects) + assert len(sources) > 0 + + compiler = get_cxx_compiler() + + # Version 1.3 is required for the `deps` directive. + config = ['ninja_required_version = 1.3'] + config.append(f'cxx = {compiler}') + if with_cuda or cuda_dlink_post_cflags: + if "PYTORCH_NVCC" in os.environ: + nvcc = os.getenv("PYTORCH_NVCC") # user can set nvcc compiler with ccache using the environment variable here + else: + if IS_HIP_EXTENSION: + nvcc = _join_rocm_home('bin', 'hipcc') + else: + nvcc = _join_cuda_home('bin', 'nvcc') + config.append(f'nvcc = {nvcc}') + + if IS_HIP_EXTENSION: + post_cflags = COMMON_HIP_FLAGS + post_cflags + flags = [f'cflags = {" ".join(cflags)}'] + flags.append(f'post_cflags = {" ".join(post_cflags)}') + if with_cuda: + flags.append(f'cuda_cflags = {" ".join(cuda_cflags)}') + flags.append(f'cuda_post_cflags = {" ".join(cuda_post_cflags)}') + maca_post_cflags_xcore1000 = cuda_post_cflags + ['--offload-arch=xcore1000'] + flags.append(f'maca_post_cflags_xcore1000 = {" ".join(maca_post_cflags_xcore1000)}') + maca_post_cflags_xcore1500 = cuda_post_cflags + ['--offload-arch=xcore1500'] + flags.append(f'maca_post_cflags_xcore1500 = {" ".join(maca_post_cflags_xcore1500)}') + flags.append(f'cuda_dlink_post_cflags = {" ".join(cuda_dlink_post_cflags)}') + flags.append(f'ldflags = {" ".join(ldflags)}') + + # Turn into absolute paths so we can emit them into the ninja build + # file wherever it is. + sources = [os.path.abspath(file) for file in sources] + + # See https://ninja-build.org/build.ninja.html for reference. + compile_rule = ['rule compile'] + if IS_WINDOWS: + compile_rule.append( + ' command = cl /showIncludes $cflags -c $in /Fo$out $post_cflags') + compile_rule.append(' deps = msvc') + else: + compile_rule.append( + ' command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags') + compile_rule.append(' depfile = $out.d') + compile_rule.append(' deps = gcc') + + if with_cuda: + cuda_compile_rule = ['rule cuda_compile'] + nvcc_gendeps = '' + # --generate-dependencies-with-compile is not supported by ROCm + # Nvcc flag `--generate-dependencies-with-compile` is not supported by sccache, which may increase build time. + if torch.version.cuda is not None and os.getenv('TORCH_EXTENSION_SKIP_NVCC_GEN_DEPENDENCIES', '0') != '1': + cuda_compile_rule.append(' depfile = $out.d') + cuda_compile_rule.append(' deps = gcc') + # Note: non-system deps with nvcc are only supported + # on Linux so use --generate-dependencies-with-compile + # to make this work on Windows too. + nvcc_gendeps = '--generate-dependencies-with-compile --dependency-output $out.d' + maca_compile_rule_xcore1000 = ['rule maca_compile_xcore1000'] + cuda_compile_rule[1:] + [ + f' command = $nvcc {nvcc_gendeps} $cuda_cflags -c $in -o $out $maca_post_cflags_xcore1000' + ] + maca_compile_rule_xcore1500 = ['rule maca_compile_xcore1500'] + cuda_compile_rule[1:] + [ + f' command = $nvcc {nvcc_gendeps} $cuda_cflags -c $in -o $out $maca_post_cflags_xcore1500' + ] + cuda_compile_rule.append( + f' command = $nvcc {nvcc_gendeps} $cuda_cflags -c $in -o $out $cuda_post_cflags') + + # Emit one build rule per source to enable incremental build. + build = [] + for source_file, object_file in zip(sources, objects): + is_cuda_source = _is_cuda_file(source_file) and with_cuda + rule = 'cuda_compile' if is_cuda_source else 'compile' + if is_cuda_source: + if source_file.endswith('_xcore1000.cu'): + rule = 'maca_compile_xcore1000' + elif source_file.endswith('_xcore1500.cu'): + rule = 'maca_compile_xcore1500' + else: + rule = 'cuda_compile' + else: + rule = 'compile' + if IS_WINDOWS: + source_file = source_file.replace(':', '$:') + object_file = object_file.replace(':', '$:') + source_file = source_file.replace(" ", "$ ") + object_file = object_file.replace(" ", "$ ") + build.append(f'build {object_file}: {rule} {source_file}') + + if cuda_dlink_post_cflags: + devlink_out = os.path.join(os.path.dirname(objects[0]), 'dlink.o') + devlink_rule = ['rule cuda_devlink'] + devlink_rule.append(' command = $nvcc $in -o $out $cuda_dlink_post_cflags') + devlink = [f'build {devlink_out}: cuda_devlink {" ".join(objects)}'] + objects += [devlink_out] + else: + devlink_rule, devlink = [], [] + + if library_target is not None: + link_rule = ['rule link'] + if IS_WINDOWS: + cl_paths = subprocess.check_output(['where', + 'cl']).decode(*SUBPROCESS_DECODE_ARGS).split('\r\n') + if len(cl_paths) >= 1: + cl_path = os.path.dirname(cl_paths[0]).replace(':', '$:') + else: + raise RuntimeError("MSVC is required to load C++ extensions") + link_rule.append(f' command = "{cl_path}/link.exe" $in /nologo $ldflags /out:$out') + else: + link_rule.append(' command = $cxx $in $ldflags -o $out') + + link = [f'build {library_target}: link {" ".join(objects)}'] + + default = [f'default {library_target}'] + else: + link_rule, link, default = [], [], [] + + # 'Blocks' should be separated by newlines, for visual benefit. + blocks = [config, flags, compile_rule] + if with_cuda: + blocks.append(cuda_compile_rule) # type: ignore[possibly-undefined] + blocks.append(maca_compile_rule_xcore1000) # type: ignore[possibly-undefined] + blocks.append(maca_compile_rule_xcore1500) # type: ignore[possibly-undefined] + blocks += [devlink_rule, link_rule, build, devlink, link, default] + content = "\n\n".join("\n".join(b) for b in blocks) + # Ninja requires a new lines at the end of the .ninja file + content += "\n" + _maybe_write(path, content) + +# Monkey patching +torch.utils.cpp_extension._write_ninja_file = _write_ninja_file with open("README.md", "r", encoding="utf-8") as fh: long_description = fh.read() @@ -47,6 +234,10 @@ SKIP_CUDA_BUILD = os.getenv("FLASH_ATTENTION_SKIP_CUDA_BUILD", "FALSE") == "TRUE # For CI, we want the option to build with C++11 ABI since the nvcr images use C++11 ABI FORCE_CXX11_ABI = os.getenv("FLASH_ATTENTION_FORCE_CXX11_ABI", "FALSE") == "TRUE" +cmdclass = {} +ext_modules = [] + +flash_mla_root = Path(__file__).parent.parent def get_platform(): """ @@ -73,72 +264,28 @@ def get_cuda_bare_metal_version(cuda_dir): return raw_output, bare_metal_version -def check_if_cuda_home_none(global_option: str) -> None: - if CUDA_HOME is not None: - return - # warn instead of error because user could be downloading prebuilt wheels, so nvcc won't be necessary - # in that case. - warnings.warn( - f"{global_option} was requested, but nvcc was not found. Are you sure your environment has nvcc available? " - "If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, " - "only images whose names contain 'devel' will provide nvcc." - ) - - def append_nvcc_threads(nvcc_extra_args): return nvcc_extra_args + ["--threads", "4"] -def generate_kernels(script_path): +def get_arch_list(build_projects: str): + projects_list = build_projects.split(':') + all_arch_list = ["xcore1000", "xcore1500"] + arch_list = [] + for project in projects_list: + if project == "C500": + arch_list.append("xcore1000") + elif project == "C600": + arch_list.append("xcore1500") + arch_list = list(set(arch_list)) + if len(arch_list) == 0: + arch_list = all_arch_list + assert 0 < len(arch_list) and len(arch_list) <= len(all_arch_list) + return arch_list - result = subprocess.run(['bash', script_path], capture_output=True, text=True) - if(result): - return result.stdout - else: - return -1 - -def get_kernes_src(): - - # generate_src_path = "build" / "kernels" - - matches = [] - for root, dirs, files in os.walk("build/kernels/bwd"): - for filename in files: - matches.append(os.path.join(root, filename)) - - for root, dirs, files in os.walk("build/kernels/fwd"): - for filename in files: - matches.append(os.path.join(root, filename)) - - for root, dirs, files in os.walk("build/kernels/fwd_split"): - for filename in files: - matches.append(os.path.join(root, filename)) - - return matches - -def find_unmatched_cu_files(cu_dir, o_dir): - unmatched_files = [] - o_lists = [] - for root, _, files in os.walk(cu_dir): - for file in files: - if file.endswith('.cu'): - cu_file_path = os.path.join(root, file) - relative_path = os.path.relpath(cu_file_path, cu_dir) - o_file_path = os.path.join(o_dir, relative_path.replace('.cu', '.o')) - - if not os.path.exists(o_file_path): - unmatched_files.append(cu_file_path) - else: - o_lists.append(o_file_path) - - return unmatched_files, o_lists -cmdclass = {} -ext_modules = [] # We want this even if SKIP_CUDA_BUILD because when we run python setup.py sdist we want the .hpp -# files included in the source distribution, in case the user compiles from source. -# subprocess.run(["git", "submodule", "update", "--init", "csrc/mctlass"]) - +# files included in the source distribution, in case the user compiles from source if not SKIP_CUDA_BUILD: print("\n\ntorch.__version__ = {}\n\n".format(torch.__version__)) TORCH_MAJOR = int(torch.__version__.split(".")[0]) @@ -151,7 +298,6 @@ if not SKIP_CUDA_BUILD: if os.path.exists(os.path.join(torch_dir, "include", "ATen", "CUDAGeneratorImpl.h")): generator_flag = ["-DOLD_GENERATOR_PATH"] - check_if_cuda_home_none("flash_mla") # Check, if CUDA11 is installed for compute capability 8.0 cc_flag = [] if CUDA_HOME is not None: @@ -173,6 +319,67 @@ if not SKIP_CUDA_BUILD: extra_objects = ['{}/lib{}.so'.format(lib_dir, l) for l in libraries] # extra_objects.extend([f for f in obj_lists if f.endswith('.o')]) + feature_flag = [] + src_file = [ + "csrc/flash_api/flash_api.cpp", + "csrc/flash_run/run_mla_fwd.cpp", + "csrc/instantiations/flash_fwd_mla_metadata.cu", + "csrc/utils/print_parameter.cpp", + "csrc/utils/process_str.cpp", + "csrc/utils/logger.cpp", + ] + xcore1000_src_file = [ + # dense + "csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_bf16_True_True_1stage_xcore1000.cu", + "csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_fp16_True_True_1stage_xcore1000.cu", + "csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_bf16_True_True_split_xcore1000.cu", + "csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m16n16_fp16_True_True_split_xcore1000.cu", + "csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m32n16_bf16_True_True_split_xcore1000.cu", + "csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m32n16_fp16_True_True_split_xcore1000.cu", + "csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu", + "csrc/instantiations/xcore1000/flash_fwd_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1000.cu", + # sparse + "csrc/instantiations/xcore1000/flash_fwd_prefill_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu", + # "csrc/instantiations/xcore1000/flash_fwd_prefill_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1000.cu", # sparse prefill only supports bf16 + "csrc/instantiations/xcore1000/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1000.cu", + "csrc/instantiations/xcore1000/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1000.cu", + ] + xcore1500_src_file = [ + # dense + # xcore1500 recommand config is blockN 32 + # "csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu", + # "csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1500.cu", + "csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n32_bf16_True_True_split_xcore1500.cu", + "csrc/instantiations/xcore1500/flash_fwd_mla_hdimqk576_hdimv512_m64n32_fp16_True_True_split_xcore1500.cu", + # sparse + "csrc/instantiations/xcore1500/flash_fwd_prefill_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu", + "csrc/instantiations/xcore1500/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_bf16_True_True_split_xcore1500.cu", + "csrc/instantiations/xcore1500/flash_fwd_sparse_mla_hdimqk576_hdimv512_m64n16_fp16_True_True_split_xcore1500.cu", + ] + + build_projects = os.getenv("FLASHATTN_BUILD_PROJECTS", "C500:C600") + arch_list = get_arch_list(build_projects) + print(f"build arch list is {arch_list}") + only_one_arch = len(arch_list) == 1 + # we support build all arch or build only one specific arch + cucc_targets = "" + for arch in arch_list: + if cucc_targets == "": + cucc_targets = arch + else: + cucc_targets = cucc_targets + "," + arch + if arch == "xcore1000": + if only_one_arch: + feature_flag.append("-DXCORE1000") + src_file.extend(xcore1000_src_file) + elif arch == "xcore1500": + if only_one_arch: + feature_flag.append("-DXCORE1500") + src_file.extend(xcore1500_src_file) + # set CUCC_TARGETS for common kernel (e.g. flash_fwd_mla_metadata.cu) + os.environ['CUCC_TARGETS'] = cucc_targets + print("CUCC_TARGETS=",os.getenv('CUCC_TARGETS')) + # HACK: The compiler flag -D_GLIBCXX_USE_CXX11_ABI is set to be the same as # torch._C._GLIBCXX_USE_CXX11_ABI # https://github.com/pytorch/pytorch/blob/8472c24e3b5b60150096486616d98b7bea01500b/torch/utils/cpp_extension.py#L920 @@ -181,24 +388,9 @@ if not SKIP_CUDA_BUILD: ext_modules.append( CUDAExtension( name="flash_mla_cuda", - sources=[ - "csrc/flash_api/flash_api.cpp", - "csrc/flash_run/run_mha_fwd.cpp", - "csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_bf16_True_True_sm80.cu", - "csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_bf16_True_True_split_sm80.cu", - "csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_fp16_True_True_sm80.cu", - "csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m16n16_fp16_True_True_split_sm80.cu", - "csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_bf16_True_True_sm80.cu", - "csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_bf16_True_True_split_sm80.cu", - "csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_fp16_True_True_sm80.cu", - "csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m32n16_fp16_True_True_split_sm80.cu", - "csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_bf16_True_True_sm80.cu", - "csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_bf16_True_True_split_sm80.cu", - "csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_fp16_True_True_sm80.cu", - "csrc/flash_api/flash_fwd_splitkv_hdimqk576_hdimv512_m64n16_fp16_True_True_split_sm80.cu", - ], + sources=src_file, extra_compile_args={ - "cxx": ["-O3", "-std=c++17", "-w"] + generator_flag, + "cxx": ["-O3", "-std=c++17", "-w"] + generator_flag + feature_flag, "nvcc": append_nvcc_threads( [ "-O3", @@ -211,14 +403,26 @@ if not SKIP_CUDA_BUILD: "--expt-relaxed-constexpr", "--expt-extended-lambda", "--use_fast_math", + "-fno-strict-aliasing", + "-mllvm", + "-enable-tbaa=false", "-mllvm", "-metaxgpu-disable-bsm-offset=0",#try to merge ldg/lds + "-mllvm", + "-metaxgpu-sink-gep=true", "-D__FAST_HALF_CVT__", "-Xclang", "-menable-no-nans", + "-mllvm", + "--metaxgpu-prefer-scsel=true", + "-mllvm", + "-metaxgpu-mmasched=true", + "-mllvm", + "-metaxgpu-internalize-symbols", ] + generator_flag + cc_flag + + feature_flag ), }, include_dirs=[ @@ -234,27 +438,13 @@ if not SKIP_CUDA_BUILD: ) ) -def get_sha(flash_mla_root: Union[str, Path]) -> str: - try: - return ( - subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=flash_mla_root) - .decode("ascii") - .strip() - ) - except Exception: - return "Unknown" - def get_package_version(): - flash_mla_root = Path(__file__).parent.parent - sha = get_sha(flash_mla_root) with open(Path(this_dir) / "flash_mla" / "__init__.py", "r") as f: version_match = re.search(r"^__version__\s*=\s*(.*)$", f.read(), re.MULTILINE) public_version = str(ast.literal_eval(version_match.group(1))) maca_version = os.environ.get("CORE_MODULE_MACA_VERSION") if maca_version: return public_version + "+metax" + maca_version - elif sha != "Unknown": - return public_version + "+git" + sha[:7] else: return public_version + "+" diff --git a/tests/test_flash_mla.py b/tests/test_flash_mla.py deleted file mode 100644 index 2091dcb..0000000 --- a/tests/test_flash_mla.py +++ /dev/null @@ -1,135 +0,0 @@ -# Adapted from deepseek-ai/FlashMLA(https://github.com/deepseek-ai/FlashMLA) -import math -import random - -import torch -import triton -import pytest - -from flash_mla import ( - get_mla_metadata, - flash_mla_with_kvcache -) - - -def scaled_dot_product_attention(query, key, value, h_q, h_kv, is_causal=False): - query = query.float() - key = key.float() - value = value.float() - key = key.repeat_interleave(h_q // h_kv, dim=0) - value = value.repeat_interleave(h_q // h_kv, dim=0) - attn_weight = query @ key.transpose(-2, -1) / math.sqrt(query.size(-1)) - if is_causal: - s_q = query.shape[-2] - s_k = key.shape[-2] - attn_bias = torch.zeros(s_q, s_k, dtype=query.dtype) - temp_mask = torch.ones(s_q, s_k, dtype=torch.bool).tril(diagonal=s_k - s_q) - attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf")) - attn_bias.to(query.dtype) - attn_weight += attn_bias - lse = attn_weight.logsumexp(dim=-1) - attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32) - return attn_weight @ value, lse - - -def cal_diff(x: torch.Tensor, y: torch.Tensor, name: str) -> None: - x, y = x.double(), y.double() - RMSE = ((x - y) * (x - y)).mean().sqrt().item() - cos_diff = 1 - 2 * (x * y).sum().item() / max((x * x + y * y).sum().item(), 1e-12) - amax_diff = (x - y).abs().max().item() - # print(f"{name}: {cos_diff=}, {RMSE=}, {amax_diff=}") - assert cos_diff < 1e-5 - - -@torch.inference_mode() -def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, causal, varlen, block_size): - - cache_seqlens = torch.full((b,), mean_sk, dtype=torch.int32) - if varlen: - for i in range(b): - cache_seqlens[i] = max(random.normalvariate(mean_sk, mean_sk / 2), s_q) - total_seqlens = cache_seqlens.sum().item() - mean_seqlens = cache_seqlens.float().mean().int().item() - max_seqlen = cache_seqlens.max().item() - max_seqlen_pad = triton.cdiv(max_seqlen, 256) * 256 - # print(f"{total_seqlens=}, {mean_seqlens=}, {max_seqlen=}") - - q = torch.randn(b, s_q, h_q, d) - block_table = torch.arange(b * max_seqlen_pad // block_size, dtype=torch.int32).view(b, max_seqlen_pad // block_size) - blocked_k = torch.randn(block_table.numel(), block_size, h_kv, d) - for i in range(b): - blocked_k.view(b, max_seqlen_pad, h_kv, d)[i, cache_seqlens[i].item():] = float("nan") - blocked_v = blocked_k[..., :dv] - - tile_scheduler_metadata, num_splits = get_mla_metadata(cache_seqlens, s_q * h_q // h_kv, h_kv) - - def flash_mla(): - return flash_mla_with_kvcache( - q, blocked_k, block_table, cache_seqlens, dv, - tile_scheduler_metadata, num_splits, causal=causal, - ) - - def ref_mla(): - out = torch.empty(b, s_q, h_q, dv, dtype=torch.float32) - lse = torch.empty(b, h_q, s_q, dtype=torch.float32) - for i in range(b): - begin = i * max_seqlen_pad - end = begin + cache_seqlens[i] - O, LSE = scaled_dot_product_attention( - q[i].transpose(0, 1), - blocked_k.view(-1, h_kv, d)[begin:end].transpose(0, 1), - blocked_v.view(-1, h_kv, dv)[begin:end].transpose(0, 1), - h_q=h_q, - h_kv=h_kv, - is_causal=causal, - ) - out[i] = O.transpose(0, 1) - lse[i] = LSE - return out, lse - - out_flash, lse_flash = flash_mla() - out_torch, lse_torch = ref_mla() - cal_diff(out_flash, out_torch, "out") - cal_diff(lse_flash, lse_torch, "lse") - - - -@pytest.mark.parametrize("b",[128]) -@pytest.mark.parametrize("s_q",[1,2]) -@pytest.mark.parametrize("mean_sk",[4096, 8192, 32768]) -@pytest.mark.parametrize("h_q",[16, 32, 64, 128]) -@pytest.mark.parametrize("h_kv",[1]) -@pytest.mark.parametrize("d",[576]) -@pytest.mark.parametrize("dv",[512]) -@pytest.mark.parametrize("causal",[True]) -@pytest.mark.parametrize("varlen",[True,False]) -@pytest.mark.parametrize("block_size",[1,4,16,64]) -@pytest.mark.parametrize("dtype",[torch.bfloat16, torch.float16]) -def test_flash_mla_checkin(b, s_q, mean_sk, h_q, h_kv, d, dv, causal, varlen, block_size, dtype): - device = torch.device("cuda:0") - torch.set_default_dtype(dtype) - torch.set_default_device(device) - torch.cuda.set_device(device) - torch.manual_seed(0) - random.seed(0) - test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, causal, varlen, block_size) - -if __name__ == "__main__": - dtype = torch.bfloat16 - device = torch.device("cuda:0") - torch.set_default_dtype(dtype) - torch.set_default_device(device) - torch.cuda.set_device(device) - torch.manual_seed(0) - random.seed(0) - - h_kv = 1 - d, dv = 576, 512 - causal = True - for block_size in [1,4,16,64]: - for b in [128]: - for s in [4096, 8192, 32768]: - for h_q in [16, 32, 64, 128]: # TP = 8, 4, 2, 1 - for s_q in [1, 2]: # MTP =1, 2 - for varlen in [False, True]: - test_flash_mla(b, s_q, s, h_q, h_kv, d, dv, causal, varlen, block_size)